61 - GRU: The Lean Version of Gated Memory

The Gated Recurrent Unit (GRU, Cho and colleagues, 2014) reaches the same idea as the LSTM with fewer parts: no separate cell state, two gates instead of three. The hidden state h acts as memory and output at once, and at every step it is updated as a blend of the old state and a new candidate state, according to a single update gate. The result is a cell with about three quarters of the parameters of an LSTM and, on the tasks tried here, equivalent performance.

The words you need

This article takes for granted the definitions of hidden state, sigmoid, gate and element-wise product introduced in the articles on RNNs and LSTMs. We add two terms. A convex blend is a weighted average (1 − z)·a + z·b with z between 0 and 1: the result always lies "between" a and b (with z = 0 it is a, with z = 1 it is b, with z = 0.5 the mean), and if a and b lie in (−1, 1) so does the blend. The candidate is the "new content" the cell could write into memory, computed at every step; how much of it actually ends up in the state is decided by a gate.

The analogy: a blackboard with a slider

The LSTM was a tape with three attendants. The GRU is a blackboard with two knobs. The first, the update gate, decides how much of the board to leave as it is and how much to wipe to write the new content; it has a single slider, so "how much I keep" and "how much I write" are tied together: the more I keep, the less I write. The second, the reset gate, decides how much of the board to look at while the new content is being prepared: if it is 0, the candidate is written as if there were no past at all.

The equations

How to read the diagram. The line at the top is the state h, which here is memory and output together. The previous state is multiplied by z (how much to keep). At the bottom, the reset gate r multiplies the previous state only to build the candidate n (the tanh): if r is close to 0 the candidate ignores the past. The candidate, multiplied by 1 − z, is added at the top. The comparison with the LSTM is immediate: the separate tape is gone, one gate is gone, and "keeping" and "writing" are controlled by the same z.

Here we follow PyTorch's convention, which is the one to know if you write real code. Careful: in Cho's original paper the role of z is inverted (h_t = (1−z)·h_{t−1} + z·candidate), and two different books may write opposite formulas for the same cell.

r_t = σ( W_ir x_t + b_ir + W_hr h_{t-1} + b_hr )          reset gate
z_t = σ( W_iz x_t + b_iz + W_hz h_{t-1} + b_hz )          update gate (how much to KEEP)
n_t = tanh( W_in x_t + b_in + r_t ⊙ (W_hn h_{t-1} + b_hn) )   candidate state
h_t = (1 − z_t) ⊙ n_t + z_t ⊙ h_{t-1}                        blend

Two details that change the numerical result. First: in PyTorch the reset multiplies W_hn h_{t-1} + b_hn, that is it is applied after the product with the matrix; the original formula applied it before (W (r ⊙ h)). The two variants are equivalent in capacity but not in the numbers, so if you copy weights between frameworks you must follow the right convention. Second: there are three n × d matrices and three n × n ones (reset, update, candidate), hence 3·n·(n + d + 2) parameters with PyTorch's two biases.

Why the gradient survives

From the blend line:

h_t = (1 − z_t) ⊙ n_t + z_t ⊙ h_{t-1}
∂h_t/∂h_{t-1} = diag(z_t) + (terms that go through the gates and through n_t)

It is the same reasoning as for the LSTM. The diagonal term diag(z_t) is a path with no recurrent matrix and no tanh derivative: if a slot of the state has z ≈ 1, its value and its gradient pass almost intact to the next step. The structural difference: in the LSTM "how much I keep" (f) and "how much I write" (i) are two independent gates, in the GRU they are complementary (z and 1 − z). Less freedom, fewer parameters.

One GRU step with the numbers in hand

With a single state slot, h_prev = 0.9 (the memory remembers a high value) and candidate n = −0.5 (the new content is negative), we apply h = (1 − z)·n + z·h_prev with three values of z. With z = 1.0 we get h = 0·(−0.5) + 1·0.9 = 0.9: the cell ignores the new and keeps the old, and this is the "direct wire" that carries the gradient. With z = 0.0 we get h = 1·(−0.5) + 0·0.9 = −0.5: the cell forgets the past and adopts the candidate, like a plain RNN. With z = 0.8 we get h = 0.2·(−0.5) + 0.8·0.9 = −0.1 + 0.72 = 0.62: a soft update, 20% novelty and 80% memory.

