Parallelizing Linear Transformers with the Delta Rule over Sequence Length

Thanks for your review!

W1 Long context experiments

Thanks for your suggestion. Models at 1B scale are currently not powerful enough to provide meaningful results on needle-in-the-haystack style long-range benchmarks. We are currently training models are larger scale (3B parameters), and will run the needle-in-the-haystack tests with these models once they have finished training. (We give the 3B model results, which have been trained on 167B tokens so far, on the other benchmarks. We find that DeltaNet performs comparably to Transformers at this scale.)

W2: reason of deltanet being slower than GLA.

Both DeltaNet and GLA scale linearly with respect to sequence length; however, the recurrent transition matrices in GLA are diagonal matrices, making them amenable to tiling techniques (commonly used to accelerate matrix multiplications on GPUs) because the hidden states are independent of each other. The recurrent transition matrices in DeltaNet are more expressive, modeling state-to-state dependencies and thus requiring "marginalizing" over the entire head dimension, making tiling less amenable, as discussed in the limitations section (lines 301-306).

Additionally, the computation of the WY representation in DeltaNet is more expensive than computing the cumulative product of decays in GLA. From Eq. 7-8, we can see that the DeltaNet chunkwise form involves more matrix multiplications than vanilla linear attention (and also GLA).

W3 Intuitive explanation and theoretical justification of delta update rule

We found [1] motivating the delta rule intuitively from the perspective of key-value associative memory: the key intuition of the delta rule is to subtract the value associated with the current input key, namely the old value, from the memory, and write a new value that is a linear combination of the old value and the current input value into the memory. This encourages the formation of key-value associations, making retrieval easier.

Regarding theoretical justification, in lines 322-323, we reference several theoretical papers that reveal the superiority of the delta rule in memory capacity in a mathematical way.

[1] Linear Transformers Are Secretly Fast Weight Programmers https://arxiv.org/pdf/2102.11174

Q1 Do you think linear models can fundamentally bridge the gap with transformers in memory-based tasks?

We don't think linear models cannot fundamentally bridge the gap with Transformers in memory-based (or recall-intensive) tasks. Several theoretical papers reveal such limitations [1, 2].

Still, we believe the use of the delta rule will push the Pareto frontier of the recall-memory tradeoff curve, a concept that was advocated in [2].

It is possible and promising to replace the subquadratic encoder in YOCO [3] with DeltaNet to fully bridge the gap with Transformers in memory tasks.

[1] RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval. https://arxiv.org/abs/2402.18510

[2] Simple linear attention language models balance the recall-throughput tradeoff https://arxiv.org/abs/2402.18668

[3] You Only Cache Once: Decoder-Decoder Architectures for Language Models https://arxiv.org/abs/2405.05254

Q2: Is there an inherent conflict between the ability to handle long context and the performance of memory-based tasks?

This is a good question, and one for which the community doesn't have a good answer for yet (in our opinion). Our sense is that memory-based tasks will always require some attention-like mechanism which "retrieves" portions of thecontxt. It is possible that subquadratic-but-superlinear attention mechanisms (e.g., those based on clustering) may enable efficient long-context modeling while still enabling good performance on memory-based tasks.

Thanks for your review!

We are adding some additional results in case they are of interest.

First, we have preliminary results with 3B models trained for 167B tokens:

We are aiming to train these models for longer (300B-1T tokens) and will report these results in the next iteration of the paper.

We have also run experiments on the recently-released MAD benchmark, a suite of synthetic tasks designed to test the capabilities of language models beyond perplexity. Here are the results:

Thanks for your review.

W1: larger scale experiments

As noted by the reviewer, it is difficult to conduct experiments at 1B+ scale. Nonetheless, we are currently running some larger-scale experiments at the 3B parameter scale. Here are some preliminary results:

As we can see, the DeltaNet results are comparable to the Transformer++ results at this scale. We will make sure to include these results once they have finished training (we are aiming for 300B-1T tokens, depending on the availability of compute). We also plan to open-source the pretrained models so that researchers can study these models in more detail.

We have also run experiments on the recently-released MAD benchmark, a suite of synthetic tasks designed to test the capabilities of language models beyond perplexity. Here are the results:

W2: The improved results compared to Mamba and GLA make use of additional architectural components: convolution and local/global attention, without them the results are comparable to the other models.

