Bidirectional RNN Concatenation

Problem Statement

Combine forward and backward hidden states from a Bidirectional RNN.

Background

In Bidirectional RNNs, we process the sequence in both directions:

Forward pass: left to right
Backward pass: right to left

The final output at each timestep is typically the concatenation of both directions' hidden states, doubling the feature dimension.

Your Task

Write a function combine_bidirectional(forward_states, backward_states) that concatenates the forward and backward hidden states along the last dimension.

Input Format

forward_states: numpy array of shape (batch_size, seq_len, hidden_dim)
backward_states: numpy array of shape (batch_size, seq_len, hidden_dim)

Output Format

Return a numpy array of shape (batch_size, seq_len, 2 * hidden_dim).

1. Background Knowledge

Bidirectional RNNs (Bi-RNNs) process sequences in both forward (left-to-right) and backward (right-to-left) directions to capture full context. At each timestep $t$ , the output is formed by concatenating forward hidden state $\overrightarrow{h_t}$ and backward hidden state $\overleftarrow{h_t}$ :

h_t = [\overrightarrow{h_t}; \overleftarrow{h_t}] \in \mathbb{R}^{2 \times d_h}

where $d_h$ is the hidden dimension. This doubles the feature size, enabling richer representations for tasks like NLP and sequence classification. Standard RNNs suffer from vanishing gradients; Bi-RNNs with LSTM/GRU mitigate this while providing bidirectional context.

Prerequisites: NumPy array operations, tensor shapes (batch, sequence, features), and understanding RNN hidden states as contextual embeddings.

2. Algorithm Approach

Use NumPy's np.concatenate() along the last axis (axis=-1 or axis=2) to merge states element-wise:

def combine_bidirectional(forward_states, backward_states):
 return np.concatenate([forward_states, backward_states], axis=-1)

This is the standard Bi-RNN output fusion in frameworks like PyTorch (torch.cat((fwd, bwd), dim=-1)) and TensorFlow. No computation beyond concatenation—purely structural.

3. Step-by-Step Strategy

Verify inputs: Confirm forward_states.shape == backward_states.shape (e.g., (batch_size, seq_len, hidden_dim)).
Concatenate: np.concatenate([forward_states, backward_states], axis=-1) yields (batch_size, seq_len, 2 * hidden_dim).
Test shape: Assert result.shape == 2 * forward_states.shape.
Edge cases: Handle batch_size=1, seq_len=1, hidden_dim=1.

Sample implementation:

import numpy as np

def combine_bidirectional(forward_states, backward_states):
 assert forward_states.shape == backward_states.shape, "Shapes must match"
 return np.concatenate([forward_states, backward_states], axis=-1)

# Test
forward = np.ones((2, 3, 4))
backward = np.zeros((2, 3, 4))
result = combine_bidirectional(forward, backward)
print(result.shape) # (2, 3, 8)

4. Common Pitfalls

Wrong axis: Concatenating on axis=0 (batches) or axis=1 (sequence) breaks shapes—always use axis=-1.
Shape mismatch: Unchecked inputs cause runtime errors; add assertion.
Memory inefficiency: Large sequences (seq_len=512) are fine since concatenation is $O(1)$ in memory (views data).
Framework confusion: In PyTorch/TF, dim=-1; NumPy uses axis=-1—consistent but verify.
Reversing backward states: Backward pass already handles reversal; don't manually reverse arrays.

5. Time & Space Complexity

Time: $O(batch\_size \times seq\_len \times hidden\_dim) = O(N)$ , linear in input size (single pass copy).
Space: $O(batch\_size \times seq\_len \times 2 \times hidden\_dim) = O(2N)$ , doubles input size (new array allocated).

Constraints fit: Max $64 \times 512 \times 512 \approx 16M$ elements (~128MB float32)—efficient on CPU/GPU.