Parallelized GRUs

The setup: three papers, one question

This talk is the measured version of that argument — one file, three probe tasks, every claim multi-seed. (TC⁰: a class of shallow parallel circuits, properly defined a few slides in — until then, read it as "provably limited.")

Sixty seconds of GRU

Two gates, both reading the previous state, blending old memory with a candidate:

$$z_t = \sigma(W_z x_t + U_z h_{t-1}) \;\;\text{(update: how much to overwrite)} \qquad r_t = \sigma(W_r x_t + U_r h_{t-1}) \;\;\text{(reset: how much memory to consult)}$$

State-dependent gating = the power and the sequential tax. (This talk calls its replacement a mixer.)

The problem, and the trick

minGRU (Feng et al., arXiv:2410.01201): make gates depend only on $x_t$ → the recurrence becomes linear in the state: $h_t = a_t \odot h_{t-1} + b_t$

The algebra: delete the state from the gates

GRU — every gate reads the state:

$$z_t = \sigma(W_z x_t + {\color{#b91c1c}\cancel{U_z h_{t-1}}}) \qquad r_t = \sigma(W_r x_t + {\color{#b91c1c}\cancel{U_r h_{t-1}}}) \qquad \tilde h_t = \tanh(W_h x_t + {\color{#b91c1c}\cancel{U_h (r_t \odot h_{t-1})}})$$

$h_t = (1-z_t)\odot h_{t-1} + z_t \odot \tilde h_t$ is nonlinear in $h_{t-1}$ → no composition rule → sequential.

minGRU — delete the red terms; the reset gate $r_t$ dies with them:

$$z_t = \sigma(W_z x_t), \quad \tilde h_t = W_h x_t \;\;\Longrightarrow\;\; h_t = \underbrace{(1-z_t)}_{a_t} \odot\, h_{t-1} \;+\; \underbrace{z_t \odot \tilde h_t}_{b_t}$$

Same update shape — but now affine in $h_{t-1}$, with $(a_t, b_t)$ computable from $x_t$ alone. Affine maps compose associatively → the scan (next slide). Still nonlinear in the inputs; linearity in the state is the only linearity the scan needs.

What is an associative scan? Start with a running sum

A scan turns a sequence into all its running combinations. If the combine operation is associative — $(x \circ y) \circ z = x \circ (y \circ z)$ — you can group work freely and compute every prefix in $\log_2 T$ doubling rounds:

Our combine op: composing affine maps

Same lattice, richer $\circ$ — each element is the affine step $(a_t, b_t)$, composed as $(a_2, b_2) \circ (a_1, b_1) = (a_2 a_1,\; a_2 b_1 + b_2)$ (apply earlier, then later):

step() still exists for O(1)-memory streaming inference — training never runs sequentially.

Interlude: groups, S₃, and "word problems"

A group = reversible operations closed under composition; Sₙ = all n! permutations of n objects. The word problem: the running product of a stream — our streaming-label setup.

S₃, S₄ are solvable — buildable from a tower of commutative layers; S₅ is not (Abel–Ruffini): the cliff two slides ahead. (Parity = the word problem of the 2-element group.)

What "constant-depth circuit" means

Picture computation as gates stacked in layers: inputs at the bottom, answer at the top. Depth = the longest chain of dependencies — the part you cannot parallelize.

TC⁰ = computable at constant depth with counting gates. The theorem (not obvious!): a fixed scan stack's function always admits such a circuit; an RNN's chain is the depth.

The price: expressivity (where Mamba lives too)

A fixed stack of parallel scans is a constant-depth parallel circuit — the class TC⁰. The transition is chosen before reading memory, and stacking layers doesn't change that:

S3 (solvable) sits inside TC⁰ — the ladder ahead is about which parameterizations realize and train it. The question: how much GRU-ness can you buy back inside the scan?

Two probe tasks, one protocol

Multi-seed means, ≥3 seeds per claim (a seed = one independent run: fresh init + data order). Rotation-family rows select the best checkpoint by val@128: longer than train (64), so it rewards generalizing solutions — shorter than every reported test length, so selection can't leak.

The data, verbatim — what a training sample actually looks like

These are real tensors from seed 0's first batch — integer ids in, integer ids out, a label at every step. Nothing else reaches the model.

Rung 1 — base minGRU (mixer="log")

Log-space scan → coefficients are positive: $a_t \in (0,1)$. The state can only decay and blend — it cannot flip, cannot rotate.

Result: chance on both tasks, at any depth.

The scan is fine. The coefficient family is the cage — a parameterization limit, no amount of data fixes it.

Rung 2 — SignedMinGRU (mixer="signed")

Linear-space scan, transition goes negative:

$$a_t \in (-1, 1) \;\text{ via } \tanh$$

• Parity: solved, held at length — the $\tanh$ asymptote is an attractor at exactly −1: nearby values get pulled back, so errors self-correct instead of accumulating
• S3: fits @64 (0.999), but the diagonal shortcut decays with length

Length generalization comes from attractors at exact solutions, not from fit accuracy.

Why diagonal can never track S3

Recovery: matrix transitions. $h_t = M_t h_{t-1} + b_t$ — same doubling scan, but composition is matrix multiplication (matrix_scan).

How a rotating state tracks the group

The dial's final pose is the running group element — a linear readout reads the label. This is exactly RotationMinGRU's transition, $M_t = R(\theta_t)\,\mathrm{diag}(1, \tanh u_t)$: the rotation supplies the turn, the $\tanh$ factor (eigenvalue → −1) supplies the flip.

Rung 3 — RotationMinGRU (mixer="rotation")

2×2 blocks: $M_t = R(\theta_t)\,\mathrm{diag}(1, \tanh u_t)$. S3 ≅ D3, the triangle's turns + mirrors — and D3 embeds in O(2), the plane's rotations + reflections → one layer can be exact.

Representable ≠ trainable — two protocol pieces:

1. Snap (straight-through estimator): forward uses $\theta$ quantized to an exact group angle; the gradient flows through the soft $\theta$ — this manufactures the attractor
2. Best-val@128 checkpoint: the exact solution is visited, not stable — keep the best visitor

0.958 mean @16× (n=8); 1 of 8 seeds lands the exact automaton. L=1 only: stacking snapped layers compounds the snap discontinuity (the code warns).

The ladder so far

One interface: MinGRUBlock(d_model, mixer=...) — stackable, mixable per layer. (Impatient for "which one do I use"? — a practical chooser table closes the deck.)

Task three: S3-hier — the operation hides in a pair

A Latin square (think Sudoku row/column rule): every operation appears exactly once per row and column — so the distribution of ops given one symbol is uniform. No shortcut exists.

S3-hier, verbatim — same stream, harder labels

The harder problem: extract then compose

S3-hier needs two jobs at once — and each half fails alone:

• L=1 rotation fails (nothing extracts)
• Homogeneous signed fails (nothing composes)
• Heterogeneous stack = right architecture...

...but with a continuous 2D composer: 1 of 12 seeds fits.

The bottleneck moved from expressivity to trainability.

Expressivity: do weights exist? (theorems.) Trainability: does SGD find them? (only seeds answer.)

What is a Givens rotation?

The smallest possible rotation: pick one coordinate plane, rotate by θ, leave every other direction exactly alone.

Rung 3½ — GivensMinGRU (mixer="givens")

Compose the atoms — three brick-wall rounds per 8-dim block, angles from one linear head:

$M_t = G_2 G_1 G_0 \in SO(8)$ — the pure rotations of 8-dim space (det +1, no reflections); continuous, no snap, same parallel scan. (Not from the three papers — built here from the 1-of-12 finding: DeltaProduct's products-of-atoms move on the rotation atom; brick-wall mesh from EUNN — Jing et al. 2017 — and Clements et al. 2016.)

Givens is the small-state group-composition specialist — its edge is map-richness at a matched state, not a ladder summit. The broad cross-task workhorse is the delta rule (next, and the benchmark round ahead).

The other lineage: the delta rule (Householder maps)

Numerical linear algebra has two orthogonal atoms — Givens rotations and Householder reflections. The second powers the DeltaNet family (β = the token's write strength):

Different lever: delta's dial is state size and Householder count (nh), not the rotation family's map-richness. Held to a matched 64-element state it trails Givens on S3-hier; grown to its native state — and across the public benchmark suite (the cross-task round ahead) — the delta family is the broad workhorse, while Givens stays the small-state specialist.

Four ways to move a state — one toy input

State tracking = making the composition of these maps mirror the group. Exactness, richness, mirror-atoms, and storage are four different bets on how to survive T steps.

Map richness is the trainability lever

Same task, same 64-element state, same protocol, same ["signed", ·] stack — only the composer changes. (Middle bar: the delta rule, held to the matched 64-element state.) The significant bracket is within-family (2D→8D rotation, p ≈ 0.009). Givens-vs-delta is suggestive only — p ≈ 0.22, params unmatched (14.6k vs 3.3k): not a like-for-like race. [Superseded at n=36 on GPU: givens 25/36 vs delta 10/36, p=0.00084 — see "Choosing a mixer."]

The trade Givens does not escape

No attractor → $\varepsilon$-sized angle errors compound as $\varepsilon T$. Reliability vs. exactness — you currently pick one.

Cost honesty

The case for the parallel composer was never CPU speed: $O(\log T)$ depth on parallel hardware, 8/12 vs 1/12 fit reliability, best fit-only length-gen profile. Update: the packaged chunked-WY forward changes the CPU number too — 0.0577 s vs eager Givens' 0.9493 s at this shape (~16×), pinned CPU stratum. GPU is a different story — next slide.

On GPU: performance and speed

Cost is at parity across all three arms (~15–18 s/seed) — the CPU stratum's 16× cost gap for Givens disappears once it runs its Triton scan path, and delta's torch.compile path is no cheaper at these shapes. Reliability still separates the mixers: delta at native state remains the most reliable composer measured (35/36); at matched 64-element state, givens is now decisively more reliable (25/36 vs 10/36, Fisher exact $p = 0.00084$).

Delta runs via torch.compile, not a shipped Triton kernel — a fusion-headroom probe found torch.compile recovers 70–91% of the eager-vs-floor headroom, so a delta Triton kernel wasn't judged worth building yet.

torch 2.8 / NVIDIA L4 / Triton 3.4 — a separate stratum, not comparable to this deck's pinned-CPU numbers.

Zooming out: three public benchmarks

Past S3-hier, the accepted benchmark round runs the whole mixer family on three public tasks. The pattern is one-sided — the delta family is the broad workhorse; Givens wins none of them:

Delta's dial is state size and Householder count (nh): a $k$-cycle needs $k-1$ reflections, so the matched nh=2 arm is 0/36 by design, and nh=4 is what first gets real traction on S5 (7/36 vs 1/36 — still partial, as theory predicts for this cliff). Givens' win stays real but narrow — small-state group-composition richness (S3-hier, 25/36 vs 10/36). Two dials, not one winner.

torch 2.8 / NVIDIA L4 / Triton 3.4 — a separate stratum, not comparable to this deck's pinned-CPU or S3-hier numbers.

Time-aware decay

Real event streams are irregularly spaced — a 3-second gap and a 3-day gap should not move memory equally. All four mixers accept an optional decay channel:

Decay: the contract

• decay="fixed" — $\lambda$ = decay_rate, a buffer, zero new parameters
• decay="learnable" — $\lambda = \mathrm{softplus}(\rho)$, a smooth ramp keeping $\lambda > 0$: per channel (log/signed), per block (rotation/givens)
• Pairing fails loud, both directions: decay enabled ⟺ delta_t provided, else ValueError
• log1p_delta=True compresses wide gap ranges before scaling

Initialization matters (measured): start near persist-by-default (decay_rate=0.05, $\gamma \approx 1$) and let gradients pull decay in where time pays. Started at heavy decay, the mechanism barely moved; with the small init it wins consistently on session-boundary tasks.

Choosing a mixer

Promoted default: delta at native/full state — most reliable trained composer measured (GPU: 35/36 fits; CPU: 12/12), cost flat in state size. Two axes still move you off the default:

GPU numbers: torch 2.8, NVIDIA L4, Triton 3.4 stratum — not comparable to this deck's pinned-CPU numbers elsewhere. Cost is at parity across arms on that stratum (~15–18 s/seed); the CPU-side chunked-WY speed advantage does not carry over.

Choosing a mixer — in code

stack = MinGRUStack(32, 256, n_layers=2, mixer=["signed", "givens"])
y, state = stack(x)                     # parallel training path
y_t, state = stack.step(x_t, state)     # O(1) streaming inference

# promoted default: full-state delta, single mixer
stack = MinGRUStack(32, 256, n_layers=2, mixer="delta")

Design lessons worth stealing

1. The scan is never the problem — the coefficient family is.
2. Length generalization = attractors at exact solutions (tanh asymptote, snap grid — or drift).
3. Commutativity is a cliff, not a slope.
4. Expressivity ≠ trainability — same automata, 1/12 vs 8/12 training odds.
5. Scalar decay composes cleanly with orthogonal structure.
6. Verify like you mean it — promoted mixers are bit-identical to their evidence code; one recorded run replicated exactly through the public path before merge.

Reproduce everything

# full self-test suite (algebra, scan≡sequential, forward≡step, decay contracts)
uv run --python 3.12 --with 'torch==2.5.1' python min_gru.py

# any single cell: task × model  (CKPT=1 = the best-val@128 checkpoint protocol)
uv run --python 3.12 --with 'torch==2.5.1' python probes.py parity minGRU-signed
CKPT=1 uv run --python 3.12 --with 'torch==2.5.1' python probes.py S3-hier minGRU-hetero-sg8

# the whole grid
uv run --python 3.12 --with 'torch==2.5.1' python probes.py all

• Repo: github.com/chris-santiago/minGRU — the README carries every number here, multi-seed and protocol-pinned
• minGRU: Feng et al., Were RNNs All We Needed? (arXiv:2410.01201) · expressivity framing: Merrill, Petty & Sabharwal (2024) · negative eigenvalues + DeltaProduct: Grazzi et al. (2025 — two separate papers) · brick-wall Givens mesh: EUNN (Jing et al. 2017), Clements et al. (2016)

Questions → the README's "ladder" section is this deck in prose.