Gated Recurrent Units (GRUs)

Long Short-Term Memory (LSTM) networks are known for effectively addressing the vanishing gradient problem and capturing long-range dependencies. However, these networks often introduce considerable computational complexity due to their architecture, featuring three distinct gates (input, forget, output) and a separate cell state. In 2014, Cho et al. proposed the Gated Recurrent Unit (GRU) as a variation that achieves similar performance on many tasks but with a simpler architecture. GRUs combine the forget and input gates into a single "update gate" and merge the cell state and hidden state.

Let's look at the structure of a GRU cell. Like other RNNs, it takes the current input $x_t$ and the previous hidden state $h_{t-1}$ to produce the next hidden state $h_t$. The magic happens through two gates: the Reset Gate ($r_t$) and the Update Gate ($z_t$).

Reset Gate

The reset gate determines how to combine the new input with the previous hidden state. Specifically, it controls how much of the previous hidden state ($h_{t-1}$) should be "forgotten" when calculating a candidate hidden state. It's computed as follows:

$$r_t = \sigma(W_r x_t + U_r h_{t-1} + b_r)$$

Here, $W_r$, $U_r$, and $b_r$ are learnable weight matrices and a bias vector for the reset gate. The sigmoid function $\sigma$ squashes the output between 0 and 1. A value close to 0 means the previous hidden state is mostly ignored, while a value close to 1 means it's mostly kept.

Update Gate

The update gate decides how much of the previous hidden state ($h_{t-1}$) should be carried forward to the new hidden state ($h_t$) versus how much of the new candidate hidden state should be used. This gate essentially combines the roles of the LSTM's forget and input gates. It's calculated similarly to the reset gate:

$$z_t = \sigma(W_z x_t + U_z h_{t-1} + b_z)$$

Again, $W_z$, $U_z$, and $b_z$ are learnable parameters, and $\sigma$ is the sigmoid function. A value of $z_t$ close to 1 means the previous state $h_{t-1}$ is largely kept, while a value close to 0 means the new candidate state is predominantly used.

Candidate Hidden State

Before computing the final hidden state, the GRU calculates a candidate hidden state ($\tilde{h}_t$). This calculation is influenced by the reset gate, which determines how much the previous hidden state $h_{t-1}$ contributes:

$$\tilde{h}_t = \tanh(W_h x_t + U_h (r_t \odot h_{t-1}) + b_h)$$

Here, $\odot$ denotes element-wise multiplication (the Hadamard product). If the reset gate $r_t$ has values close to 0, the contribution from $h_{t-1}$ is effectively erased, allowing the candidate state to be based primarily on the current input $x_t$. $W_h$, $U_h$, and $b_h$ are another set of learnable weights and bias. The $\tanh$ function helps regulate the values within the network, typically squashing them between -1 and 1.

Final Hidden State

Finally, the update gate $z_t$ mediates between the previous hidden state $h_{t-1}$ and the candidate hidden state $\tilde{h}_t$ to produce the final hidden state $h_t$ for the current time step:

$$h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t$$

This equation acts like a weighted average. If $z_t$ is close to 1, the candidate state $\tilde{h}_t$ contributes more, effectively updating the hidden state with new information. If $z_t$ is close to 0, the previous hidden state $h_{t-1}$ is retained more, allowing information to pass through unchanged across multiple time steps. This mechanism is how GRUs can maintain long-range dependencies.

A simplified view of information flow within a GRU cell. $x_t$ is the input, $h_{t-1}$ is the previous hidden state. The reset gate ($r_t$) influences the candidate state ($\tilde{h}_t$), and the update gate ($z_t$) combines the candidate state with the previous state to produce the final hidden state $h_t$.

GRUs vs. LSTMs

GRUs are often seen as a more streamlined alternative to LSTMs.

• Fewer Gates: GRUs have two gates (reset, update) compared to LSTM's three (input, forget, output).
• No Separate Cell State: GRUs don't maintain a distinct cell state ($c_t$) like LSTMs; they directly modify the hidden state ($h_t$).
• Efficiency: Fewer parameters and operations generally mean GRUs can be slightly faster to train and require less memory than LSTMs of similar size.

In practice, the choice between LSTM and GRU often depends on the specific dataset and task. Neither consistently outperforms the other across all scenarios, although GRUs have gained popularity due to their relative simplicity and comparable performance.

Here's a PyTorch snippet showing the core calculations for a single GRU step (assuming inputs x_t, h_tm1 and pre-defined weight/bias tensors):

import torch
import torch.nn.functional as F

# Example tensors (batch_size, input_size/hidden_size)
# Replace with actual dimensions and initialized weights/biases
batch_size = 1
input_size = 10
hidden_size = 20

x_t = torch.randn(batch_size, input_size)
h_tm1 = torch.randn(batch_size, hidden_size) # h_{t-1}

# --- Assume weight matrices (W_*, U_*) and biases (b_*) are defined ---
# Example initialization (replace with actual learned parameters)
W_r = torch.randn(input_size, hidden_size)
U_r = torch.randn(hidden_size, hidden_size)
b_r = torch.randn(hidden_size)

W_z = torch.randn(input_size, hidden_size)
U_z = torch.randn(hidden_size, hidden_size)
b_z = torch.randn(hidden_size)

W_h = torch.randn(input_size, hidden_size)
U_h = torch.randn(hidden_size, hidden_size)
b_h = torch.randn(hidden_size)
# ---------------------------------------------------------------------

# Reset Gate calculation
r_t = torch.sigmoid(x_t @ W_r + h_tm1 @ U_r + b_r)

# Update Gate calculation
z_t = torch.sigmoid(x_t @ W_z + h_tm1 @ U_z + b_z)

# Candidate Hidden State calculation
h_tilde_t = torch.tanh(x_t @ W_h + (r_t * h_tm1) @ U_h + b_h)

# Final Hidden State calculation
h_t = (1 - z_t) * h_tm1 + z_t * h_tilde_t

print("Previous hidden state shape:", h_tm1.shape)
print("Current hidden state shape:", h_t.shape)

This simplified structure, while effective, still relies on sequential processing. The computation for time step $t$ depends on the result from time step $t-1$. This inherent sequential dependency limits parallelization during training and remains a bottleneck for processing very long sequences, setting the stage for the non-recurrent attention mechanisms used in Transformers.