LSTMs and GRUs

A vanilla RNN updates its hidden state as $h_t = \tanh(W_h h_{t-1} + W_x x_t + b)$. During backpropagation through time, the gradient of the loss with respect to $h_{t-k}$ involves a product of $k$ Jacobians, each roughly $W_h^\top \operatorname{diag}(\tanh'(\cdot))$. Since $|\tanh'| \le 1$ and is typically much smaller away from the origin, and since repeated multiplication by the same (or similarly-scaled) matrix tends to shrink or blow up singular values geometrically, the gradient magnitude scales roughly like $\lambda^k$ for some effective factor $\lambda$. For $\lambda < 1$ this vanishes exponentially in $k$; for $\lambda > 1$ it explodes. This is the vanishing/exploding gradient problem, and it means vanilla RNNs cannot reliably learn dependencies spanning more than a few dozen steps. The LSTM (Long Short-Term Memory) fixes this by introducing a separate cell state $c_t$ that is updated by addition rather than repeated nonlinear transformation:

Here $f_t$ (the forget gate) and $i_t$ (the input gate) are vectors in $(0,1)^d$, produced by sigmoid layers, and $\tilde c_t$ is a candidate update. Because this update is additive rather than a matrix-multiply-then-squash, the gradient $\partial c_t / \partial c_{t-1} = f_t$ (elementwise), with no forced contraction from a shared weight matrix. If the network learns $f_t \approx 1$ for the dimensions that must persist, gradients flow through those dimensions almost unchanged across arbitrarily many steps — this is the mechanism, not a heuristic, that mitigates vanishing gradients. A separate hidden state (output) $h_t$ is derived from the cell state via an output gate $o_t$: $h_t = o_t \odot \tanh(c_t)$. All three gates and the candidate are computed from the same inputs $(x_t, h_{t-1})$ via learned affine transforms followed by nonlinearities — sigmoids for gates (bounding them to act as soft switches in $[0,1]$), $\tanh$ for the candidate content (bounding value magnitude to $[-1,1]$ for stability). The GRU (Gated Recurrent Unit) simplifies this design by merging the cell state and hidden state into a single vector $h_t$, and using only two gates: an update gate $z_t$ (playing a role similar to a combined forget/input gate) and a reset gate $r_t$ (controlling how much past information is used when computing the candidate). This yields comparable gradient-flow benefits with fewer parameters and one fewer state vector to carry, at the cost of slightly less fine-grained control over memory (no independent input/forget/output decisions).

LSTM cell. At each timestep $t$, given input $x_t \in \mathbb{R}^n$, previous hidden state $h_{t-1} \in \mathbb{R}^d$, and previous cell state $c_{t-1} \in \mathbb{R}^d$, compute:

GRU cell. Given $x_t$ and $h_{t-1}$ only (no separate cell state):

The GRU's $h_t$ update is structurally identical to the LSTM's $c_t$ update — a convex combination controlled by a gate — which is exactly what preserves the additive gradient path. The reset gate $r_t$ additionally lets the network "forget" past context specifically when forming the candidate, without discarding $h_{t-1}$ itself in the final blend. ``` function LSTM_STEP(x_t, h_prev, c_prev, params):

# params bundles W_, U_, b_ for f, i, o, c f_t = sigmoid(params.Wf @ x_t + params.Uf @ h_prev + params.bf) i_t = sigmoid(params.Wi @ x_t + params.Ui @ h_prev + params.bi) o_t = sigmoid(params.Wo @ x_t + params.Uo @ h_prev + params.bo) c_tilde = tanh(params.Wc @ x_t + params.Uc @ h_prev + params.bc) c_t = f_t c_prev + i_t c_tilde # elementwise; additive update h_t = o_t tanh(c_t) return h_t, c_t

function GRU_STEP(x_t, h_prev, params): z_t = sigmoid(params.Wz @ x_t + params.Uz @ h_prev + params.bz) r_t = sigmoid(params.Wr @ x_t + params.Ur @ h_prev + params.br) h_tilde = tanh(params.Wh @ x_t + params.Uh @ (r_t h_prev) + params.bh) h_t = (1 - z_t) h_prev + z_t * h_tilde return h_t

function RUN_SEQUENCE(inputs x_1..x_T, params, cell_type): h = zeros(d) c = zeros(d) # unused for GRU outputs = [] for t in 1..T: if cell_type == LSTM: h, c = LSTM_STEP(x_t, h, c, params) else: h = GRU_STEP(x_t, h, params) outputs.append(h) return outputs # feed into loss / next layer ``` Training proceeds exactly as with a vanilla RNN: unroll the recurrence across $T$ timesteps, compute a loss on the outputs, and run Backpropagation Through Time to get gradients with respect to all shared weight matrices $W_*, U_*, b_*$. The only difference from a vanilla RNN is the richer local computational graph at each timestep — more nodes, but each still $O(d^2)$ work and the chain rule still applies mechanically.

Time complexity. At each timestep, an LSTM computes four affine transforms of size $O(d(d+n))$ (each gate/candidate uses a $d \times d$ recurrent matrix and a $d \times n$ input matrix), where $d$ is hidden size and $n$ is input size. So each step costs $O(d^2 + dn)$, and a full sequence of length $T$ costs $O(T(d^2+dn))$. A GRU has three affine transforms instead of four, so it is a constant factor ($\approx 3/4$) cheaper per step but the same asymptotic order. Both match vanilla RNNs in asymptotic time complexity — the constant factor (3× or 4× more matrix multiplies) is the real cost of gating, not a complexity-class change.

Space complexity. Per-step memory is $O(d)$ for the hidden state (plus $O(d)$ for the cell state in LSTM); storing all intermediate activations for backpropagation through time costs $O(Td)$, since every gate/candidate value at every timestep must be cached for the backward pass. Parameter count is $O(d^2+dn)$ for GRU and $\frac{4}{3}$ times that for LSTM, independent of $T$.

Correctness of the gradient-flow fix — the key argument. Consider the cell-state recurrence $c_t = f_t \odot c_{t-1} + i_t \odot \tilde c_t$. The partial derivative of $c_t$ with respect to $c_{t-1}$, holding gates fixed, is:

Compare this to the vanilla RNN's $\partial h_t/\partial h_{t-1} = \operatorname{diag}(\tanh'(\cdot)) W_h$, which involves the same shared matrix $W_h$ at every step — its eigenvalues, raised to the $k$-th power over $k$ steps, dictate exponential growth or decay independent of what the task needs. In contrast $\operatorname{diag}(f_t)$ involves no matrix multiplication at all, just an elementwise gate that the network learns per-dimension, per-timestep. If the network sets $f_t \approx \mathbf 1$ (learned, e.g., by biasing $b_f$ to be large and positive at initialization — a well-known practical trick), then

