PyTorch GRU

GRU stands for Gated Recurrent Unit. It is a streamlined version of LSTM that uses only two gates instead of three and merges the cell state and hidden state into one. GRU often matches LSTM performance on many tasks while training faster and using less memory.

GRU vs LSTM at a Glance

               LSTM              GRU
─────────────────────────────────────────────────
States         h_t + c_t         h_t only
Gates          3 (forget,        2 (reset, update)
               input, output)
Parameters     More              Fewer (faster)
Performance    Excellent         Comparable to LSTM
Memory use     Higher            Lower

The Two GRU Gates

Reset Gate

The reset gate (r_t) decides how much of the previous hidden state to forget when computing the new candidate state. When r_t is near 0, the gate nearly ignores past memory — useful for resetting after a topic change in a sentence.

Update Gate

The update gate (z_t) blends the old hidden state with the new candidate state. When z_t is near 1, the old hidden state is kept mostly unchanged — the GRU ignores the current input. When z_t is near 0, the GRU replaces the old state with the new candidate.

GRU Step Diagram

Inputs: x_t (current), h_{t-1} (previous hidden)

Reset gate:     r_t = σ(W_r·[h_{t-1}, x_t])
Update gate:    z_t = σ(W_z·[h_{t-1}, x_t])
Candidate:      h̃_t = tanh(W·[r_t ⊙ h_{t-1}, x_t])
New hidden:     h_t = (1 - z_t) ⊙ h_{t-1} + z_t ⊙ h̃_t
                         ↑                   ↑
                   keep old state      add new candidate

⊙ = element-wise multiply
σ = sigmoid

PyTorch nn.GRU

import torch
import torch.nn as nn

gru = nn.GRU(
    input_size=10,
    hidden_size=32,
    num_layers=1,
    batch_first=True
)

x  = torch.rand(8, 20, 10)        # batch=8, seq_len=20, features=10
h0 = torch.zeros(1, 8, 32)        # initial hidden state (no cell state needed)

output, h_n = gru(x, h0)          # returns output sequence + final hidden state

print(output.shape)   # torch.Size([8, 20, 32])
print(h_n.shape)      # torch.Size([1, 8, 32])

GRU for Sequence Classification

class GRUClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes, num_layers=1):
        super().__init__()
        self.gru = nn.GRU(input_size, hidden_size,
                          num_layers=num_layers,
                          batch_first=True)
        self.fc  = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        _, h_n = self.gru(x)            # only need final hidden state
        out = self.fc(h_n[-1])          # h_n[-1] = last layer's hidden state
        return out

model = GRUClassifier(input_size=20, hidden_size=64, num_classes=3)
x = torch.rand(16, 30, 20)   # 16 sequences, 30 steps, 20 features
print(model(x).shape)         # torch.Size([16, 3])

Stacked and Bidirectional GRU

# Stacked GRU: 2 layers
gru_stacked = nn.GRU(10, 32, num_layers=2, batch_first=True, dropout=0.3)

# Bidirectional GRU
gru_bi = nn.GRU(10, 32, batch_first=True, bidirectional=True)

x = torch.rand(8, 20, 10)
out, h = gru_bi(x)
print(out.shape)   # torch.Size([8, 20, 64])  ← 32 forward + 32 backward

Real Example: Time Series Forecasting

class TimeSeriesGRU(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, output_size=1):
        super().__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers=2,
                          batch_first=True, dropout=0.2)
        self.fc  = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        output, _ = self.gru(x)
        # Use hidden state at final time step for next-step prediction
        return self.fc(output[:, -1, :])

model = TimeSeriesGRU(input_size=1, hidden_size=64, output_size=1)

# Input: 32 sequences, each 50 time steps, 1 feature (univariate)
x = torch.rand(32, 50, 1)
pred = model(x)
print(pred.shape)   # torch.Size([32, 1])  ← predicted next value

When to Choose GRU Over LSTM

Choose GRU when:
  → Dataset is small or medium (fewer params helps prevent overfitting)
  → Training speed matters (GRU trains faster)
  → Memory is limited
  → Your task shows similar performance with both (common)

Choose LSTM when:
  → Very long sequences (LSTM's cell state may retain more)
  → Task requires precise long-range memory (rare)
  → You have abundant data and compute

In practice: try both; GRU is the better starting point.

GRU Parameter Count vs LSTM

For input_size=D, hidden_size=H:

GRU parameters:
  3 × (D×H + H×H + H)   ← 3 matrices (for 2 gates + candidate)

LSTM parameters:
  4 × (D×H + H×H + H)   ← 4 matrices (for 3 gates + candidate)

GRU has 25% fewer parameters than LSTM for same D and H.

Summary

GRU is a simplified version of LSTM with two gates — reset and update — and a single hidden state. The reset gate controls how much past context to use when forming a new candidate state. The update gate blends old and new hidden states. GRU has 25% fewer parameters than LSTM, trains faster, and achieves comparable performance on most tasks. Use nn.GRU the same way as nn.LSTM but without the cell state. The return value is (output, h_n) instead of (output, (h_n, c_n)). GRU is an excellent default choice for sequence modeling tasks.