The behaviour of the cell therefore lies entirely in the number z, learned at every step as a function of input and state.

Code by hand

import numpy as np

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))

class GRUCell:
    """GRU with PyTorch's conventions (reset applied AFTER the hidden matrix, z = how much to KEEP)."""
    def __init__(self, d, n, seed=0):
        rng = np.random.default_rng(seed)
        self.n = n
        self.W_ih = rng.normal(0, 0.3, (3 * n, d))      # rows: [reset | update | new]
        self.W_hh = rng.normal(0, 0.3, (3 * n, n))
        self.b_ih = rng.normal(0, 0.1, 3 * n)
        self.b_hh = rng.normal(0, 0.1, 3 * n)

    def step(self, x, h):
        n = self.n
        gi = self.W_ih @ x + self.b_ih
        gh = self.W_hh @ h + self.b_hh
        r = sigmoid(gi[0:n] + gh[0:n])                   # reset gate: how much past to use for the candidate
        z = sigmoid(gi[n:2 * n] + gh[n:2 * n])           # update gate: how much old state to keep
        cand = np.tanh(gi[2 * n:] + r * gh[2 * n:])      # candidate state
        return (1 - z) * cand + z * h                    # convex mix of new and old

cell = GRUCell(d=3, n=4)
rng = np.random.default_rng(1)
seq = rng.normal(size=(6, 3))
h = np.zeros(4)
for x in seq:
    h = cell.step(x, h)
print(h.round(3))

The structure mirrors the LSTM, but the matrices are three instead of four and are already split into W_ih (acts on the input) and W_hh (acts on the state), as in PyTorch. In step we compute gi (the input's contribution) and gh (the state's contribution); the gates r and z add the two contributions and go through the sigmoid. The candidate is the most delicate line: np.tanh(gi[2*n:] + r * gh[2*n:]) multiplies the reset by the state's contribution alone, which is PyTorch's convention. The last line is the blend: (1 - z) * cand + z * h. With a high z the state hardly changes.

Then with the library, and counting the parameters

We copy the weights into nn.GRU (the rows of the matrices are stacked in the order reset, update, new) and compare the parameters of the three types of recurrent layer with input 2 and state 64.

import torch
import torch.nn as nn

gru = nn.GRU(input_size=3, hidden_size=4, batch_first=True)
with torch.no_grad():
    for name, value in [("weight_ih_l0", cell.W_ih), ("weight_hh_l0", cell.W_hh),
                        ("bias_ih_l0", cell.b_ih), ("bias_hh_l0", cell.b_hh)]:
        getattr(gru, name).copy_(torch.tensor(value, dtype=torch.float32))

_, h_n = gru(torch.tensor(seq, dtype=torch.float32).unsqueeze(0))
print(np.abs(h_n[0, 0].detach().numpy() - h).max())          # ~1e-8: same function

def n_params(m):
    return sum(p.numel() for p in m.parameters())

for name, cls in [("RNN", nn.RNN), ("GRU", nn.GRU), ("LSTM", nn.LSTM)]:
    print(name, n_params(cls(2, 64, batch_first=True)))        # input 2, hidden 64

The loop over the names copies the four matrices/vectors from our cell into nn.GRU: unlike the LSTM, nothing has to be split here, because we had already kept input and state separate. The value 1.36e-08 is the float32/float64 numerical error. The last loop instantiates the three types of recurrent layer with the same sizes (input 2, state 64) and counts the parameters with sum(p.numel() ...): 4352 for the RNN, 13056 for the GRU, 17408 for the LSTM. The ratios 1 : 3 : 4 are the number of transformations (one, three, four).

1.3557848121958926e-08
RNN 4352
GRU 13056
LSTM 17408