so gradients backpropagate through the cell-state path across arbitrarily many steps without forced exponential decay. This is not a guarantee that gradients never vanish (if the task genuinely requires $f_t \ll 1$ at some step, that path should attenuate), but it removes the structural obstruction that plagued vanilla RNNs: the network now has the freedom, via learnable gates, to choose long memory when the task demands it. The same reasoning applies to the GRU's $h_t = (1-z_t)\odot h_{t-1} + z_t \odot \tilde h_t$, where $\partial h_t/\partial h_{t-1}$ has a $\operatorname{diag}(1-z_t)$ term along the direct path.

Tradeoffs. Gating does not eliminate vanishing gradients in the output gate / hidden-state path (LSTM's $h_t$ still passes through $\tanh$ and gate multiplications every step), and very long sequences (thousands of steps) can still be difficult. LSTMs have more parameters and compute than GRUs for the same hidden size, which matters at scale; GRUs sometimes underperform LSTMs on tasks needing finer memory control (e.g., counting tasks benefit from LSTM's independent output gate). Both remain strictly sequential across time — step $t$ needs $h_{t-1}$ — so they cannot be parallelized across the sequence dimension the way transformer self-attention can, which is the main reason they were later superseded in large-scale settings.

Consider a toy LSTM with hidden size $d=2$, scalar input $n=1$, tracking whether a running total of a stream of numbers has ever exceeded a threshold (a task requiring long memory: the "exceeded" fact must persist regardless of later inputs). Suppose after training, the relevant gate weights approximately realize: forget gate saturates near $1$ once the flag dimension is set (protecting the memory), input gate opens briefly at the triggering timestep, then closes. Trace with hand-picked (simplified, illustrative) parameter values, hidden dimension 1 (the "flag" unit), input $x_t$ = the stream value, threshold behavior baked into weights:

Read the trace as a story: at $t=2$ the input crosses the threshold, the input gate opens ($i_2=0.85$) and lets a large candidate value ($\tilde c_2 = 0.95$) into the cell state, jumping $c_t$ from near $0$ to $0.834$. From $t=3$ onward, even though the raw inputs drop back to small values, the forget gate stays close to 1 ($0.98, 0.99$), so the cell state is barely eroded ($0.834 \to 0.822 \to 0.814$) — the network has "remembered" that the threshold was crossed. The output gate $o_t$ grows over time, gradually exposing more of this memory into $h_t$, which a downstream classifier reads to answer "has the threshold ever been exceeded?" A vanilla RNN attempting the same task would have to re-encode this fact through a $\tanh$-squashed matrix multiply at every single step, and — because $\tanh'(\cdot) \le 1$ always — that signal would erode multiplicatively rather than being protected by a near-1 gate. For a gradient-flow comparison, suppose over 20 further steps the forget gate averages $\bar f = 0.98$ per dimension. The surviving fraction of the gradient signal reaching back 20 steps along the pure cell-state path is approximately $\bar f^{20} = 0.98^{20} \approx 0.667$ — a mild, controllable attenuation. Contrast this with a vanilla RNN where the effective per-step contraction factor is often empirically well below $0.5$ (products of $\tanh'$ terms near saturated regions and shrinking singular values), giving $0.5^{20} \approx 9.5\times 10^{-7}$ — six orders of magnitude smaller. That gap is the entire practical payoff of gating.

The LSTM was invented to directly solve the vanishing gradient problem that had been formally analyzed for recurrent networks in the early 1990s. Its constant-error-carousel idea — a self-loop with weight 1 acting as an uninterrupted gradient highway — was the conceptual seed for the additive cell-state update used today. The GRU emerged nearly two decades later as a leaner alternative designed for the same purpose but with a smaller parameter budget, becoming popular during the era of encoder-decoder machine translation systems.

Gated recurrent architectures powered production systems for roughly a decade: Google's neural machine translation system (pre-transformer) used stacked LSTMs; Apple's Siri and various speech-recognition pipelines used LSTM acoustic models; handwriting recognition, music generation, and time-series forecasting systems widely adopted LSTM/GRU cells. They remain a strong choice today for streaming or low-latency inference (constant per-step compute, no need to attend over a growing context window), for small-to-medium datasets where transformers' data hunger is a liability, and as building blocks inside larger systems such as Encoder-Decoder Architectures used in translation and summarization before attention mechanisms took over that role. Understanding gating is also essential background for appreciating why attention and transformers were adopted — attention solves the sequential-bottleneck and long-range-dependency problems that gating only partially addressed, without the sequential $O(T)$ dependency chain that limits Recurrent Neural Networks parallelism.

Ask the tutor: Can you derive, step by step, how the LSTM's forget-gate gradient path compares numerically to a vanilla RNN's gradient decay over 50 timesteps?