PaperlayerView exact revision on arXiv

Source-linked AI summary

Parallelizing Linear Transformers with the Delta Rule over Sequence Length

Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, Yoon Kim

arXiv:2406.06484v6cs.LGcs.CL

TL;DR

Linear-time models reduce attention’s cost but remain limited on in-context retrieval, while existing DeltaNet training does not parallelize over sequence length. This paper introduces a hardware-efficient parallel training algorithm based on a memory-efficient Householder representation, scales DeltaNet to 1.3B parameters trained on 100B tokens, and reports stronger language-modeling and zero-shot performance than linear-time baselines, with hybrid variants outperforming transformer baselines.

  • Problem

    Linear attention and state-space models offer linear-time alternatives, but they underperform transformers on in-context retrieval, while DeltaNet’s existing training algorithm is sequential and hardware-inefficient.

  • Method

    The paper parallelizes DeltaNet training across sequence length by reparameterizing its recurrence and using a memory-efficient representation for products of Householder matrices.

  • Results

    DeltaNet outperforms strong linear-recurrent baselines in language modeling and downstream performance, while hybrid DeltaNet-attention models outperform strong transformer baselines.

  • Takeaways & Limitations

    The algorithm makes DeltaNet scalable to moderate-scale language-modeling settings while retaining strong performance on recall-oriented evaluations.

  • Takeaways & Limitations

    DeltaNet’s training speed still lags behind GLA, and its length generalization is limited relative to GLA and RetNet.

Abstract

from arXiv · show

Transformers with linear attention (i.e., linear transformers) and state-space models have recently been suggested as a viable linear-time alternative to transformers with softmax attention. However, these models still underperform transformers especially on tasks that require in-context retrieval. While more expressive variants of linear transformers which replace the additive update in linear transformers with the delta rule (DeltaNet) have been found to be more effective at associative recall, existing algorithms for training such models do not parallelize over sequence length and are thus inefficient to train on modern hardware. This work describes a hardware-efficient algorithm for training linear transformers with the delta rule, which exploits a memory-efficient representation for computing products of Householder matrices. This algorithm allows us to scale up DeltaNet to standard language modeling settings. We train a 1.3B model for 100B tokens and find that it outperforms recent linear-time baselines such as Mamba and GLA in terms of perplexity and zero-shot performance on downstream tasks. We also experiment with two hybrid models which combine DeltaNet layers with (1) sliding-window attention layers every other layer or (2) two global attention layers, and find that these hybrids outperform strong transformer baselines.

1 Introduction

Linear-time sequence models avoid attention’s quadratic cost and KV-cache growth, but they remain weaker on recall-intensive tasks. This work makes DeltaNet trainable in parallel over sequence length and scales it to larger language-modeling settings.

  • Motivation: Attention provides accurate sequence modeling and efficient GPU training, but its quadratic sequence-length complexity makes it expensive.Hardware-aware attention methods still require a KV cache whose size grows linearly with sequence length.
  • Motivation: Linear attention removes the KV cache and enables constant-memory inference, yet recall-intensive performance can lag behind transformers.Gated variants improve language-modeling competitiveness but still underperform transformers on some retrieval-oriented tasks.
  • Implementation: The parallel DeltaNet layer is released in the FLASHLINEARATTENTION library.The passage identifies the implementation as publicly available.
  • Prior Work: DeltaNet improves associative recall with a delta-rule update, but its original sequential training algorithm was hardware-inefficient and difficult to scale.Earlier evidence covered synthetic tasks and small-scale language modeling or translation.
  • Contribution: The proposed algorithm parallelizes DeltaNet’s forward and backward passes across sequence length using a memory-efficient representation of Householder-matrix products.The method enables scaling DeltaNet to moderate-scale language modeling and reports stronger results than Mamba and GLA.

2 Background

Linear attention rewrites attention as a matrix-state linear recurrence, trading quadratic parallel computation for constant-memory inference and linear recurrent computation. DeltaNet replaces additive memory updates with a delta rule that can correct or remove stored associations, while chunkwise parallelization addresses sequential training bottlenecks.

  • Linear Attention: Linear attention replaces the softmax exponential kernel with a feature-map dot product and represents the computation as a matrix-valued linear RNN.The feature map is φ: R^d → R^n.
  • Linear Attention: The simplified linear transformer updates its state additively and computes each output by multiplying the state with the current query.Its recurrent complexity is O(Ld^2).
  • Training Forms: The fully parallel form uses a causal mask and matrix multiplications, while the recurrent form uses fewer FLOPs but cannot parallelize across sequence length.The parallel form costs O(L^2d + Ld^2), whereas the recurrent form costs O(Ld^2).
  • Chunkwise Parallel Form: Chunkwise parallelization propagates information between chunks while avoiding materialization of intra-chunk states, trading sequential computation against sequence-level parallelism.With chunk size C, its complexity is O(LCd + Ld^2); C = L gives the fully parallel form and C = 1 gives the recurrent form.
  • DeltaNet: Additive updates can cause key collisions when sequence length exceeds state dimension because old associations are difficult to deallocate.An effective update should remove less important associations based on interactions between the new key and memory contents.
  • DeltaNet: DeltaNet uses a delta update that adjusts memory according to the difference between the current prediction and target value.The update can be interpreted as one-step online regression optimization, with β_t controlling the learning rate or writing strength.
  • DeltaNet Training: The original DeltaNet training algorithm was strictly sequential, motivating an equivalent chunkwise algorithm for hardware-efficient large-scale training.Theoretical parallel scan would require materializing a two-dimensional hidden state at every time step, causing substantial memory I/O unless the state is small.