The GRU has 13,056 parameters against the LSTM's 17,408 (75%, exactly 3 out of 4 as the formula predicts) and three times the plain RNN. It is a count, not a measure of quality.

A test on the adding problem

We reuse the benchmark from the previous article (two marked numbers to add, baseline 0.167) with a GRU, 4000 training steps, clipping at 1.0, one seed.

def adding_batch(batch, T):
    values = torch.rand(batch, T)
    markers = torch.zeros(batch, T)
    for b in range(batch):
        markers[b, torch.randperm(T)[:2]] = 1.0
    return torch.stack([values, markers], dim=-1), (values * markers).sum(dim=1)

# Same adding-problem benchmark as the LSTM article, with a GRU instead of an LSTM
class NetGRU(nn.Module):
    def __init__(self, hidden=64):
        super().__init__()
        self.rec = nn.GRU(2, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)
    def forward(self, x):
        out, _ = self.rec(x)
        return self.head(out[:, -1]).squeeze(-1)

def run_gru(T, steps=4000, seed=0):
    torch.manual_seed(seed)
    net = NetGRU()
    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
    for _ in range(steps):
        x, y = adding_batch(64, T)
        loss = nn.functional.mse_loss(net(x), y)
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
    with torch.no_grad():
        x, y = adding_batch(1000, T)
        return nn.functional.mse_loss(net(x), y).item()

for T in (50, 100):
    print(f"T={T:3d}  GRU {run_gru(T):.4f}")

The code redefines adding_batch (identical to the one in the previous article), so same task, same baseline (0.167), same 4000 steps. The only difference is nn.GRU in place of nn.LSTM, and the absence of the special bias initialisation. That it works "first time" is consistent with the diag(z) term seen above, but remember the limit: one seed, a synthetic task. It does not prove that the GRU is better than the LSTM, it proves that it is not worse on this task.

T= 50  GRU 0.0001
T=100  GRU 0.0001

The GRU solves the task at 50 and at 100 steps, with errors of the same order as the LSTM's (0.0005 and 0.0002 in the previous experiment; with a single seed the differences between 0.0001 and 0.0005 are not significant). Note that here we did not have to set any special bias: the GRU has no forget gate to initialise.

Looking inside: what the update gate does after training

So far we have asserted that a high z means "I keep" and a low z "I write". Let us check by looking at the gate in a trained network. We train a 32-unit GRU on the adding problem at T = 50 and then, with the learned weights, redo the computation by hand recording z at every step, for 500 sequences. Then we compare the mean value of z at the steps where the marker was present (those to be written into memory) with the one at the steps without a marker.

# Train a GRU on the adding problem (T=50), then LOOK at its update gate z during a forward pass.
torch.manual_seed(0)
gru_net = NetGRU(hidden=32)
opt = torch.optim.Adam(gru_net.parameters(), lr=1e-2)
for _ in range(1500):
    x, y = adding_batch(64, 50)
    loss = nn.functional.mse_loss(gru_net(x), y)
    opt.zero_grad(); loss.backward()
    nn.utils.clip_grad_norm_(gru_net.parameters(), 1.0)
    opt.step()

@torch.no_grad()
def update_gate_trace(gru, x):
    """Recompute the GRU by hand with the trained weights and return z at every step (PyTorch convention)."""
    n = gru.hidden_size
    W_ih, W_hh = gru.weight_ih_l0, gru.weight_hh_l0
    b_ih, b_hh = gru.bias_ih_l0, gru.bias_hh_l0
    h = torch.zeros(x.size(0), n)
    zs = []
    for t in range(x.size(1)):
        gi = x[:, t] @ W_ih.T + b_ih
        gh = h @ W_hh.T + b_hh
        r = torch.sigmoid(gi[:, :n] + gh[:, :n])
        z = torch.sigmoid(gi[:, n:2 * n] + gh[:, n:2 * n])          # z = how much to KEEP
        cand = torch.tanh(gi[:, 2 * n:] + r * gh[:, 2 * n:])
        h = (1 - z) * cand + z * h
        zs.append(z)
    return torch.stack(zs, dim=1)                                    # (batch, T, n)

