Stacked RNNs › Computational Cost Scaling

Stacking layers has a direct and significant impact on computational cost. Understanding these costs helps you make informed decisions about architecture design.

For a single LSTM layer processing a sequence of length $T$ with hidden size $d$, the computation at each timestep involves matrix multiplications of size $O(d^2)$ (the recurrent matrix $W_h$) and $O(d \cdot m)$ (the input matrix $W_x$). The total time for one forward pass is:

$$O_{\text{layer}} = O\!\left(T \cdot (d^2 + d \cdot m)\right)$$

For a stacked model with $L$ layers and the same hidden size $d$ (and noting that after the first layer the "input" is also size $d$):

$$O_{\text{stacked}} = O\!\left(T \cdot L \cdot d^2\right)$$

Time scales linearly with the number of layers. A 4-layer model takes approximately 4 times as long per batch as a 1-layer model with the same hidden size.

During training, you must store intermediate activations for backpropagation. For a single-layer LSTM:

• Hidden states: $O(T \cdot d)$ per sample
• Cell states: $O(T \cdot d)$ per sample
• Gate activations: $O(T \cdot d)$ per sample (4 gates)

Total memory per sample: $O(T \cdot d)$

For an $L$-layer stack, memory scales as $O(T \cdot L \cdot d)$: you must store the full state history for every layer to compute gradients.

For a 4-layer LSTM processing sequences of length 100 with hidden size 512:

• Memory for activations: $100 \times 4 \times 512 \times 4 \approx 800$K floats per sample
• At batch size 32: roughly 100 MB just for LSTM activations (in float32)

This is why long sequences with deep stacks require careful batch size management.

Parameter count for a stacked LSTM with $L$ layers, input size $m$, and hidden size $d$:

$$P = 4\!\left[(m + d) \cdot d + d\right] + 4(L-1)\!\left[(2d) \cdot d + d\right]$$

Simplified: the first layer has $O(d \cdot m + d^2)$ parameters and each subsequent layer has $O(d^2)$ parameters. Total parameters scale as $O(L \cdot d^2)$ for $L \gg 1$.