Unlocking Feature Learning in Gated Delta Networks at Scale

Yifeng Liu; Quanquan Gu

Source-linked AI summary

Unlocking Feature Learning in Gated Delta Networks at Scale

Yifeng Liu, Quanquan Gu

arXiv:2606.04048v1Jun 2, 2026cs.LGcs.AI

TL;DR

Scaling LLMs motivates efficient architectures and principled hyperparameter transfer, but µP extensions to structured linear models remain limited. This paper derives µP scaling rules for Gated Delta Networks by analyzing forward, gating, and recurrent-state coordinate sizes, then validates them in language-model pre-training. The formulation transfers learning rates across widths under AdamW and SGD, while standard parametrization does not.

Problem
Scaling LLMs is computationally expensive, Transformer attention is quadratic in sequence length, and µP extensions to structured linear models remain limited.
Method
The paper derives Gated Delta Network µP rules by propagating coordinate-size estimates through the forward pass, gating mechanisms, and recurrent state dynamics.
Results
Gated Delta Network µP enables zero-shot learning-rate transfer across widths under both AdamW and SGD, whereas standard parametrization fails to transfer.
Takeaways & Limitations
The derived formulation empirically validates the theoretical scaling analysis for Gated Delta Networks under both optimizers.
Takeaways & Limitations
The derivations approximate the gating analysis by suppressing SiLU activations, and approximation quality degrades when pre-activations are far from zero.

Abstract

from arXiv · show

Training and scaling Large Language Models demand enormous computational resources, motivating both efficient sub-quadratic architectures and principled hyperparameter tuning methods. While the Maximal Update Parametrization ($μ$P) has enabled zero-shot hyperparameter transfer for standard Transformers, its extension to linear models, particularly those with structured state transitions and complicated architectures, remains largely unexplored. By rigorously propagating coordinate-size estimates through the forward pass, gating mechanisms, and recurrent state dynamics, we derive the scaling rules for Gated Delta Network. Experiments on language-model pre-training confirm that our configurations enable stable learning-rate transfer across model widths under both AdamW and SGD, whereas standard parametrization fails to transfer, validating the correctness and practical utility of our analysis.

1 Introduction

Scaling LLMs increases training cost and exposes the quadratic cost of Transformer attention, motivating efficient linear architectures and principled hyperparameter transfer. This paper derives and validates a µP formulation for Gated Delta Networks.

Larger models require more computation because optimal performance increases with model size and training budgets grow steeply.
Transformer softmax self-attention has quadratic sequence-length complexity, making long-context training and inference increasingly costly.
Linear attention and recurrent state-space models address these efficiency limitations through linear structures and recurrent inference.
µP supports zero-shot hyperparameter transfer from small proxy models to large target models by identifying feature-learning parametrizations.
The paper derives complete µP scaling rules for Gated Delta Networks, including non-standard scaling for gating weights and scalar gating parameters.
Pre-training experiments show zero-shot learning-rate transfer under both AdamW and SGD, whereas standard parametrization fails to transfer.

2 Related Work

Related work develops efficient sequence models that replace quadratic attention with linear or recurrent mechanisms, alongside methods for accelerating hyperparameter search and transfer. µP provides a principled parametrization for transferring hyperparameters across widths.

Efficient Sequence Models: Linear attention rewrites softmax attention with a kernel to enable O(1) per-step inference.
Efficient Sequence Models: SSMs model long-range dependencies with recurrent state spaces, including structured and input-selective transition mechanisms.
Efficient Sequence Models: Delta-rule architectures update fast-weight structures and form part of a broader family of efficient linear sequence models.
Hyperparameter Transfer: Hyperparameter-transfer research includes search acceleration, cross-task transfer, and approaches based on standard parametrization.
Hyperparameter Transfer: µP extends this line by providing width-aware parametrizations intended to support hyperparameter transfer as models scale.

3 Preliminaries

Gated Delta Net combines delta-rule state updates with gating, normalization, convolutions, and nonlinearities. µP addresses width-scaling instability by preserving bounded, non-vanishing feature updates through width-dependent initialization and learning-rate choices.

3.1 Gated Delta Net: Gated Delta Net uses query, key, and value vectors in a recurrent latent-state update, with αt controlling retention and βt controlling writing strength.
3.1 Gated Delta Net: Short convolutions and SiLU activations follow the query, key, and value projections, while L2 normalization and RMSNorm stabilize the architecture.
3.2 µP theory: Under standard parametrization, constant learning rates can cause diverging activations and gradients as width grows, requiring O(1/d) scaling.
3.2 µP theory: µP specifies width-dependent initializations and learning rates so feature updates remain bounded and non-vanishing as width approaches infinity.
3.2 µP theory: A vector has Θ(d^a)-sized coordinates when each coordinate has variance Θ(d^2a) as width increases.
3.2 µP theory: µP requires Θ(1)-sized pre-activations, O(1) outputs, and maximally large parameter updates that avoid divergence.

4 The µP Forward Analysis of Gated Delta Net