3 Parallelizing DeltaNet Across the Sequence Dimension

The section develops a chunkwise-parallel training algorithm for DeltaNet by representing its recurrent updates with compact products of generalized Householder transformations. This avoids materializing matrix-valued hidden states and enables sequence-level parallelism while retaining a tunable trade-off between computation and sequential depth.

  • Memory-efficient recurrence: Naïvely computing pseudo-values requires explicitly materializing the previous state, costing O(d2) memory; the proposed construction reduces this to O(d) memory.The construction obtains the pseudo-values without storing each matrix-valued hidden state.
  • Memory-efficient recurrence: DeltaNet’s additive update can be represented through generalized Householder transformations, enabling compact WY representations of their products.This representation avoids materializing matrix-sized hidden states at every time step during parallel training.
  • Chunkwise parallelization: Computing all pseudo-values sequentially still costs O(L2d), so the method derives a chunkwise-parallel form that propagates information between chunks and computes intra-chunk states in parallel.Chunkwise matrices and recurrences preserve the memory-efficient representation while exposing sequence-level parallelism.
  • Hardware-efficient implementation: The implementation rewrites recurrent operations into matrix multiplications, uses forward substitution for triangular inverses, and recomputes hidden states during backpropagation to save GPU memory.The resulting kernels are implemented by adapting FLASHLINEARATTENTION.
  • Hardware-efficient implementation: Chunkwise speed-ups increase with sequence length L and head dimension dhead, benefiting from sequence-level parallelism and higher GPU occupancy.The comparison uses Triton implementations of recurrent and chunkwise-parallel forms across sequence lengths and head dimensions.
  • Fully parallel form: The fully parallel DeltaNet form is avoided for training because its required matrix inverse scales cubically with sequence length.The derived attention matrix may still be useful for interpretability studies of recurrent models.
  • DeltaNet transformer: The DeltaNet transformer replaces self-attention while retaining standard transformer modules, and uses SiLU feature maps with L2-normalized key/query vectors for stability.At βt = 1, the transition becomes a projection that erases one subspace while preserving the other d − 1 subspaces, supporting targeted forgetting.

4 Empirical Study

The empirical study evaluates DeltaNet against recurrent and Transformer baselines on synthetic recall, language modeling, zero-shot reasoning, and training-throughput benchmarks. DeltaNet generally improves language-modeling and downstream performance, while hybrid variants outperform Transformer++ baselines and linear-time models train faster than Transformers on longer sequences.

  • Experimental setup: DeltaNet is compared with Transformer++, RetNet, Mamba, and GLA across synthetic and real-world language-modeling evaluations.Synthetic benchmarks include MQAR, MAD, and RegBench; language-modeling evaluations include perplexity, zero-shot reasoning, and recall-intensive tasks.
  • Synthetic benchmarks: DeltaNet performs perfectly on the hardest MQAR setting and outperforms convolutional Mamba in the low-dimension setting.These experiments use two DeltaNet heads without convolutions.
  • Synthetic benchmarks: DeltaNet is better at MAD recalling tasks than other architectures, especially Fuzzy Recall, but struggles on the Memorize task.The MAD suite probes synthetic token-manipulation capabilities, while RegBench results are deferred to the appendix.
  • Language modeling: At 340M parameters, DeltaNet outperforms GLA on recall-intensive tasks, whereas at 1.3B it underperforms GLA because of poorer state-size scalability.The reported recall-intensive tasks are SWDE, SQuAD, and FDA; state size is identified as important for these tasks.
  • Language modeling: Both sliding-window and global-attention DeltaNet hybrids outperform strong Transformer++ baselines.The sliding-attention hybrid interleaves sliding-window attention every other layer, while the global-attention hybrid uses full attention on two layers.
  • Training throughput: All linear-time models outperform Transformers for longer-sequence training, while DeltaNet trains close to GLA and significantly faster than Mamba.Throughput comparisons use 1.3B models on a single H100.

5 Discussion and Related Work