Thanks for the comment! Note that Mamba uses conv layers by default, while GLA does not. This is why we trained our own GLA+conv baselines. Our main experiments compare DeltaNet+conv against Mamba+conv and GLA+conv, i.e., we give all models the chance to use convolution layers for a fair and meaningful comparison. In this setting, DeltaNet outperforms both Mamba and GLA. Morevoer, these depthwise separable "short convolution" layers are cheap both in terms of the number of parameters and compute, and hence we think that this convolution primitive is a practical approach to modeling local interactions.

Q1 Effect of Chunk size C

The computational complexity of the WY representation is O((N/C) * C^3) = O(NC^2). If C is too large, WY representation computation will be very expensive (recall that WY computation is fully recurrent).

If the chunk size is too small, and we adapt the materialization version of FlashLinearAttention, we need to write more hidden states to HBMs, resulting in a high I/O burden and thereby slowing down the actual running speed. If we adapt the non-materialization version of FlashLinearAttention, it will lack sequence-level parallelism, which is not desirable in large-scale training. Please refer to Section 3.3 of the GLA paper for more discussions.

To provide some intuition, we measured the running time (forward and backward pass) on a single A100 GPU while varying the chunk size $C$:

The experimental settings were as follows: sequence length $\text{seq\_len} = 4096$, batch size $B = 2$, head dimension $\text{d\_head} = 128$, model dimension $\text{d\_model} = 2048$, and number of heads $\text{num\_head} = 16$. As we can see, the moderate chunk size of 32 performs the best. When the chunk size is less than 32, the I/O cost surpasses the WY representation cost, and vice versa.

Q2: Chunk state

Yes this is correct! Concretely, we adapt the "non-materialization" version of FlashLinearAttention for chunkwise DeltaNet implementation. In this version, the hidden states of each chunk are materialized on HBMs in the forward pass, then discarded to save memory. During the backward pass, the hidden states of each chunk are recomputed.

Q3: DeltaNet w/o shortconv in large scale

Thank you for the suggestion. We agree that this will be an interesting experiment to run. We are currently using all our resources for the 3B experiments, but we will run this ablation once the 3B experiments are done.

Minor

Thanks for identifying the typos/errors! We will fix these.

Thanks for your feedback!

The potential issue is that DeltaNet without convolution is in several tasks worse than GLA without convolution. Any insight on why that is?

Local information plays a critical role in NLP tasks. As highlighted in [3], incorporating more local features can significantly enhance the performance of linear attention. Many recent studies also demonstrate that additional local attention mechanisms greatly improve linear attention.

Short convolutions are one method to enhance local information, while the gating mechanism in GLA imposes a strong inductive bias towards local context. DeltaNet’s update rule, on its own, doesn’t prioritize local information, so without short convolutions, it may underperform in NLP tasks. However, when short convolutions are applied, DeltaNet’s ability to leverage local context is improved. Since GLA already incorporates gating mechanisms to enhance local information, short convolutions may not provide additional benefits.

Additionally, it’s important to note that both the delta rule and the vanilla linear attention update rules can be seen as "context-based addressing" mechanisms. [1] emphasizes the importance of "location-based addressing," and [2] also highlights that a drawback of linear attention is its "lack of precise local token shifts and comparisons." The gating mechanism in GLA addresses this by emphasizing nearby tokens, while short convolutions can be considered another form of location-based addressing.

• [1] https://arxiv.org/abs/1410.5401 Neural Turing Machines

• [2] https://arxiv.org/abs/2402.18668 Simple linear attention language models balance the recall-throughput tradeoff

• [3] https://arxiv.org/abs/2312.11135 Linear Attention via Orthogonal Memory

For completeness, I'd suggest adding some discussion about it to emphasize that it's essential

Thank you for your suggestion! We will include this discussion in the next iteration of the paper draft to emphasize its importance.

Thank you for answering my questions and providing additional results!

• The experiments with 3B model provide evidence that performance is comparable to that of a Transformer as model size increases to 3B. Given the compute constraints, my concern has been addressed adequately.

• Regarding the convolutional component, I am not challenging the fairness of the comparison. The potential issue is that DeltaNet without convolution is in several tasks worse than GLA without convolution. Any insight on why that is? For completeness, I'd suggest adding some discussion about it to emphasize that it's essential to the performance of DeltaNet and reporting the results of DeltaNet w/o convolution for the 1.3B model, which are currently missing.

• I found the results about the effect of chunk size insightful; it'd be good to include these results in the final version.

My questions and concerns have been answered adequately except the second bullet point above. Hence, I decided to keep my original scores.