The forward analysis propagates coordinate-size estimates through projections, recurrent state dynamics, readout, and gating. It yields normalization and initialization rules that preserve µP behavior, while identifying an approximation caused by SiLU activations.

The analysis derives µP conditions by propagating coordinate-size estimates through the forward pass and gating mechanisms.
The hidden state is assumed to have Θ(1) coordinate size and updates of the same order, with SiLU temporarily suppressed for the analysis.
Under µP initialization, query, key, and value projections plus short convolutions preserve Θ(1) coordinate size.
Contractive recurrent dynamics make the latent state’s steady-state variance track the write updates under the stated gating assumptions.
The readout from latent state and query produces a smaller coordinate scale before RMSNorm, so a √d-multiplier is inserted beforehand.
With the multiplier or an equivalent query RMSNorm, standard µP applies to projection weights apart from those governing αt and βt.
The gating analysis treats αt and βt as nonlinear transformations rather than Gaussian variables, requiring separate treatment of their weights and scalars.
Under µP, gating weights use initial variance 1/d while scalar gating parameters use width-independent initial variance.

5 The µP Analysis of Gated Delta Net under SGD

The SGD analysis derives learning-rate scalings for Gated Delta Net by tracking gradient magnitudes through recurrent states, normalized projections, and gating. Under short effective memory, projection matrices share Θ(1) scaling, while gating parameters require distinct scalings.

Notation and assumptions: Under SGD, absent gradient normalization, different Gated Delta Net weights receive different gradient magnitudes and therefore require different learning-rate scalings.The feature-learning condition requires each hidden-state update to remain Θ(1).
Notation and assumptions: The analysis assumes short effective memory, Leff = O(1), so the BPTT tail remains an O(1) factor relative to the direct gradient.When αt → 1, Leff can instead grow as O(T), requiring learning rates to be scaled down by 1/Leff.
Gradient of query, key and value projections: The value, key, and query projection matrices all use SGD learning-rate scaling Θ(1).For the value projection, the resulting update produces a Θ(1) change in vt, satisfying the feature-learning condition; the same analysis extends to key and query projections.
Gradient of the Gating: Gating weight matrices require a different learning-rate scaling from the projection matrices under SGD.The gating analysis treats direct state contributions while bounding recurrent future-state effects under the short-memory assumption.
Summary: The complete SGD µP formulation is summarized by learning-rate rules that differ from the standard µP law for gating weights and scalar gating parameters.Initialization variances and forward multipliers remain identical to the AdamW formulation; the learning-rate row is the relevant difference.
Gradient of the Gating: Scalar gating parameters require updates that induce a Θ(1) change in the gating pre-activation.Because their gradients are smaller, the scalar learning-rate scaling differs from the matrix parameters.

6 Experiments

The experiments evaluate Gated Delta Networks across widths, optimizers, learning rates, and parametrization schemes. Under µP, optimal learning rates transfer across widths, whereas standard parametrization does not.

6.1 Experiment details: Models use eight layers and six attention heads, with widths spanning d ∈ {256, 512, 1024, 1536} for AdamW and d ∈ {256, 512, 768, 1024} for SGD.These settings correspond to approximately 21M–342M non-embedding parameters.
6.1 Experiment details: The architecture uses query/key head dimension d/8, value head dimension d/4, kernel size 4, MLP size 4d, and tied input/output embeddings.
6.1 Experiment details: Training uses FineWeb-Edu 100B for 20k steps with 480-sequence batches, sequence length 1024, and one NVIDIA H100 80GB GPU.The training corpus contains approximately 9.83B tokens in total.
6.2 Experiment results: For AdamW, the optimal learning rate is consistently the same across all four widths under µP, while it shifts substantially with width under SP.Validation loss is reported after subtracting the best loss at each width across learning rates.
6.2 Experiment results: SGD shows the same qualitative pattern: learning rates fail to transfer under SP and vary substantially under original µP, but transfer under the proposed µP configuration.

7 Conclusion

The paper derives a µP-style parametrization for Gated Delta Networks and validates it through language-model pre-training. The proposed formulation transfers learning rates across widths under AdamW and SGD, unlike standard parametrization.

7 Conclusion: The analysis derives µP scaling rules for Gated Delta Networks, including distinct scaling for gating weights and scalar gating parameters under SGD.
7 Conclusion: Pre-training confirms zero-shot learning-rate transfer under both AdamW and SGD with the proposed µP formulation.
7 Conclusion: Standard parametrization fails to transfer learning rates across widths in the reported experiments.
7 Conclusion: The authors suggest that these derivations may inform scaling-law research for other linear or hybrid architectures.

A.1 Derivation of the cumulative latent space

The derivation analyzes the recurrent latent state by propagating coordinate-size estimates through the state update and its variance recurrence. Under the stated assumptions, the latent state coordinates remain at Θ(1/d) scale.