x, y = adding_batch(500, 50)
z = update_gate_trace(gru_net.rec, x)                                # (500, 50, 32)
marked = x[:, :, 1] == 1                                             # positions of the two markers
print("mean z (keep) at MARKED steps  :", round(z.mean(dim=-1)[marked].mean().item(), 3))
print("mean z (keep) at UNMARKED steps:", round(z.mean(dim=-1)[~marked].mean().item(), 3))
print("final loss:", round(nn.functional.mse_loss(gru_net(x), y).item(), 5))

mean z (keep) at MARKED steps  : 0.672
mean z (keep) at UNMARKED steps: 0.869
final loss: 8e-05

Explaining the code. update_gate_trace does not use nn.GRU as a black box: it reads the trained weights (weight_ih_l0 and so on) and repeats the computation of our cell by hand, step by step, saving z at every step. It is the same code as the NumPy cell from before, but in PyTorch and on batches. Then marked = x[:, :, 1] == 1 builds a boolean mask of the steps with a marker, and z.mean(dim=-1) averages over all 32 units. Result: at the steps with a marker the gate keeps on average 67.2% of the old state, at the steps without a marker 86.9%. The difference goes in the expected direction (at a marker the network leaves more room for new information), but it is not a clean switch: it is the mean over 32 units, many of which do nothing specific, and they are means over a single trained model. A finer inspection would look unit by unit. In any case it is a simple way not to trust the story blindly: gates are numbers you can read, and it is worth doing so when a model behaves strangely.

The GRU as a simplified LSTM: a table of correspondences

Comparing the equations, the GRU is obtained from the LSTM with three simplifications. Here they are in a table.

LSTM                                   GRU
cell state c + hidden state h          a single state h
forget gate f and input gate i         a single gate z:  keep = z,  write = 1 − z
output gate o (filters c towards h)    none: the whole state is the output
candidate g from [h, x]                candidate n from [r ⊙ h, x]  (reset gate)
3 gates + candidate = 4 transformations 2 gates + candidate = 3 transformations

There are three consequences. First, the GRU exposes all of its state at every step: it cannot keep a piece of information "hidden" from the rest of the network as the LSTM does with the output gate. Second, it has no unlimited memory: the state is a convex blend of values in (−1, 1), so it stays in (−1, 1), whereas the LSTM's c can grow. Third, the reset r is a novelty with no equivalent in the LSTM: it lets the candidate "start from scratch" ignoring memory, while keeping the memory in the state.

Checked with four calculations

A GRU with input 5 and state 10 has, with PyTorch's two biases, 3·10·(10+5+2) = 510 parameters. If z = 1 in all units and at all steps, h_t = h_{t−1}: the cell ignores the inputs and keeps the initial state (zero), a "frozen" cell. If r = 0 and z = 0 the candidate does not see the memory and the state is entirely replaced by the candidate, h_t = tanh(W x_t + b), that is a network with no memory (a dense layer applied at every step). A forget-gate bias is not needed because there is no separate forget gate; it still makes sense to initialise the update gate's bias towards positive values to start with a tendency to "keep", but the results above show that the GRU learned without it too.

A race at equal parameters, over three seeds

Saying "the GRU is lighter" is one thing, saying "at equal parameters it is also faster to train" is another claim and has to be measured. Let us run an honest race on the adding problem at T = 50: GRU with state 64 (13,121 parameters, head included), LSTM with state 64 (17,473) and LSTM with state 55 (13,036, practically the same number of parameters as the GRU). Metric: the steps needed for the moving average of the loss over the last 100 steps to drop below 0.05 (the metric of the previous lesson), with three seeds per configuration.

class AnyNet(nn.Module):
    def __init__(self, kind, hidden):
        super().__init__()
        self.rec = {"GRU": nn.GRU, "LSTM": nn.LSTM}[kind](2, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)
        if kind == "LSTM":
            with torch.no_grad():
                self.rec.bias_ih_l0[hidden:2 * hidden] = 1.0
    def forward(self, x):
        out, _ = self.rec(x)
        return self.head(out[:, -1]).squeeze(-1)

