Unlocking Feature Learning in Gated Delta Networks at Scale

When scaling Large Language Models (LLMs) to massive sizes, engineers face a daunting bottleneck. How do you know which hyperparameters—like the learning rate—will work on a giant model if you can only afford to train a tiny one? Traditionally, finding the right settings requires expensive "grid searches." This is essentially trial and error performed at scale.

Current state-of-the-art methods like Maximal Update Parametrization ($\mu$P) allow researchers to tune settings on small "proxy" models. They can then transfer these settings to large ones with near-perfect accuracy. However, $\mu$P was designed for standard Transformer architectures. As the industry moves toward more efficient, linear recurrent models like Gated Delta Networks (GDN) to handle longer contexts, these existing scaling rules begin to break.

A new study from researchers at UCLA addresses this gap. The authors report that applying standard scaling rules to GDN leads to training instability. It also prevents successful hyperparameter transfer. By deriving a new, mathematically rigorous $\mu$P formulation specifically for the GDN architecture, they have unlocked "zero-shot" learning-rate transfer. This ensures that what works for a small model will reliably work for a massive one.

The breakdown of standard scaling

The push for efficiency has led to the rise of linear recurrent models. Unlike the standard Transformer, which has quadratic complexity (the computational cost grows exponentially with the length of the text), linear models like GDN maintain a constant cost per step. This makes them highly attractive for long-context applications.

However, the internal mechanics of GDN are significantly more complex than a standard Transformer. Instead of simple attention mechanisms, GDN relies on a "delta rule" to update a latent state (a compressed mathematical representation of the sequence history). It also utilizes data-dependent gating to control how much information is written or erased at each step.

The authors find that these updates happen through a recurrent state that evolves over time. Because of this, standard ways we scale weights and learning rates fail. If an engineer uses "Standard Parametrization" (SP)—the common practice of using fixed initialization and uniform learning rates—the model's activations and gradients tend to diverge as the model gets wider. This failure is visible in the researchers' experiments. Under SP, the optimal learning rate shifts drastically as the model width increases [, Figure 2]. This makes it impossible to predict the behavior of a large model from a small one.

Propagating coordinate sizes through gates

To solve this, the authors employ a technique called coordinate-size estimation. Think of this as tracking the "voltage" of individual components in a circuit as you increase the size of the system. If the voltage spikes too high, the circuit fries. If it drops too low, nothing happens. In neural networks, we want the "voltage" (the variance of the activations and gradients) to remain stable regardless of how wide the model becomes.

The researchers' approach involves several rigorous steps:

1. Forward Pass Analysis: They trace how the magnitude of vectors moves through the projection matrices ($W_q, W_k, W_v$) and the gating mechanisms. They determine that the L2-normalized query and key vectors have a specific "coordinate size" of $\Theta(1/\sqrt{d})$. Here, $d$ represents the model width.
2. Latent State Dynamics: They analyze the recurrent update rule. This involves both a "write" term (adding new information) and an "erase" term (removing old information). They prove that the latent state $S_t$ maintains a stable variance of $\Theta(1/d)$ per entry.
3. Gating Calibration: Most crucially, they look at the gating parameters ($\alpha_t$ and $\beta_t$). These decide how much to update the state. They discover that these parameters behave differently than standard weights.
4. Backward Pass Derivation: They propagate these estimates backward through the chain rule. This determines exactly how the learning rate must scale for every weight class. This ensures "feature learning" (the process where the model learns meaningful patterns rather than just memorizing noise).

Successful zero-shot transfer

The effectiveness of this new math is confirmed by pre-training experiments on the FineWeb-Edu 100B dataset. The authors compare their proposed $\mu$P formulation against both the standard approach and a previous attempt at $\mu$P for structured state space models.

The results are striking. For the AdamW optimizer, the authors report that their $\mu$P configuration enables the optimal learning rate to remain almost perfectly constant across all tested model widths .

This means the setting found for a 256-width model works for a 1536-width model. For the SGD optimizer, the improvement is even more dramatic. While the original $\mu$P configuration failed to transfer effectively, the authors' specialized formulation achieved nearly perfect zero-shot transfer .

Beyond just the loss curves, the authors used "diagnostic probes" to verify the internal physics of the model. They tracked the Root Mean Square (RMS) of the activations. Under their $\mu$P configuration, the scaled quantities remained constant across different widths [, Figure 4].

They also verified that the gradients for the hidden states followed the predicted $\Theta(1/d)$ scaling. This provides empirical proof that their theoretical derivation was correct [, Figure 6].

Limits of the derivation

While the paper provides a robust framework, it is important to recognize its boundaries. The theoretical derivation relies on a "short effective memory" assumption. Specifically, the authors assume that the gating values $\alpha_s$ decay the past state quickly. This prevents the influence of very distant tokens from overwhelming the gradient calculations. If a model must remember things over extremely long horizons without forgetting, this mathematical approximation might lose accuracy.

Furthermore, the analysis treats SiLU activations (a common non-linear function in modern AI) as being "suppressed" to facilitate the math. In real-world training, these activations introduce non-Gaussian statistics. Essentially, the "noise" in the system is not perfectly bell-shaped. The authors acknowledge that while this approximation is adequate for the initializations used in $\mu$P, it is not a perfect representation of the complex, non-linear reality of a running model.

The verdict: A new standard for linear models

If you are building or training large-scale Gated Delta Networks, the verdict is clear. Do not use standard initialization or vanilla $\mu$P. The authors have demonstrated that doing so will lead to unstable training and wasted compute.

The paper provides a specific recipe for success. For those using the AdamW optimizer, the gating weights can be treated similarly to standard hidden weights. However, for those using SGD, you must apply a non-standard $\Theta(1/\sqrt{d})$ scaling to the gating weight matrices. You must also apply a $\Theta(\sqrt{d})$ scaling to the scalar gating parameters.

The research is highly actionable. The code is reportedly available; see the paper for the canonical link. This work effectively bridges the gap between the theoretical promise of efficient linear architectures and the practical necessity of stable, scalable training.

Figures from the paper