Slowrun and Gated Delta Net

# Slowrun and Gated Delta Net 
## Table of Contents

1. Introduction
2. Experiments: ProRes, Attention Residuals, and NOBLE
3. The Attention Mechanism Landscape
4. Why Hybrid Attention Matters for Data-Constrained Training
5. Softmax and Linear Attention: Strengths, Weaknesses, and the Accumulation Problem
6. The Delta Rule and State Tracking with Negative Eigenvalues
7. Gated Delta Networks: Architecture and Initialization
8. Hybrid Attention: Empirical Results and the Compute-Quality Frontier
9. Results Summary
10. References

## Introduction
This is a follow up post from my last [post](https://hackmd.io/@l_WDq7lkQq29Pz-KD1JPNA/H1kVdYdtbe), where I document my experience speedrunning the slowrun. 
TL;DR, slowrun is how low of a val loss can you extract out in data constrained environment (100M Tokens)

## My Experiments since last time: 

Most of my experiments this time have been related to implementing papers and seeing how they fare on the challenge. I also did multiple sweeps, just to see what is the best possible config, in terms of both time and val loss. I tried a bunch of papers like ProRes, Attention Residuals, Noble, and Gated DeltaNet, and found out that the Gated DeltaNet is a promising path to improve data efficiency.
> Note: Baseline for below experiments was 3.365

### Experiment: [ProRes](https://arxiv.org/abs/2603.05369)

Progressive Reisdual warmup(ProRes) is a technique that enables the "Early Layers learn first" by having a warmup on the layer connections defined by:
$$
x_{l+1} = x_l + \alpha(l,t)\cdot \mathcal{F}(\mathrm{Norm}(x_l)).
\
$$

Here, $\alpha(l,t) \in \mathbb{R}$ is a predefined scalar that depends on the layer index $l$ and the training step $t$. An example schedule is:

$$
\alpha(l,t) = \min\left(\frac{t}{T \times l}, 1\right), \qquad l = 1, \ldots, L
\
$$

Essentially this translates to the later layers would start getting information and contributing to learning once the earlier layers are stable. 

They claim strong results: <div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-1"></div>

however on slowrun, they did not fare so well at the time: 
Run : Residual Warmup (ProRes, warmup_frac=0.25)
- **Val loss: 3.38366** INCONCLUSIVE (+0.019 vs baseline, within seed variance)
- **Train loss: 3.000** (vs baseline ~3.060 — noticeably better)
- **Hypothesis:** Progressive residual warmup stabilises early training by gradually enabling deeper layers. α(l,t) = min((t+1)/(T×l), 1) with T = 0.25 × num_iterations / 16 ≈ 48 steps. Layer 1 warms up by step 48 (1.6%), layer 16 by step 763 (25%).
- **Result:** Val loss 3.38366 — nominally 0.001 below baseline, but within the ±0.002 seed variance band.
- **Why it likely didn't help:**
  1. **Zero-init projections already achieve the same effect.** Both `c_proj` (attn) and `c_proj` (MLP) are zero-init, meaning sublayer outputs start at zero regardless of α. ProRes is redundant with zero-init — the model already starts as a shallow (identity residual) network and gradually learns deeper representations as projection weights grow from zero.
  2. **Interaction with `resid_lambdas=1.1`.** The residual amplification already biases early layers to dominate. ProRes suppresses deep layers further, but those layers were already contributing little at init due to zero-init projections.
  3. **Train loss improved (+0.060) but val loss didn't.** This suggests ProRes may slightly improve optimization (smoother early gradients) but the benefit doesn't transfer to generalisation, the model already converges well without it.
  4. **warmdown-ratio=0.6 confound.** This run used 0.6 warmdown (not the baseline's 0.5), which was previously shown to be slightly worse (Run 5). The ProRes benefit may have been cancelled by the suboptimal warmdown ratio.

### Experiment: [Attention Residuals](https://arxiv.org/abs/2603.15031)

Attention Residuals provide a unified view of time and depth by treating residual connections across layers as an attention mechanism over previous layer states. Instead of using a fixed additive skip connection, each layer selectively aggregates information from earlier layers using learned, layer-specific attention weights.

The core formulation is:

$$
h_l = \alpha_{0 \to l}\cdot h_1 + \sum_{i=1}^{l-1} \alpha_{i \to l}\cdot f_i(h_i).
\
$$

Here, $\alpha_{i \to l}$ are layer-specific attention weights satisfying:

$$
\sum_{i=0}^{l-1} \alpha_{i \to l} = 1.
$$

This means the input to layer $l$ is a weighted combination of the initial representation and transformed outputs from all preceding layers. Since network depth is typically much smaller than sequence length, applying attention over depth remains computationally feasible.

#### Full Attention Residuals

The attention weights are defined through a kernel function $\phi$ as:

$$
\alpha_{i \to l} = \frac{\phi(q_l, k_i)}{\sum_{j=0}^{l-1} \phi(q_l, k_j)}.
\tag{2}
$$

For each layer $l$, the query and keys/values are defined as:

$$
q_l = w_l,\qquad
k_i = v_i =
\begin{cases}
h_1, & i = 0 \\
f_i(h_i), & 1 \leq i \leq l-1
\end{cases}
\tag{3}
$$

The input to layer $l$ then becomes:

$$
h_l = \sum_{i=0}^{l-1} \alpha_{i \to l}\cdot v_i.
\tag{4}
$$

In this formulation, $q_l = w_l$ is a layer-specific learnable query vector. The kernel $\phi$ determines how strongly each previous layer contributes to the current layer input. A common choice is a softmax-style kernel, which gives normalized attention over depth.

This mechanism allows each layer to dynamically retrieve information from earlier layers instead of relying on a uniform residual path. In practice, this gives the model a more flexible way to control information flow across depth.
<div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-2"></div>


Block Attention Residuals reduce the cost of attention over depth by partitioning the \(L\) layers into \(N\) blocks. Instead of attending over every previous layer output individually, the model first compresses each block into a single representation and then applies attention only over these block-level summaries. This lowers memory and communication overhead from \(O(Ld)\) to \(O(Nd)\).

Within each block B_n, the block representation is formed by summing the outputs of the layers in that block:

$$
b_n = \sum_{j \in \mathcal{B}_n} f_j(h_j).
\tag{5}
$$

Inter-block attention then uses these block representations, together with the token embedding, as the sources for the current layer. So instead of retrieving information from all prior layers one by one, the model retrieves from a much smaller set of summarized block states.

This gives a practical tradeoff between expressiveness and efficiency. Full Attention Residuals are more flexible, but Block Attention Residuals preserve most of the benefit while making computation and memory much cheaper, especially at larger depths. I did experiments with both full attention residuals and block attention residuals, with mixed results. 

- **Val loss: 3.596**  (worse by +0.231 vs baseline 3.365)
- **Train loss: 2.980**| **Wall time: 21.2m** (way over 15m budget)
- **Changes:**
  - Block AttnRes added as supplement (preserving baseline resid_lambdas + x0_lambdas)
  - Per-layer `attn_res_projs` (pseudo-query vectors, zero-init) + `attn_res_lambdas` (blend gates, zero-init)
  - Block boundaries every 4 layers → softmax attention over block-level representations
  - Formula: `V = stack(block_reps + [x])`, `h = softmax_attn(V)`, `x = x + λ_i * (h - x)`
  - Zero-init should make model start as pure baseline
- **Failure analysis:**
  1. **MFU cratered (−8pp):** `torch.stack` + `torch.einsum` on variable-length tensors broke torch.compile graph caching. Each block boundary creates a different stack size, forcing recompilation.
  2. **Wall time +40%:** Extra ops per layer (stack, einsum, softmax) add ~85ms/step even with small stacks.
  3. **Val loss:** Despite zero-init lambdas, the attn_res_lambdas learned nonzero values early and destabilized the residual stream. The softmax attention over block reps diluted the most-recent-layer signal that resid_lambdas=1.1 amplifies.
  4. **Interaction with torch.compile:** The dynamic stack sizes likely prevented effective fusion/optimization.


### Experiment: [NOBLE](https://arxiv.org/abs/2603.06492)

NOBLE(Non-linear low rank branch for linear enhancement)  augments a standard linear layer with a low-rank nonlinear branch, aiming to make parameter-efficient adaptation more expressive than standard LoRA-style linear low-rank updates. Instead of only adding a linear low-rank bypass, NOBLE introduces a nonlinear transformation in the bottleneck space.

The modified layer is defined as:

$$
f_{\mathrm{NOBLE}}(x) = xW + b + \sigma(xW_{\mathrm{down}})W_{\mathrm{up}}
$$

Here, $W_{\mathrm{down}} \in \mathbb{R}^{d_{\mathrm{in}} \times r}$, $W_{\mathrm{up}} \in \mathbb{R}^{r \times d_{\mathrm{out}}}$, and $r \ll \min(d_{\mathrm{in}}, d_{\mathrm{out}})$ is the bottleneck rank. The key idea is that the low-rank branch is no longer purely linear; instead, it can model richer feature transformations through the nonlinearity $\sigma$.

#### CosNet: Recommended Nonlinearity

The paper recommends **CosNet**, a cosine-based bottleneck nonlinearity:

$$
\sigma_{\cos}(h) = \cos\!\big(\omega \odot (M \cdot \cos(\omega_1 \odot h + \phi_1)) + \phi_2\big)
$$

Here, $h \in \mathbb{R}^r$ is the bottleneck representation, $M \in \mathbb{R}^{r \times r}$ is a learned mixing matrix, and the cosine frequencies and phases are learnable. This gives the low-rank branch a bounded, smooth, and periodic activation structure, which the paper argues is better suited than standard activations like ReLU or GELU in this setting.

#### Key Design Choices

NOBLE uses near-zero initialization for $W_{\mathrm{up}}$, so the nonlinear branch contributes very little at initialization and training starts close to the original pretrained model. It also slightly reduces the initialization scale of the main weight matrix $W$, leaving more room for the nonlinear low-rank branch to become useful during adaptation.

The optimizer setup is also adjusted: elevated learning rates are used for $W_{\mathrm{up}}$ and the CosNet mixing parameters, while $W_{\mathrm{down}}$ stays at the base learning rate. This reflects the fact that different parts of the low-rank branch need different update scales to train effectively.

#### Computational Overhead

Compared to standard low-rank adaptation, NOBLE introduces a modest increase in parameters and step time due to the nonlinear branch and mixing parameters. But the intended tradeoff is better adaptation quality and fewer training steps, making the extra per-step cost worthwhile when the nonlinear bottleneck improves convergence.

I performed a sweep for this to check the best possible config: 
NOBLE Sweep (8 runs)
- **Baseline:** val loss 3.3539 in this sweep's baseline run
- **Hypothesis:** Low-rank nonlinear branches on MLP projections capture input-dependent modulation that linear projections miss. In the multi-epoch regime, this richer function class should help generalization by reducing representational bottlenecks.
- **Sweep design:** 8 configs varying rank (r=16, 32), target (proj, mlp, all), alpha (0.005, 0.01), gamma (0.3, 0.5)
- **Results:**

| Run | Config | Val Loss | Δ vs Baseline | Train Time | Overhead |
|-----|--------|----------|---------------|------------|----------|
| 00 | Baseline (no NOBLE) | 3.365 | — | 14.50m | — |
| 01 | r16, proj | 3.3615 | −0.0035 | 15.62m | +7.7% |
| 02 | r32, proj | 3.3569 | −0.0081 | 15.66m | +8.0% |
| **03** | **r16, mlp** | **3.3495** | **−0.0155** | **16.78m** | **+15.7%** |
| 04 | r32, mlp | 3.3507 | −0.0143 | 16.84m | +16.1% |
| 05 | r16, proj, α=0.005 | 3.3611 | −0.0039 | 15.61m | +7.7% |
| 06 | r16, proj, γ=0.5 | 3.3620 | −0.0030 | 15.62m | +7.7% |
| 07 | r16, all | 3.3523 | −0.0127 | 19.27m | +32.9% |

**Key findings:**
  1. **MLP-targeted NOBLE wins (−0.0155), proj-only shows marginal gains (−0.003 to −0.008).** The zero-init on c_proj means the NOBLE branch's near-zero init adds signal to an already-zero output, creating a small but persistent perturbation. MLP gate/fc projections don't have zero-init, so the NOBLE branch integrates cleanly.
  2. **r16 > r32** for MLP target: lower rank provides better implicit regularization in the data-limited regime.
  3. **All-projections (Q/K/V/c_proj/gate/fc) is marginal (−0.0127) at huge cost (+33%).** The Q/K/V branches don't help — attention already has the gate mechanism for modulation.
  4. **Wall time is the primary blocker.** Best config (r16_mlp) takes 16.78m, exceeding the 15-min tiny track budget by 12%.
  
Apart from this I did multiple sweeps(layer duping, warmdown schedules) results of which were a mixed bag. 

At this time I came across the [Olmo Hybrid Paper](https://allenai.org/papers/olmo-hybrid) and this tweet: 
<div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-3"></div>
and this essentially led me to wonder if the problem was architectural, which eventually got me to the hybrid attention and gated delta net rabbit hole. 

## Hybrid Attention:
Hybrid Attention refers to using both linear attention and softmax attention in the same model. In the version that I tested on slowrun, we use a particular variant for alternating known as the gated delta net with negative eigenvalues. 
### The Attention Mechanism Landscape and data efficiency

| Mechanism | Complexity | State Size | Can Forget | Can Correct | Can Track State | Best For |
|-----------|-----------|------------|------------|-------------|-----------------|----------|
| Softmax | $O(T^2 d)$ | $O(Td)$ KV cache | N/A | N/A | Via patterns | Precise retrieval, deep reasoning |
| Linear Attn | $O(Td^2)$ | $O(d_k d_v)$ | ❌ | ❌ | ❌ | — (too weak in practice) |
| Gated Linear (Mamba) | $O(Td^2)$ | $O(d_k d_v)$ | ✅ | ❌ | ❌ | Long-range dependency, fast inference |
| Delta Rule | $O(Td^2)$ | $O(d_k d_v)$ | ❌ | ✅ | ❌ | Associative storage |
| **GDN** | $O(Td^2)$ | $O(d_k d_v)$ | ✅ | ✅ | ✅ | **Everything except exact retrieval** |
| **Hybrid (Softmax + GDN)** | Mixed | Mixed | ✅ | ✅ | ✅ | **Optimal trade-off** |

We will walk through all the terms in the table, and motivation for this architectural change in further sections. 

## Why This Matters for the Slowrun Regime

The Slowrun benchmark is an extreme case of the data-limited, multi-epoch regime:

- **~100M unique tokens, 12–16 epochs** — each sequence is seen over a dozen times
- **Tight compute budgets** — 15 min (Tiny), 60 min (Limited)
- **Fixed architecture size** — 16 layers / 1024d (Tiny), 30 layers / 1792d (Limited)

In this regime, the data efficiency advantage of hybrid attention is **amplified**:

1. **Each epoch of repetition benefits the recurrent state more than attention.** The delta rule's error-correction mechanism means repeated sequences refine the recurrent state's associations ($v_t - h_{t-1}^\top k_t$ shrinks with repeated exposure).

2. **The small token budget makes the inductive bias argument strongest.** With only 100M unique tokens, the generalization gap dominates. In our regime where we have very few tokens to start with, the relative advantage should be even larger.

3. **The tight compute budget penalizes O(T²) redundancy.** Every softmax layer that could be replaced by a faster GDN layer frees compute for more training steps or larger batch sizes. 

4. **Multi-epoch training is implicit fine-tuning.** The model is effectively fine-tuning on its own pretraining data through repeated exposure. 

Let's now get into the detailing of the table, and the different attention mechanisms:


### Softmax attention
Given an input sequence $X = [x_1, x_2, \ldots, x_T] \in \mathbb{R}^{T \times d}$, softmax attention computes:

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

$$\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$$

For causal (autoregressive) language modelling, the attention matrix is masked so token $t$ can only attend to positions $\leq t$:

$$o_t = \sum_{s=1}^{t} \frac{\exp(q_t^\top k_s / \sqrt{d_k})}{\sum_{j=1}^{t} \exp(q_t^\top k_j / \sqrt{d_k})} \cdot v_s$$

Each output $o_t$ is a **weighted average** of all past value vectors, where the weights are determined by query-key similarity. The softmax normalization ensures the weights sum to 1, creating a probabilistic retrieval over the full history.

#### Strengths

- **Exact retrieval.** Any token can attend to any other token with arbitrary precision. If token 500 needs information from token 3, softmax attention can route it there in a single layer.
- **No information bottleneck.** The $T \times T$ attention matrix preserves all pairwise relationships.
- **Well-understood optimization.** Decades of engineering (Flash Attention, KV caching, GQA,MLA) have made softmax attention extremely efficient in practice.

#### Weaknesses

- **$O(T^2)$ complexity.** Computing the full attention matrix requires $O(T^2 \cdot d)$ FLOPs and $O(T^2)$ memory. For long sequences (T > 8k), this dominates training cost.

- **No persistent state.** Softmax attention recomputes everything from scratch at every token. It cannot incrementally update a belief. Every time a new token arrives, the model must re-examine the entire context.
- **Redundant re-reading.** In multi-epoch training, the model sees the same sequences 12–16 times. Softmax attention processes each occurrence identically, it cannot build on prior exposure because it has no recurrent state.
- **Poor inductive bias for state tracking.** Tracking whether a sentence is negated, whether we're inside a quotation, or whether a conditional clause is still open requires maintaining Boolean state across many tokens. Softmax can only approximate this by learning specific attention patterns.
#### The KV Cache and Its Implications

During autoregressive inference, softmax attention uses a KV cache to avoid re-computing keys and values for past tokens. The cache grows linearly: $O(T \cdot d_k \cdot L)$ for $L$ layers. For a 30-layer model with 14 heads of dimension 128 at sequence length 2048:

$$\text{KV cache} = 2 \times 30 \times 14 \times 128 \times 2048 \times 2 \text{ bytes} \approx 440 \text{ MB}$$

This linear growth in context length is manageable for moderate sequences but becomes prohibitive for long-context applications (100k+ tokens). Linear attention mechanisms, by contrast, maintain a **fixed-size** state regardless of sequence length.

<div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-4"></div>



### Linear Attention

####  The Core Idea

Linear attention replaces the softmax with a feature map $\phi$:

$$o_t = \frac{\sum_{s=1}^{t} \phi(q_t)^\top \phi(k_s) \cdot v_s}{\sum_{s=1}^{t} \phi(q_t)^\top \phi(k_s)}$$

The key insight: by associativity, we can rewrite this as a recurrence over a **state matrix** $h_t \in \mathbb{R}^{d_k \times d_v}$:

$$h_t = h_{t-1} + \phi(k_t) \otimes v_t$$

$$o_t = h_t^\top \phi(q_t) \quad \bigg/ \quad z_t^\top \phi(q_t) \quad \text{where } z_t = z_{t-1} + \phi(k_t)$$

This converts $O(T^2)$ attention into an $O(T)$ recurrence. The state $h_t$ is a compressed summary of all past key-value pairs.

#### The Problem: Naive Accumulation

The basic linear attention update $h_t = h_{t-1} + k_t \otimes v_t$ is a pure accumulator. It can only add information, never subtract/modify it. This creates three fundamental problems:

The update rule $h_t = h_{t-1} + k_t \otimes v_t$ only ever adds to the state, never modifies it. When the model encounters a contradictory update, the old association just sits there, and the state ends up holding both the old and new value superimposed, resulting in conflicting state maps. After thousands of tokens, the state becomes a noisy average of everything seen so far, with no way to surface the relevant signal. Tracking something structural, like whether you're inside a negated clause, requires flipping the polarity of a stored representation, which an accumulator simply cannot do. All three of these failures show up constantly in real language, which explains why early linear attention models so dramatically underperformed softmax attention. 
<div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-5"></div>

These limitations explain why early linear attention models dramatically underperformed softmax attention. The state mechanism was too simple. The solution is to allow the updates to be richer, and should allow the modification of the state easily. 

The delta rule fixes this by only writing the **error**:

$$W_t = W_{t-1} + k_t \otimes \underbrace{(v_t - W_{t-1}^\top k_t)}_{\text{prediction error}}$$

The term $W_{t-1}^\top k_t$ is what the memory **already predicts** for key $k_t$. If the prediction is correct (error = 0), nothing is written. If it's wrong, only the difference is added. This is an online gradient descent step on the squared prediction error $\|v_t - W^\top k_t\|^2$ with learning rate 1.

When applied to sequence modelling, the delta rule transforms the state update into:

$$h_t = h_{t-1} + k_t \otimes \beta_t \cdot (v_t - h_{t-1}^\top k_t)$$

Where $\beta_t \in [0, 1]$ is a token-dependent "write strength" (learning rate). This gives the recurrence three key properties that naive linear attention lacks:

1. **Self-correcting.** If the state already stores the correct value for a key, it's left unchanged. Only new or incorrect information triggers an update.
2. **Interference reduction.** The error signal $v_t - h_{t-1}^\top k_t$ subtracts the contribution of previously stored associations.
3. **Capacity efficiency.** The delta rule achieves near-optimal storage capacity for associative memories. For a state of size $d_k \times d_v$, it can reliably store $\sim d_k$ associations.

The write strength $\beta_t$ controls how aggressively the state is updated. Different $\beta$ values produce qualitatively different dynamics:

| $\beta$ | Eigenvalue $\lambda = 1 - \beta$ | Behavior |
|---------|--------------------------------|----------|
| 0 | 1 | **Read-only**: state unchanged, pure retrieval |
| 0.5 | 0.5 | **Gentle accumulation**: new info blended with old |
| 1.0 | 0 | **Full overwrite**: old value for this key erased, new value written |
| 1.5 | −0.5 | **Sign-flipping**: state polarity inverts along the key direction |
| 2.0 | −1 | **Full sign-flip**: maximum polarity inversion |

The standard delta rule constrains $\beta \in [0, 1]$, which limits the recurrence to non-negative eigenvalues. This is sufficient for associative storage but insufficient for good state tracking. 

### The State-Tracking Problem

Consider parsing the following sentence:

> "The cat that the dog that barked chased ran away."

To correctly assign "ran" to "the cat" (not "the dog"), the model must track a **stack** of nested relative clauses. Each opening clause pushes a state; each verb pops one. This requires:

1. **Persistent state** that survives across many tokens
2. **Selective modification**: push/pop without corrupting other stored states
3. **Polarity awareness**: distinguishing "inside clause" from "outside clause"




#### How Softmax Attention Handles It

Softmax attention can solve state tracking by learning specific attention patterns: "when I see a verb, attend back to the most recent unmatched noun." But this solution is **fragile**:

- It requires the model to dedicate specific attention heads to tracking, reducing capacity for other tasks.
- The pattern must be learned entirely from data there's no structural inductive bias toward stack-like behavior.
- Longer or more deeply nested structures require proportionally more heads and layers, since each head can only implement one simple routing pattern.

Empirically, transformer language models struggle with systematic generalization on state-tracking benchmarks . They can solve specific instances seen in training but fail on longer or differently structured inputs.

#### How Recurrent Models Handle It (Better, With Negative Eigenvalues)

A recurrent model with state $h_t$ can track Boolean state changes if its transition matrix can have **negative eigenvalues**. Consider a simple example: tracking negation.

With a scalar state $s_t$ and transition $s_t = \lambda \cdot s_{t-1} + u_t$:

- If $\lambda > 0$: the state accumulates. Seeing "not" adds a negative signal, but it decays. After enough tokens, the negation is forgotten.
- If $\lambda < 0$: the state **flips sign** at each step. Seeing "not" triggers $\lambda < 0$, which flips the accumulated representation. A second "not" flips it back. The polarity change is **structural**, not dependent on the magnitude of learned weights.
<div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-6"></div>

This is why the extension to $\beta \in [0, 2]$ (allowing $\lambda = 1 - \beta < 0$) is crucial. It gives the recurrence the structural ability to track Boolean state changes without relying on carefully tuned weight magnitudes.
<div class="fig-mount" data-src="/gdn_blog_figures.html" data-fig="fig-7"></div>
### Gated Delta Networks (GDN)

#### The Full Update Rule


[GDN](https://arxiv.org/abs/2412.06464) combines three mechanisms into a single recurrence:

$$h_t = \gamma_t \cdot h_{t-1} + k_t \otimes \left[\beta_t \cdot (v_t - h_{t-1}^\top k_t)\right]$$

$$o_t = h_t^\top q_t$$

Where:
- $h_t \in \mathbb{R}^{d_k \times d_v}$ is the **recurrent state** (associative memory), one per head
- $q_t, k_t \in \mathbb{R}^{d_k}$ are L2-normalized query and key vectors
- $v_t \in \mathbb{R}^{d_v}$ is the value vector
- $\gamma_t = \exp(g_t) \in (0, 1]$ is the **forget gate** (Mamba2-style)
- $\beta_t = \sigma(b_\text{proj}(x_t) + b_\text{bias}) \times 2.0 \in [0, 2]$ is the **write strength**

Each of the three moving parts addresses a problem we've already seen. Working through them in that order:

**The delta rule write** is the fix for naive accumulation. Rather than writing $v_t$ directly into the state, it computes $\Delta_t = \beta_t \cdot (v_t - h_{t-1}^\top k_t)$, where $h_{t-1}^\top k_t$ is what the state already predicts for key $k_t$. Only the prediction error gets written. If the state already has the right answer, nothing changes.

**The forget gate** handles what error-correction alone can't: state saturation over long sequences. Even a perfectly error-correcting state fills up eventually. $\gamma_t = \exp(g_t)$ is always in $(0, 1]$, computed as:

$$g_t = -\exp(A_{\log}) \cdot \text{softplus}(a_\text{proj}(x_t) + dt_\text{bias})$$

When $\gamma_t \approx 1$ the state is fully retained; when $\gamma_t \approx 0$ it's nearly wiped. $A_{\log}$ sets the timescale, larger values mean faster decay and shorter memory. The gate is input-dependent, so the model learns when to hold on and when to let go.

**The write strength $\beta_t$** is what closes the loop on the state-tracking problem. Recall the nested clause example, tracking whether "ran" belongs to "the cat" or "the dog" requires flipping the polarity of a stored representation, not just updating its magnitude. When $\beta_t > 1$, the correction term overshoots, inverting the state's polarity along the key direction. This is the negative eigenvalue mechanism in practice. The standard delta rule caps $\beta$ at 1 and can't do this; GDN extends the range to $[0, 2]$, which is what makes boolean state tracking structurally possible rather than something the model has to approximate through learned weight patterns.

There are two smaller pieces worth noting. Before the recurrence, GDN runs depthwise 1D convolutions (kernel size 4) over q, k, v with SiLU activation:

$$q = \text{SiLU}(\text{Conv1d}(q_\text{proj}(x))), \quad k = \text{SiLU}(\text{Conv1d}(k_\text{proj}(x))), \quad v = \text{SiLU}(\text{Conv1d}(v_\text{proj}(x)))$$

These give the model a local 4-token window before the recurrence handles longer-range dependencies. The paper flags this as load-bearing ("do not turn it off"), though my own experiments found removing it only costs ~0.002 val loss when most layers are softmax. The output is then gated by an input-dependent signal:

$$o_t = \text{RMSNorm}(h_t^\top q_t) \odot \text{SiLU}(g_\text{proj}(x_t))$$

This lets the model suppress the recurrent output entirely when it isn't useful, rather than always passing it forward.


### Expanding the Recurrence: What Each Component Contributes

Writing out the full update and comparing to simpler alternatives:

| Model | Update Rule | Can Forget? | Can Correct? | Can Flip? |
|-------|------------|-------------|-------------|-----------|
| Linear Attention | $h_t = h_{t-1} + k_t \otimes v_t$ | ❌ | ❌ | ❌ |
| Gated Linear (Mamba) | $h_t = \gamma_t h_{t-1} + k_t \otimes v_t$ | ✅ | ❌ | ❌ |
| Delta Rule | $h_t = h_{t-1} + k_t \otimes \beta(v_t - h_{t-1}^\top k_t)$ | ❌ | ✅ | ❌ |
| **GDN ($\beta \leq 1$)** | $h_t = \gamma_t h_{t-1} + k_t \otimes \beta(v_t - h_{t-1}^\top k_t)$ | ✅ | ✅ | ❌ |
| **GDN ($\beta \leq 2$)** | $h_t = \gamma_t h_{t-1} + k_t \otimes \beta(v_t - h_{t-1}^\top k_t)$ | ✅ | ✅ | ✅ |

Each row strictly increases in representational power. The full GDN with $\beta \in [0, 2]$ is the first linear-complexity model that can selectively forget, error-correct, and track Boolean state.

### Complexity Comparison

| | Softmax Attention | GDN (Recurrent) | GDN (Chunk-Parallel) |
|--|------------------|-----------------|---------------------|
| Training FLOPs | $O(T^2 d)$ | $O(T d_k d_v)$ | $O(T \cdot C \cdot d_k d_v)$ |
| Training Memory | $O(T^2)$ | $O(d_k \cdot d_v)$ | $O(T/C \cdot d_k \cdot d_v)$ |
| Inference (per token) | $O(T d)$ | $O(d_k d_v)$ | $O(d_k d_v)$ |
| Inference State | $O(T d)$ KV cache | $O(d_k d_v)$ | $O(d_k d_v)$ |

Where $C$ is the chunk size for the Triton-optimized chunk-parallel algorithm. In practice, $C = 64$ and the chunk-parallel mode achieves near-hardware-peak throughput on modern GPUs.




### Initialization: Why It Matters Enormously

#### The Degenerate $\beta = 1$ Problem

With standard initialization (zero-mean projections, no bias), all heads start at:

$$\beta_\text{init} = \sigma(0) \times 2 = 1.0$$

At $\beta = 1$, the transition eigenvalue is $\lambda_k = \gamma(1 - 1) = 0$. The state is fully overwritten at every token. There is no accumulation, no state-tracking which means that every head behaves identically as a memoryless lookup. The model must discover role specialization purely from gradient signal, wasting early training on learning something that can be easily initialised. 

#### The Dead Head Problem

With $A_{\log}$ initialized from $\log(\text{Uniform}(0, 16))$, roughly 30% of heads get $A > 8$, which gives $\gamma = \exp(-A) < 0.0003$. These heads are basically dead, having zero memory retention. 

### The Fix: Structural Diversity from Step 1

We fix both problems with three small changes (adding only 392 parameters total):

1. **Per-head $\beta$ bias:** $\text{beta_bias} = \text{linspace}(-1.5, 1.5, H)$, giving $\beta_\text{init} \in [0.36, 1.64]$. Half the heads start as accumulators ($\beta < 1$), half as state-trackers ($\beta > 1$).

2. **Tighter $A$ range:** $A \sim \text{Uniform}(2, 8)$, so $\gamma \in [0.0003, 0.14]$ — all heads have meaningful memory retention.

3. **Bias on $b_\text{proj}$:** Matching the NVlabs reference implementation. Allows per-head $\beta$ offset to be learned independently of the input projection.


### Hybrid Attention: The Next Frontier

#### The Core Thesis

Neither pure softmax attention nor pure linear attention is optimal. The evidence from our experiments — and from recent work on OLMo-hybrid (Ai2), Jamba (AI21), Zamba (Zyphra), and Griffin (Google DeepMind), converges on a clear conclusion:

> **The optimal architecture interleaves softmax attention layers with linear attention layers in a fixed ratio.**


#### Why Pure Approaches Fail

**Pure softmax** is inefficient. It recomputes all pairwise relationships at every layer, even though most layers only need local or incremental context. Our experiments show that replacing half the softmax layers with GDN (alternating pattern) **improves** quality by −0.007 val loss. The softmax layers were doing redundant work.

**Pure linear attention** lacks precision. Our all-GDN (30/30 layers) configuration was catastrophically worse (+0.065 val loss). Without any softmax layers, the model cannot perform exact retrieval — the compressed state inevitably loses fine-grained information. Deep layers need to do precise reasoning over specific tokens, which requires the full $O(T^2)$ attention mechanism.

###  The Mathematical Argument for Hybridization

Consider a hybrid model with $L$ layers, of which $L_s$ use softmax and $L_g$ use GDN ($L = L_s + L_g$).

**Softmax layers** provide a function $f_s: \mathbb{R}^{T \times d} \to \mathbb{R}^{T \times d}$ where each output position can depend on **any** input position with **arbitrary** precision. The attention matrix $A \in \mathbb{R}^{T \times T}$ has rank up to $T$ , it's really expressive.

**GDN layers** provide a function $f_g: \mathbb{R}^{T \times d} \to \mathbb{R}^{T \times d}$ where the output at position $t$ depends on positions $\leq t$ through the compressed state $h_t \in \mathbb{R}^{d_k \times d_v}$. The effective "attention matrix" has rank at most $d_k$ , it's a low-rank approximation thus reducing the expressivity. But it runs in $O(T)$ instead of $O(T^2)$.

Not all layers require high expressivity: 

- **Early layers** extract local features (n-grams, syntax patterns). These are well-served by the delta rule's recurrent state, which naturally compresses local patterns into its key-value memory.
- **Middle layers** build intermediate representations through incremental refinement. The delta rule's error-correction helps here as it updates the running representation only when new information contradicts what's already stored.

- **Deep layers** perform global reasoning. These **need** exact retrieval: the model must attend to specific tokens with high precision. Only softmax attention can provide this.

Our experimental data confirms this layerwise division:

| GDN Placement | Val Loss | Interpretation |
|---------------|----------|----------------|
| Encoder-only (layers 0–14) | −0.001 | GDN in early layers ≈ neutral (softmax already efficient here) |
| **Alternating (every other)** | **−0.007** | **Best: each GDN has a softmax neighbor for correction** |
| Decoder-only (layers 15–28) | +0.022 | GDN in deep layers **hurts**  |
| All layers (0–29) | +0.065 | No softmax -> no precision -> catastrophic |

#### The OLMo-Hybrid Pattern

The OLMo-hybrid architecture (Ai2, 2025) formalizes this insight with a **3:1 ratio**: for every 4 layers, 3 use linear attention and 1 uses full softmax attention. The softmax layer is placed at the end of each group, acting as a "correction" step that refines the compressed representations built by the linear layers.

For our 16-layer tiny model, this gives 12 GDN + 4 softmax layers:
- GDN: layers 0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14
- Softmax: layers 3, 7, 11, 15

The final layer (15) is always softmax to ensure the model's output has access to exact full-context retrieval before the prediction head.

## Data Efficiency: Why Hybrids Win in the Multi-Epoch Regime

In data-limited training (our regime: 12–16 epochs over ~100M tokens), the model sees every sequence multiple times. This creates a specific optimization landscape where hybrid attention has a structural advantage:

**Argument 1: Inductive bias reduces hypothesis space.**

Softmax attention is maximally flexible; it can learn any token-to-token routing pattern. But in a data-limited regime, this flexibility is a liability: the model has too many degrees of freedom and can overfit to noise. GDN's recurrent state provides a strong inductive bias towards fixing that. 

**Argument 2: Error correction exploits repetition.**

The delta rule's error-correcting property becomes more valuable with data repetition. On the first pass through a sequence, the recurrent state builds initial associations. On subsequent passes, the prediction error $v_t - h_{t-1}^\top k_t$ is smaller (the state already partially knows the answer), so the update is finer-grained. This is analogous to iterative refinement — each epoch polishes the recurrent state's associations rather than overwriting them. Softmax attention, by contrast, processes each pass the same way it cannot build on prior exposure.

**Argument 3: Negative eigenvalues amortize structural learning.**

In multi-epoch training, the model needs to learn structural patterns (negation, scope, agreement) that are invariant across examples. GDN's negative eigenvalue mechanism provides a structural prior for these patterns: $\beta > 1$ heads are pre-configured for state tracking from initialization.

### Empirical Evidence: Epoch-by-Epoch Dynamics

Our alternating GDN experiments reveal a striking pattern:

| Epoch | Softmax-only | Hybrid (Alternating GDN) | Δ |
|-------|-------------|--------------------------|---|
| 1 | 4.260 | 4.263 | +0.003 (GDN slightly worse) |
| 2 | 3.881 | 3.843 | **−0.038** (GDN surges ahead) |
| 3 | 3.723 | 3.703 | −0.020 |
| 4 | 3.675 | 3.643 | −0.032 |
| 5 | 3.587 | 3.565 | −0.022 |
| 6 | 3.349 | 3.342 | −0.007 |

GDN starts slightly behind (epoch 1) — the recurrent state needs time to learn its associations. But by epoch 2, the error-correction mechanism kicks in and GDN surges ahead by −0.038. This advantage persists through all subsequent epochs but narrows during warmdown (−0.007 at epoch 6), suggesting that the delta rule's learned associations matter less when the learning rate approaches zero.
A more detailed proof on how expressivity of hybrid attention helps in data efficiency and scaling laws can be found in the sections 6B and 6E of the [Olmo Hybrid Paper](https://allenai.org/papers/olmo-hybrid)

Over the full 12-epoch training, GDN achieves 3.2458 best val loss vs 3.2526 for the softmax-only baseline (−0.0068). This is a meaningful gap that would place higher on the leaderboard if it didn't exceed the time limit. 

### The Compute-Quality Frontier

The fundamental tension in hybrid attention is quality vs. cost:

| Architecture | Quality (val loss) | Wall Time | Quality per FLOP |
|-------------|-------------------|-----------|-----------------|
| Pure softmax | 3.2526 | 1.00× | baseline |
| Hybrid Alt14 (conv) | 3.2458 (−0.007) | 1.36× | worse (gain < cost) |
| Hybrid Alt14 (noconv) | ~3.248* | 1.23× | borderline |
| Hybrid Alt8 (noconv) | 3.253 | 1.16× | worse (minimal gain) |

\* Projected from 3-epoch speed data.

## Results Summary Section
### ProRes
| Run | Config                                       | Val Loss | Δ vs Baseline | Interpretation |
| --- | -------------------------------------------- | -------: | ------------: | -------------- |
| 18  | Residual Warmup (ProRes, `warmup_frac=0.25`) |  3.38366 |        +0.019 | Inconclusive   |

### Attention Residuals
| Method                    | Val Loss | Δ vs Baseline | Train Loss | Wall Time | Verdict  |
| ------------------------- | -------: | ------------: | ---------: | --------: | -------- |
| Block Attention Residuals |    3.596 |        +0.231 |      2.980 |     21.2m | Rejected |

### NOBLE Sweep
| Run | Config              | Val Loss | Δ vs Baseline | Train Time | Overhead |
| --- | ------------------- | -------: | ------------: | ---------: | -------: |
| 00  | Baseline (no NOBLE) |    3.365 |             — |     14.50m |        — |
| 01  | r16, proj           |   3.3615 |       −0.0035 |     15.62m |    +7.7% |
| 02  | r32, proj           |   3.3569 |       −0.0081 |     15.66m |    +8.0% |
| 03  | r16, mlp            |   3.3495 |       −0.0155 |     16.78m |   +15.7% |
| 04  | r32, mlp            |   3.3507 |       −0.0143 |     16.84m |   +16.1% |
| 05  | r16, proj, α=0.005  |   3.3611 |       −0.0039 |     15.61m |    +7.7% |
| 06  | r16, proj, γ=0.5    |   3.3620 |       −0.0030 |     15.62m |    +7.7% |
| 07  | r16, all            |   3.3523 |       −0.0127 |     19.27m |   +32.9% |



### Hybrid Attention Placement Study

> **Baseline note:** The GDN experiments were compared against the best val loss in the repository at the time of the run, not the 3.365 baseline above. The deltas here are relative to the softmax-only hybrid baseline (3.2526) and are not directly comparable to the numbers in the experiments above.

| GDN Placement                   | Val Loss Change | Interpretation      |
| ------------------------------- | --------------: | ------------------- |
| Encoder-only (early layers)     |          −0.001 | Nearly neutral      |
| Alternating (every other layer) |          −0.007 | Best configuration  |
| Decoder-only (late layers)      |          +0.022 | Worse than baseline |
| All layers GDN                  |          +0.065 | Strong degradation  |

### Hybrid Attention Epoch-wise Dynamics
| Epoch | Softmax-only | Hybrid (Alternating GDN) |      Δ |
| ----- | -----------: | -----------------------: | -----: |
| 1     |        4.260 |                    4.263 | +0.003 |
| 2     |        3.881 |                    3.843 | −0.038 |
| 3     |        3.723 |                    3.703 | −0.020 |
| 4     |        3.675 |                    3.643 | −0.032 |
| 5     |        3.587 |                    3.565 | −0.022 |
| 6     |        3.349 |                    3.342 | −0.007 |

### Hybrid Attention Compute–Quality Frontier
| Architecture           | Quality (Val Loss) | Wall Time | Quality per FLOP             |
| ---------------------- | -----------------: | --------: | ---------------------------- |
| Pure softmax           |             3.2526 |     1.00× | baseline                     |
| Hybrid Alt14 (conv)    |             3.2458 |     1.36× | worse tradeoff than baseline |
| Hybrid Alt14 (no conv) |             ~3.248 |     1.23× | borderline                   |
| Hybrid Alt8 (no conv)  |              3.253 |     1.16× | minimal benefit              |




## References and recommended reads

- Yang et al. (2024). "Gated Delta Networks: Improving Mamba2 with Delta Rule." arXiv:2412.06464. ICLR 2025.
- Ben-Gurion et al. (2024). "Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues." arXiv:2411.12537.
- Dao & Gu (2024). "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality." (Mamba2)
- Katharopoulos et al. (2020). "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention."
- Widrow & Hoff (1960). "Adaptive Switching Circuits." IRE WESCON Convention Record.
- OLMo Team (2025). "2 OLMo 2 Furious." arXiv:2501.00656. COLM 2025.
- Lieber et al. (2024). "Jamba: A Hybrid Transformer-Mamba Language Model." AI21 Labs. arXiv:2403.19887.
- De et al. (2024). "Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models." Google DeepMind. arXiv:2402.19427.
- Waleffe et al. (2024). "An Empirical Study of Mamba-based Language Models." NVIDIA. arXiv:2406.07887.
- Glorioso et al. (2024). "Zamba: A Compact 7B SSM Hybrid Model." Zyphra. arXiv:2405.16712.
- Ren et al. (2024). "Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling." Microsoft. arXiv:2406.07522. ICLR 2025.
- Dong et al. (2024). "Hymba: A Hybrid-head Architecture for Small Language Models." NVIDIA. arXiv:2411.13676.
- Arora et al. (2024). "Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff." (Based) Hazy Research. arXiv:2402.18668.
- Goldstein et al. (2024). "GoldFinch: High Performance RWKV/Transformer Hybrid with Linear Pre-Fill and Extreme KV-Cache Compression." arXiv:2407.12077.


## Acknowledgements

Thanks to Akshay and Samip from Qlabs for the compute and the amazing slowrun directory!

Thanks to Vatsa, Pushkar, Matt and others for proofreading the blog:)