def solve_steps(kind, hidden, seed, T=50, max_steps=4000, threshold=0.05):
    torch.manual_seed(seed)
    net = AnyNet(kind, hidden)
    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
    recent = []
    for step in range(1, max_steps + 1):
        x, y = adding_batch(64, T)
        loss = nn.functional.mse_loss(net(x), y)
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)
        opt.step()
        recent.append(loss.item()); recent = recent[-100:]
        if step >= 100 and sum(recent) / 100 < threshold:
            return step
    return None

count = lambda m: sum(p.numel() for p in m.parameters())
configs = [("GRU", 64), ("LSTM", 64), ("LSTM", 55)]     # LSTM(55) has about as many parameters as GRU(64)
for kind, hidden in configs:
    n_par = count(AnyNet(kind, hidden))
    steps = [solve_steps(kind, hidden, seed) for seed in range(3)]
    print(f"{kind:4s} hidden={hidden:3d} params={n_par:6d}  steps to solve per seed: {steps}")

GRU  hidden= 64 params= 13121  steps to solve per seed: [423, 393, 399]
LSTM hidden= 64 params= 17473  steps to solve per seed: [1027, 841, 947]
LSTM hidden= 55 params= 13036  steps to solve per seed: [708, 694, 1035]

Explaining the code. AnyNet is a network with a GRU or an LSTM and a linear head (with the forget-gate bias at 1 for the LSTM, so as not to disadvantage it). solve_steps is the same training as before (Adam, learning rate 0.003, clipping 1.0) with a stopping condition. count sums the parameters. The list configs chooses the three configurations. Result: the GRU solves the task in about 400 steps on all three seeds (393-423), the LSTM with state 64 in about 940 (841-1027), the parameter-matched LSTM in 694-1035. The GRU is faster by a factor of two or more in this experiment, and the difference is much larger than the variability between seeds. Careful about what can be said: it holds for this task, at this length, with this learning rate and this LSTM initialisation. It is not a law. But it is a datum that the argument "the GRU is just a poorer LSTM" does not predict.

Plausible explanations, not verified separately: the GRU has a simpler direct path for the state (a single gate controls keeping and writing), has fewer parameters to coordinate and does not have to learn to manage two states. To be sure you would need ablations. The practical message stands: if your training budget is limited, the GRU deserves to be tried first; if validation shows that the LSTM generalises better, switch to that.

GRU or LSTM: how to choose

• Parameters and speed. The GRU is smaller and a bit faster per step; with small datasets it has less risk of overfitting.
• Performance. The empirical comparison by Chung and colleagues (2014) on music and speech modelling and the one by Jozefowicz and colleagues (2015) show no absolute winner: it depends on the task.
• Counting. Weiss, Goldberg and Yahav (2018) show that an LSTM can implement an unbounded counter, whereas the state of a GRU, a convex blend of values in (−1, 1), stays bounded.
• Rule of thumb. Start from the GRU for speed, move to the LSTM if validation justifies it, and treat the choice as a hyperparameter to measure.

The four criteria do not weigh the same: the first two say that a priori there is no winner, the third points to the only case where the difference is structural (tasks that require counting), the fourth turns everything into a procedure. In practice the decision is made with a trial on a fixed budget, not by reading other people's tables.

Real traps

The first trap is opposite formulas between sources: whether z means "keep" or "write" changes from source to source, so check the code and not only the formula. The second is the reset before or after the matrix: copying weights between frameworks without checking it gives different results without any error. The third is the forgotten initial state: the GRU has a single state, so it is easier to carry forward between batches, but it must be detached from the graph (h.detach()) if you reuse it. The last is thinking that fewer gates mean less optimisation: clipping remains useful.

One GRU step with the numbers, in PyTorch's convention

A single unit: input x = 1.0, previous state h = 0.5. Weights (on the state w_h, on the input w_x, bias): reset (0.2, 0.4, 0), update (0.3, 0.5, 0.5), candidate (w_x = 0.8, w_h = 0.6, 0).

