RNN Single Step Forward
Problem Statement
Compute the hidden state for a single timestep in a standard RNN.
Background
The update rule for a vanilla RNN is:
atβ=tanh(Waaββ atβ1β+Waxββ xtβ+baβ)
Where:
- atβ is the new hidden state
- atβ1β is the previous hidden state
- xtβ is the input at timestep t
- Waaβ is the hidden-to-hidden weight matrix
- Waxβ is the input-to-hidden weight matrix
- baβ 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:
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]
[0.9354, 0.9801] (approximately)
tanh(W_aa @ prev_h + W_ax @ x + b_a) = tanh([1.7, 2.3])
Constraints:
- All inputs are valid numpy arrays with compatible dimensions
- hidden_dim and input_dim are between 1 and 512
1. Background Knowledge
Recurrent Neural Networks (RNNs) are neural networks designed for sequential data processing, maintaining a hidden state atβ that captures information from previous timesteps. The vanilla RNN update follows:
atβ=tanh(Waaβatβ1β+Waxβxtβ+baβ)
Key prerequisites:
- Linear algebra: Matrix-vector multiplication (e.g., WaaββRnhβΓnhβ, atβ1ββRnhβ).
- Activation functions: tanh(z)=cosh(z)sinh(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: Waaββ atβ1β.
- Compute input-to-hidden contribution: Waxββ xtβ.
- Add bias: z=(Waaβatβ1β)+(Waxβxtβ)+baβ.
- Apply activation: atβ=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:
- Waxβx=[[0.5,0.6],[0.7,0.8]]β [1,2]=[1.7,2.3] (since Waaβatβ1β=0, baβ=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 (nhβ,niββ€512); matrix mult. is GPU-friendly in full RNNs.