GRU

The Gated Recurrent Unit (GRU) is a recurrent cell that uses learned gates to control how much of the past hidden state is kept and how much new information is written. It is a lighter alternative to the LSTM that fights vanishing gradients without a separate cell state.

Motivation

Vanilla Recurrent Neural Networks update the hidden state with $h_t=\tanh (W_h{h}_{t-1}+W_x{x}_t+b)$. Repeatedly multiplying by $W_h$ makes gradients shrink or blow up over long sequences, so the network cannot learn long range dependencies. Gating Mechanisms fix this by letting the cell copy its state forward almost unchanged when the gates say so, creating a near identity path for gradients.

The Two Gates

A GRU has an update gate $z_t$ and a reset gate $r_t$, both computed with sigmoids so their entries live in $(0,1)$:

$$\begin{array}{rl}{z}_t& =\sigma (W_z{x}_t+U_z{h}_{t-1}+b_z)\\ r_t& =\sigma (W_r{x}_t+U_r{h}_{t-1}+b_r)\\ {\tilde{h}}_t& =\tanh (W_h{x}_t+U_h(r_t\odot h_{t-1})+b_h)\\ h_t& =(1-z_t)\odot h_{t-1}+z_t\odot {\tilde{h}}_t\end{array}$$

• The reset gate $r_t$ decides how much of the previous state feeds into the candidate ${\tilde{h}}_t$. When $r_t\to 0$ the cell forgets the past and behaves like it is reading a fresh sequence.
• The update gate $z_t$ is a learned interpolation between the old state and the candidate. When $z_t\to 0$ the state is copied verbatim, which is exactly the identity path that preserves gradients.

Interpolation view

The final line is a convex combination: $h_t=\text{lerp}(h_{t-1},{\tilde{h}}_t,z_t)$. This single blend replaces the LSTM’s separate forget and input gates.

GRU vs LSTM

Aspect	GRU	LSTM
Gates	2 (update, reset)	3 (forget, input, output)
Extra state	none	separate cell state $c_t$
Parameters	fewer	more
Output control	exposes full $h_t$	gated via output gate

GRUs train faster and use less memory. LSTMs sometimes edge them out on tasks that need very long memory because the protected cell state $c_t$ is never squashed by a $\tanh$ on the recurrent path.

PyTorch Implementation

We batch the three input projections and the two recurrent projections into single matrices, then split with einops.rearrange.

import torch
import torch.nn as nn
from einops import rearrange
 
class GRUCell(nn.Module):
    def __init__(self, d_in, d_hidden):
        super().__init__()
        # stack z, r, h projections -> 3 * d_hidden
        self.W = nn.Linear(d_in, 3 * d_hidden)
        self.U = nn.Linear(d_hidden, 3 * d_hidden, bias=False)
        self.d_hidden = d_hidden
 
    def forward(self, x, h_prev):
        # x: (b, d_in), h_prev: (b, d_hidden)
        xz, xr, xh = rearrange(self.W(x), "b (g d) -> g b d", g=3)
        uz, ur, uh = rearrange(self.U(h_prev), "b (g d) -> g b d", g=3)
 
        z = torch.sigmoid(xz + uz)
        r = torch.sigmoid(xr + ur)
        h_cand = torch.tanh(xh + r * uh)
        return (1 - z) * h_prev + z * h_cand

Reset gate placement

The reset gate multiplies the recurrent term $U_h{h}_{t-1}$, not the raw input. A common bug is applying $r_t$ to $h_{t-1}$ before both the $U_h$ and $U_z$, $U_r$ projections, which corrupts the gate computations. Only the candidate’s recurrent contribution is reset.

Practical Notes

• Initialize the update gate bias slightly negative so the cell starts by remembering (small $z_t$), which stabilizes early training.
• GRUs are strong baselines for small to medium sequence tasks; Transformers dominate once data and compute are large.