r = σ(0.4·1.0 + 0.2·0.5)       = σ(0.50)  = 0.622       (I use 62% of the past for the candidate)
z = σ(0.5·1.0 + 0.3·0.5 + 0.5) = σ(1.15)  = 0.760       (I KEEP 76% of the old state)
n = tanh( 0.8·1.0 + r · 0.6·0.5 ) = tanh(0.8 + 0.622·0.3) = tanh(0.987) = 0.756   (candidate)
h' = (1 − z)·n + z·h = 0.240·0.756 + 0.760·0.5 = 0.181 + 0.380 = 0.562

Reading. The new state (0.562) is a blend: 24% comes from the candidate (0.756) and 76% from the old state (0.5). Since z is high, the state changes little, from 0.5 to 0.562. Had z been 0.1 the new state would have been 0.9·0.756 + 0.1·0.5 = 0.730: almost entirely the candidate. Note also where the reset comes in: it multiplies only the state's contribution inside the candidate (r · 0.6 · 0.5), not the input and not the old state that is blended in at the end. This is what lets the GRU "start from scratch" in computing the candidate while still keeping the old state for the blend.

Frequently asked questions

My book writes h = z·h_prev + (1 − z)·n, PyTorch does not. Who is right? Nobody is wrong: it is a convention. In PyTorch z is the fraction of the old state that is kept (as in our code); in the work of Cho and colleagues (2014) z weighted the candidate. The behaviour is identical up to renaming z → 1 − z. When you copy weights between libraries, check the convention with a numerical test like the one in the lesson.

Does Keras also have the reset "after" the matrix? It depends on the version and on a parameter: in recent versions of TensorFlow/Keras the GRU layer has an option called reset_after that chooses between the two variants (apply the reset before or after the state's matrix), with a default that favours the one compatible with cuDNN. If you bring weights from PyTorch to Keras check that option, or the numbers will not match.

Is the GRU always better than the LSTM because it has fewer parameters? No: in our experiments it learned faster (≈400 steps against ≈950 at T = 50), but the literature finds no absolute winner, and for tasks that require counting or memorising precisely the LSTM may be preferable. Measure on your own data.

How many gates does a GRU really have? Two gates (reset and update) and a candidate: three linear transformations, against the LSTM's four (three gates plus candidate). Hence the 3 to 4 ratio in the parameters (13,056 against 17,408 with input 2 and state 64).

Can I stack bidirectional GRUs? Yes: nn.GRU(..., num_layers=2, bidirectional=True); the output of the first layer has size 2H and becomes the input of the second. The rule from the article on bidirectional RNNs applies: only if the sequence is complete at the moment of the decision.

Exercises with solutions

Before moving on, try to solve these exercises with pen and paper; the solutions (checked with the code) are right below each one.

Exercise 1. How many parameters does nn.GRU(10, 20) have with PyTorch's two biases? Solution: 3·20·(20 + 10 + 2) = 1,920.

Exercise 2. With z = 0.9, candidate n = −1 and previous state h = 0.5 (PyTorch convention), what is the new state? Solution: (1 − 0.9)·(−1) + 0.9·0.5 = −0.1 + 0.45 = 0.35: the state changes little because z is high.

Exercise 3. What happens if z = 1 everywhere? Solution: h_t = h_{t−1} at every step: the cell ignores the inputs and keeps the initial state forever, the limiting case in which the gradient passes without any loss. Exercise 4. At equal sizes, what is the ratio between the parameters of a GRU and those of an LSTM? Solution: 3 to 4, because the GRU has three transformations (reset, update, candidate) against the LSTM's four (three gates plus candidate).

What you can do now

You can write the GRU equations in PyTorch's convention, explain why the diag(z) term preserves the gradient like the LSTM's forget gate, and follow a complete step with every number in view. You can copy the weights of a hand-written cell into nn.GRU and count the parameters (ratio 1 : 3 : 4 between RNN, GRU and LSTM), read the update gate in a trained network and compare GRU and LSTM on the steps needed to solve a task. The next step is reading the sequence in both directions, with bidirectional RNNs.