GRU Cell - Sequence Models

GRU Cell

What you'll build

One time step of a Gated Recurrent Unit, matching PyTorch's nn.GRUCell exactly (weight layout, gate order and reset-gate placement).

Background

The GRU (Cho et al., 2014) merges the LSTM's cell and hidden state into one vector and uses two gates instead of three. It has 3 weight blocks instead of 4 (25% fewer parameters) and usually trains about as well. Its update rule is a learned interpolation between the old state and a candidate.

Formula

gi = x @ W_ih.T + b_ih              # (N, 3H), split into  i_r, i_z, i_n
gh = h_prev @ W_hh.T + b_hh         # (N, 3H), split into  h_r, h_z, h_n

r = sigmoid(i_r + h_r)              # reset gate
z = sigmoid(i_z + h_z)              # update gate
n = tanh(i_n + r * h_n)             # candidate state
h = (1 - z) * n + z * h_prev

Conventions to get right

• Gate order r, z, n along the first axis of the weights (rows 0:H, H:2H, 2H:3H), as in nn.GRUCell.weight_ih / weight_hh. Weights are (out, in), hence .T.
• The update gate z keeps the old state: z = 1 copies h_prev, z = 0 takes the candidate n. Some papers swap the roles of z and 1 - z; PyTorch uses the form above.
• Reset gate placement: r multiplies h_n = h_prev @ W_hn.T + b_hn, including the bias b_hn, after the matrix product. The original paper computes (r * h_prev) @ W_hn.T instead, which gives different numbers. Follow PyTorch. This is why the input and hidden projections must stay separate (you can't add gi + gh before splitting).

Inputs

Tensor	Shape
x	(N, D)
h_prev	(N, H)
W_ih	(3H, D)
W_hh	(3H, H)
b_ih, b_hh	(3H,)

Expected return format

The new hidden state h: a float array of shape (N, H).

Edge cases and conventions

• A sigmoid helper is provided. If |h_prev| < 1, the new h stays in (-1, 1) (a convex mix of values in that range).

Constraints

NumPy only; vectorised over the batch (two matmuls total).

Follow-up

• Count the parameters of nn.GRUCell(D, H) (3(HD + H^2 + 2H)) and compare with an LSTM.
• When would you pick a GRU over an LSTM, and vice versa?

Input: x = [[0.3, -0.1]], h_prev = [[1.0, -2.0]], all weights and biases zero
Output: array([[0.5, -1.0]])

Input: D = H = 1, x = [[1.0]], h_prev = [[0.5]], W_ih = [[0], [1], [1]], W_hh = [[1], [0], [2]], biases 0
Output: array([[0.614295]])

import numpy as np

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


def gru_cell(x, h_prev, W_ih, W_hh, b_ih, b_hh):
    """
    One GRU step (gate order r, z, n, as in nn.GRUCell).

    Args:
        x: Input (N, D)
        h_prev: Previous hidden state (N, H)
        W_ih: Input weights (3H, D), row blocks [r; z; n]
        W_hh: Hidden weights (3H, H), row blocks [r; z; n]
        b_ih, b_hh: Biases (3H,)

    Returns:
        h: New hidden state (N, H)
    """
    # Your code here
    pass