A.1 Derivation of the cumulative latent space: The derivation assumes zero-mean inputs, bounded gating scalars, and statistical independence between the previous state and current inputs at initialization.
A.1 Derivation of the cumulative latent space: Independence and zero-mean assumptions eliminate cross-terms when computing the expected squared state update.
A.1 Derivation of the cumulative latent space: The contraction factor is approximately E[α_t^2] and is strictly less than one, so the state variance follows a convergent geometric progression.
A.1 Derivation of the cumulative latent space: The limiting element-wise variance of the latent state is Θ(1/d), yielding Θ(1/√d) coordinate size.
A.1 Derivation of the cumulative latent space: Under AdamW, coordinate-wise gradient normalization makes effective update magnitudes Θ(1) across weight classes, including the gating matrices.

B.2 Derivation for Main Projection Weights (Wq, Wk, Wv, Wo)

The derivation determines learning-rate scaling for projection and gating weights by requiring each optimization step to produce an O(1) feature shift. For AdamW, the main projection scaling is governed by the fan-in dimension.

B.2 Derivation for Main Projection Weights (Wq, Wk, Wv, Wo): For a main projection z = W x_t, the gradient and Adam update determine the feature shift caused by changing W.
B.2 Derivation for Main Projection Weights (Wq, Wk, Wv, Wo): Because each hidden-state coordinate is Θ(1), its ℓ1 norm scales as Θ(d), amplifying the aggregate projection update.
B.2 Derivation for Main Projection Weights (Wq, Wk, Wv, Wo): The main projection learning rate must scale as 1/d to enforce the feature-learning condition |∆z_i| = Θ(1).
B.2 Derivation for Main Projection Weights (Wq, Wk, Wv, Wo): The gating-weight derivation likewise chooses its learning rate so that the pre-activation shift |∆z_α,t| remains Θ(1).
B.2 Derivation for Main Projection Weights (Wq, Wk, Wv, Wo): Under AdamW, learning-rate scaling is dictated by fan-in rather than output dimension, placing both matrix types under a 1/n_ℓ−1 multiplier.

B.4 Derivation for Scalar Parameters (alog and b)

The scalar-parameter analysis derives width-independent learning-rate scaling for achieving Θ(1) feature drift, then verifies the resulting µP dynamics empirically across model widths.

Scalar-parameter derivation: The gradient with respect to scalar bias b equals the pre-activation gradient, and Adam's scale-invariant update directly shifts the feature.Because b enters additively, its feature change equals the parameter update.
Scalar-parameter derivation: The scalar learning rate must remain independent of width to produce the required Θ(1) feature drift.
Experimental setup: The analysis evaluates models with widths d ∈ {256, 512, 768, 1024} during pre-training using the same optimal SGD learning rate of 0.4.
Forward-pass verification: Across widths, the scaled query quantities RMS(eq)/√dhead × √d remain approximately constant under the proposed µP configuration, unlike SP and original µP.The d = 768 setting was omitted because of instability across all configurations.

C.2 Verification of Backward Pass Gradient Scaling

The backward-pass diagnostics test whether gradients follow the coordinate-size scaling predicted by the proposed µP formulation, while additional probes examine gating and recurrent-state stability.

Hidden-state gradients: Under the proposed and original µP configurations, ∂L/∂h follows Θ(1/d) coordinate scaling, whereas SP produces Θ(1/√d) and scaling instability.
Query-key gradients: Both ∂L/∂eq and ∂L/∂ek follow Θ(1/d) coordinate scaling under the proposed configuration, while original µP deviates and transfers imperfectly.
Gating assumptions: The gating analysis assumes αt and βt avoid trivial saturation, such as vanishing completely or remaining strictly 1.0.
Gating stability: β remains centered near 0.5 across widths, with many tokens triggering strong writes, while zα,t and zβ,t retain Θ(1) standard deviation.The zβ,t analysis omits d = 768 because of instability across all configurations.
State-space dynamics: The cosine similarity between qt and kt stays around 0 as t increases and converges during training, supporting the independent-query-key assumption and stable state dynamics.

D Dynamic analysis of AdamW

AdamW diagnostics show that the proposed µP configuration preserves forward-pass scaling, stabilizes hidden-state gradients, and maintains nontrivial gating and recurrent-state dynamics across widths.

Experimental setup: The AdamW experiments analyze widths d ∈ {256, 512, 1024, 1536} using the same optimal learning rate of 8 × 10^-3.
Forward-pass scaling: Across widths, the scaled query and key quantities remain approximately constant under µP, unlike SP, supporting stable variance in the recurrent state St.
Backward-pass stability: Under µP, hidden-state gradients are much more stable across layers than under SP in AdamW experiments.Adam(W) normalization with µP also ensures Θ(1) hidden-state updates.
Gating stability: βt remains around 0.5 across widths, many tokens trigger strong writes, and zα,t and zβ,t retain Θ(1) standard deviation under AdamW.
State-space dynamics: The cosine similarity between qt and kt remains around 0 and converges during training, supporting independent query-key behavior and stable state dynamics under AdamW.