The paper frames linear recurrent models through associative matrix-valued state updates, contrasting cheap elementwise recurrences with richer but more expensive structured matrix interactions. It also identifies limitations in the unifying framework and DeltaNet’s current efficiency and length-generalization trade-offs.

  • Associative RNN framework: Associative matrix-valued recurrences unify several linear recurrent models and enable parallel scan when their update operator is associative.Parallel scan computes all states in O(log L) steps with O(L) work, excluding associative-operation costs.
  • Associative RNN framework: Recent practical models favor Hadamard-product updates because expensive associative operators can prevent language-model training at scale.Elementwise recurrence costs O(dn) per update, whereas unrestricted matrix multiplication costs O(dn^2).
  • DeltaNet and structured transitions: DeltaNet uses structured rank-one transition matrices to model interactions beyond elementwise recurrence while retaining parameter efficiency and emphasizing associative recall.The adopted parameterization sets D = I, a_t = β_tk_t, and b_t = k_t; broader diagonal-plus-low-rank forms are left for future work.
  • Framework scope: The framework is not claimed to capture all autoregressive subquadratic transformations, including models with unstructured matrix recurrences or exotic associative operators.The authors emphasize that efficient, hardware-friendly training algorithms remain an important criterion for unifying frameworks.
  • Limitations and future work: The work’s algorithm parallelizes DeltaNet training across sequence length, but training still lags GLA because state-to-state dependencies require marginalization over the head dimension.This overhead may constrain DeltaNet’s memory size and recall-intensive performance, although block-diagonal transitions are suggested as a possible improvement.
  • Limitations and future work: DeltaNet has limited length generalization compared with GLA and RetNet, and to some extent Mamba, which can extrapolate beyond training length.The authors speculate that missing explicit decay factors contribute to this limitation and suggest adding a gating term.

6 Related Work

Related work connects linear transformers to associative memory and highlights a tension between expressive recurrent enhancements and sequence-length parallelism.

  • Scope: The paper situates its contribution within recent work on linear-time sequence models and associative memory mechanisms.The related-work discussion connects memory capacity, recurrence design, and hardware parallelism.
  • Associative memory: Linear transformers can be viewed as iterated Hopfield networks, with vanilla Hebbian-like updates having limited memory capacity.Higher-order polynomial and exponential-kernel variants are related to efforts to enhance memory capacity.
  • Parallelism and expressiveness: Recurrent DeltaNet enhancements and related models improve expressiveness but cannot be parallelized across sequence length.This pattern suggests a fundamental trade-off between parallelism and expressiveness.

7 Conclusion

The paper presents a sequence-length-parallel DeltaNet training algorithm that speeds up modern-hardware training and supports moderate-scale language modeling, where DeltaNet performs well against linear-recurrent baselines.

  • Conclusion: The proposed algorithm parallelizes DeltaNet training across sequence length and achieves significant speed-ups against existing implementations on modern hardware.This enables scaling DeltaNet to moderate-scale language-modeling settings.
  • Conclusion: DeltaNet performs well compared with recent linear-recurrent baselines in moderate-scale language modeling.The conclusion summarizes the empirical comparison without specifying a single benchmark value.

A Experiments Continued

The continued experiments describe the hardware and optimization setup for 340M and 1.3B language-modeling runs.

  • Training setup: 340M models were trained on 15 billion tokens, while 1.3B models were trained on 100 billion tokens using 8 H100 GPUs.The corresponding batch sizes were 0.5M and 2M tokens, respectively.
  • Training setup: Both model sizes used AdamW with a peak learning rate of 3 × 10^-4 and cosine learning-rate schedules.Warm-up lasted 0.5 billion tokens for 340M models and 1 billion tokens for 1.3B models.

A.2 Synthetic tasks

RegBench evaluates in-context language learning by requiring models to infer a probabilistic language from context and predict subsequent tokens.

  • RegBench uses sequences containing 10 to 20 strings drawn from a distinct language defined by a probabilistic finite automaton.The task requires inferring the underlying language from the context on the fly.
  • The benchmark is designed to assess the in-context language learning capability of different model architectures.
  • Testing evaluates models on predicting the next token of generated testing sequences.

B.1 WY representation derivation

This subsection presents the WY representation derivation in the context of the synthetic RegBench accuracy figure and focuses on the first chunk for notational simplicity.

  • Figure 7 reports accuracy (%) on the RegBench benchmark.
  • The supplied material provides the benchmark figure caption and a first-chunk scope note, but no further WY derivation details.
  • The derivation discusses only the first chunk to reduce notational clutter.

B.2 UT transform derivation

The UT transform derivation rewrites DeltaNet’s recursive updates in matrix form and connects this computation to efficient chunkwise training and hybrid model design.

  • UT transform derivation: The derivation shows that the UT-transform matrix formulation is equivalent to the recursive update equations.
  • UT transform derivation: The recursive computation of w_r[t] is converted into a matrix system that can be solved for W[t].
  • UT transform derivation: The same derivation applies to U[t] by replacing K[t] with V[t] in the final step.
  • Chunkwise training: Chunkwise linear-attention algorithms provide a precedent for exact outputs and hardware-aware I/O optimization in linear attention.
  • Hybrid models: Hybrid architectures combine linear recurrent layers with local, sliding-window, or global attention, and DeltaNet can be combined with classic attention.
  • Chunkwise training: The forward pass of the chunkwise DeltaNet algorithm is illustrated with Pytorch-like code, omitting batch-size and head dimensions.
Loading 2406.06484v6…