Minimal GRU (MinGRU) - A simplified GRU with a single gate.

Implements the MinGRU from "Were RNNs All We Needed?" (Feng et al., 2024). MinGRU strips the GRU down to its essential component: a single forget/update gate. This makes it parallel-scannable during training while preserving the core gating mechanism that makes GRUs effective.

Key Innovations

• Single gate: Only one gate z_t controls interpolation (vs 3 in standard GRU)
• No hidden-to-hidden: Gate depends only on input, not previous hidden state
• Parallel scannable: The simplified recurrence admits a parallel prefix scan
• ~30 lines of core logic: Drastically simpler than standard GRU

Equations

z_t = sigmoid(linear_z(x_t))           # Update gate (input-only)
candidate_t = linear_h(x_t)            # Candidate (no hidden dependency)
h_t = (1 - z_t) * h_{t-1} + z_t * candidate_t  # Interpolation

Architecture

Input [batch, seq_len, embed_dim]
      |
      v
[Input Projection] -> hidden_size
      |
      v
+---------------------------+
|     MinGRU Layer          |
|  z = sigmoid(W_z * x)    |
|  c = W_h * x             |
|  h = (1-z)*h + z*c       |
+---------------------------+
      | (repeat num_layers)
      v
[Layer Norm] -> [Last Timestep]
      |
      v
Output [batch, hidden_size]

Usage

model = MinGRU.build(
  embed_dim: 287,
  hidden_size: 256,
  num_layers: 4,
  dropout: 0.1
)

References

• Paper: https://arxiv.org/abs/2410.01201