RNN Single Step Forward

mediumFREERecurrent Neural Networks (RNNs)Deep LearningDeep Learning 1

Problem Statement

Compute the hidden state for a single timestep in a standard RNN.

Background

The update rule for a vanilla RNN is:

$a_t = \tanh(W_{aa} \cdot a_{t-1} + W_{ax} \cdot x_t + b_a)$

Where:

$a_t$ is the new hidden state
$a_{t-1}$ is the previous hidden state
$x_t$ is the input at timestep $t$
$W_{aa}$ is the hidden-to-hidden weight matrix
$W_{ax}$ is the input-to-hidden weight matrix
$b_a$ is the bias vector

Your Task

Write a function rnn_step(prev_hidden, input_vec, W_aa, W_ax, b_a) that computes and returns the new hidden state using np.tanh.

Input Format

prev_hidden: numpy array of shape (hidden_dim,)
input_vec: numpy array of shape (input_dim,)
W_aa: numpy array of shape (hidden_dim, hidden_dim)
W_ax: numpy array of shape (hidden_dim, input_dim)
b_a: numpy array of shape (hidden_dim,)

Output Format

Return a numpy array of shape (hidden_dim,) representing the new hidden state.

Example

Input

prev_hidden=[0,0], input_vec=[1,2], W_aa=[[0.1,0.2],[0.3,0.4]], W_ax=[[0.5,0.6],[0.7,0.8]], b_a=[0,0]

Output

[0.9354, 0.9801] (approximately)

Reasoning

Inputs: the previous hidden state $a_{t-1} = [0, 0]$ (2 hidden units), the current input $x_t = [1, 2]$ (2 features), two $2 \times 2$ weight matrices and a zero bias. Each row of a weight matrix belongs to one hidden unit.
Recurrent term: $W_{aa} a_{t-1}$ multiplies every weight by $0$ , so it is $[0, 0]$ . At the first timestep there is no memory yet.
Input term: row 1 of $W_{ax}$ gives $0.5 \cdot 1 + 0.6 \cdot 2 = 1.7$ ; row 2 gives $0.7 \cdot 1 + 0.8 \cdot 2 = 2.3$ .
Pre-activation: add the two terms and $b_a = [0, 0]$ to get $[1.7, 2.3]$ .
Squash: applying $\tanh$ to each entry turns $1.7$ into $\approx 0.9354$ and $2.3$ into $\approx 0.9801$ , giving [0.9354, 0.9801], a vector with one value per hidden unit.
Notice that both values sit close to $1$ : $\tanh$ maps any real number into $(-1, 1)$ and saturates for large inputs, so the second unit's larger pre-activation gains only a little.

Constraints

All inputs are valid numpy arrays with compatible dimensions
hidden_dim and input_dim are between 1 and 512

Learn this in

Introduction to LLMsLLM Study Plan Introduction to NLPNLP Study Plan

More from Deep Learning 1

Early Stopping Checkeasy L2 Regularization Penaltyeasy Inverted Dropout Forwardmedium Binary Cross-Entropy Losseasy Momentum Update Stepmedium

Concept primer

1. Background Knowledge

Recurrent Neural Networks (RNNs) are neural networks designed for sequential data processing, maintaining a hidden state $a_t$ that captures information from previous timesteps. The vanilla RNN update follows:

$a_t = \tanh(W_{aa} a_{t-1} + W_{ax} x_t + b_a)$

Key prerequisites:

Linear algebra: Matrix-vector multiplication (e.g., $W_{aa} \in \mathbb{R}^{n_h \times n_h}$ , $a_{t-1} \in \mathbb{R}^{n_h}$ ).
Activation functions: $\tanh(z) = \frac{\sinh(z)}{\cosh(z)}$ squashes values to $[-1, 1]$ , enabling non-linearity.
NumPy basics: Array shapes, broadcasting, np.tanh, np.dot.
RNN context: This single-step forward pass is the core computation unrolled in full sequences for tasks like language modeling.

2. Algorithm Approach

The direct matrix computation approach implements the RNN equation:

Compute hidden-to-hidden contribution: $W_{aa} \cdot a_{t-1}$ .
Compute input-to-hidden contribution: $W_{ax} \cdot x_t$ .
Add bias: $z = (W_{aa} a_{t-1}) + (W_{ax} x_t) + b_a$ .
Apply activation: $a_t = \tanh(z)$ .

This mirrors forward propagation in RNN training (BPTT unrolls it across time). No loops or optimization needed—just algebraic operations.

3. Step-by-Step Strategy

def rnn_step(prev_hidden, input_vec, W_aa, W_ax, b_a):
 # Step 1: Hidden-to-hidden: shape (hidden_dim, hidden_dim) @ (hidden_dim,) -> (hidden_dim,)
 hidden_contrib = np.dot(W_aa, prev_hidden)

 # Step 2: Input-to-hidden: shape (hidden_dim, input_dim) @ (input_dim,) -> (hidden_dim,)
 input_contrib = np.dot(W_ax, input_vec)

 # Step 3: Pre-activation: element-wise sum + bias
 z = hidden_contrib + input_contrib + b_a

 # Step 4: Apply tanh activation
 return np.tanh(z)

Verification with sample:

$W_{ax} x = [[0.5,0.6],[0.7,0.8]] \cdot [1,2] = [1.7, 2.3]$ (since $W_{aa} a_{t-1} = 0$ , $b_a = 0$ ).
Matches expected $\tanh([1.7, 2.3])$ .

4. Common Pitfalls

Shape mismatches: Ensure np.dot inputs align—(m,n) @ (n,) yields (m,). Use prev_hidden.shape for debugging.
Broadcasting errors: + b_a works due to NumPy broadcasting, but verify all are (hidden_dim,).
Using @ vs dot: For 2D matrices/vectors, both work; prefer np.dot for clarity with 1D vectors.
Forgetting tanh: Linear output without np.tanh breaks non-linearity.
In-place ops: Avoid z +=... if inputs are views; use explicit addition.
Dimension assumptions: Code must handle hidden_dim ≠ input_dim (1-512).

5. Time & Space Complexity

Aspect	Complexity	Details
Time	O(n_h² + n_h n_i)	Dominated by W_{aa} a_{t-1} (n_h²) and W_{ax} x_t (n_h n_i), where n_h = hidden_dim, n_i = input_dim. \tanh is O(n_h).
Space	O(n_h)	Temporary vectors; inputs reused. No recursion/sequences stored.

Scalability: Efficient for constraints ( $n_h, n_i \leq 512$ ); matrix mult. is GPU-friendly in full RNNs.

solution.py

Python 3.13

Test results