Linearizing Large Language Models

Thanks for taking the time to give detailed feedback and suggestions! We agree with them and feel that incorporating these comments will make our paper stronger. Due to length limitations, we kept our responses short and had to omit some questions. We are happy to provide more details during the discussion period.

---

I found the paper lacking an important baseline: continually train the original model [...]

Good point! We include these results below. Here, we uptrain Llama2-7B on 20B tokens of RefinedWeb, which is the same setting used to uptrain Llama2-SUPRA.

All the evaluations are preserved with continual training except for MMLU which shows a slight drop when uptraining the softmax Llama2 (45.8->43.8). The MMLU performance drop is much wider for Llama2-SUPRA (45.8->24.9).

This tells us that the MMLU performance drop can be attributed to both the architecture (linearization procedure) and the uptraining process (data, hyperparameters), with the linearization procedure being the bigger contributor.

The perplexity values should be reported for all experiments.

Previous works have relied heavily on perplexity to report good performance of linear models at context extension. We have grown skeptical about the pertinence of these results when evaluating long-context tasks and noticing discrepancies. Nonetheless, we agree that this is one of the most important metrics and that including these results would strengthen our paper.

Below we compare the average cross entropy as a function of the token position in the sequence for Llama2-Uptrain-10B vs Llama2-SUPRA on a val set (C4 subset).

This adds further evidence to our claim that linear uptrains perform poorly at higher context lengths – in early tokens, the gap between Llama2-uptrain and Llama2-SUPRA is small, but as we go to later tokens, the loss for Llama2-uptrain continues to decrease while the loss for Llama2-SUPRA stays high.

In the final version, we will include perplexity values for all our other experiments and models.

We would like to thank you for increasing the score after our first answer. For completeness, we wanted to add answers to the other questions that you raised. We are happy to discuss more if you have any additional questions.

5. It would be interesting to try data-dependent decays proposed in Gated Linear Attention.

We actually tried this previously and noticed only a slight difference.

• 1B SUPRA; 10B RW tokens; HellaSwag = 57.0
• 1B GLA; 10B RW tokens; HellaSwag = 57.6

We believe the data dependent decay suffers from the same problem as fixed decay, which we mention in the paper (decays too fast and does not allow long context attention). We are continuing to look into this for future work, but this would require a thorough study of gating mechanisms and is out of the scope of this current work.

1. In table1, it seems that continual training for more tokens could bridge the gap between linear and softmax attention.

As training is pushed from 20B tokens to 70B tokens, the results do improve (Table 1). They would probably improve further with even longer training. However, training for more tokens would bring the cost closer to pretraining. In our experiments, our approach seems most relevant in the context of uptraining with a modest compute budget. With more available compute, it is likely that other approaches such as Mamba would yield better results at lower cost and would therefore be preferable.

2. In table3, why do you only report the results of HellaSwag and ARC? Do you have a specific reason? What about other metrics?

For Table 3, we only considered 1B models for ablations. Since 1B models are weaker than the 7B models in Table 1, we opted to evaluate on easier test sets such as HellaSwag and ARC. Evaluating on a difficult set like MMLU would not be very meaningful because the results are just close to random guessing.

4. It would be interesting to use the attention distillation loss proposed in Hedgehog for learning the feature map?

Thank you for bringing this up! We were aware of the Hedgehog paper and in fact cited it in our related work. However, there is no released code beyond the pseudocode attached to the paper. This made it time-consuming to implement and test, but we agree that it is an interesting baseline and will try to include it in our camera ready version.

More specifically, a naive implementation of their approach would require two full computations of the attention matrices to compute the loss that the authors introduce. This is known to be memory intensive and we would not be able to train at scale with this approach without an involved implementation of memory efficient attention with an added attention loss.

Furthermore, in our Appendix, we show that our trained SUPRA matrices do not actually approximate the softmax matrices. The RetNet paper shows how the normalization of the attention factor that allows to approximate the softmax can lead to instabilities. When uptraining 1B models, we did experience these instabilities in practice. Thus, while we lack the experiments that would give a definitive answer to your question, we believe that the Hedgehog approach of trying to very closely approximate softmax may not be the correct approach for uptraining such linear models at scale.

Unfortunately, the Hedgehog paper does not publish evaluations on the common benchmarks like HellaSwag, ARC, and PIQA which does not allow a direct comparison of our respective results.

3. Could you release an anonymous huggingface repo for the 7B mamba model?

Yes, we plan to release our models and code together with this submission. For the 7B mamba, we have created an anonymous model on Hugging Face https://huggingface.co/mambacolm/mambacolm

Thanks as well for including the missing references! We will include these in the final version of the paper.

Thank you for your positive review! We aimed to provide a comprehensive view of the linearization process and examine its viability as an alternative to pretraining linear models. We are glad you appreciated our work and specifically the comparison of the strengths and limitations compared to softmax transformers.

---

Kernel Choice: The kernel choice was ablated in Table 3. We tried a few different things but the highest performance was the RetNet-like kernel. We will also clarify the choices of the kernel and add further ablations in our final version. Thank you for the suggestion!

Regarding Table 2: We agree with the reviewer that evaluating on a subset is not ideal. We apologize for underestimating the time it would take to compute the 16k results and not filling the full test at submission time. For this specific case, evaluating on the full evaluation set gave us very similar results and insights as our reported results. We will include complete results in the final version.

Thanks for your review and for appreciating our experimental results. Our goal was to find a better way to linearize LLMs while still preserving strong performance. We focused on this last aspect as we found that the full picture of evaluations of linear models was missing from the literature.

---

Kasai et al. is actually doing the same thing [...] The experiments conducted in this paper are not newly enabled because of the method proposed.

While our method for linearization is based on T2R (Kasai et al), both our experiments and our discussion significantly differ from that work. We argue that our paper, while working in the same framework as T2R, enables scaling. Some design choices may seem obvious in hindsight, but scaling is a nontrivial task. In Table 3, we show that naively scaling up T2R is unstable and fails. We did the work of testing different design choices and finding ways to outperform most concurrent linear models with a lower training budget (Table 1).

In addition, we argue that the value of our work lies not only in architectural elements but also in our thorough experimental evaluation. For instance, we explore performance of transformers and linear models (not just our models but also others) on long-context tasks and tasks like MMLU. This is something that previous papers like RetNet or T2R did not address.

The ablation test shown in Table 2 does not seem to support that the proposed method has any specific strength.

Our goal is not to produce linear models that match softmax transformer performance. Rather, we aim to provide a holistic view of the strengths/weaknesses of our linearization process. As you mentioned, Table 2 shows that our models do not perform well on long-context. At the same time, Table 2 also shows that linear models in general (Mamba, RWKV) do not perform well on long-context, and that is something that we wanted to highlight in our paper. Many previous works imply that linear models perform well on long-context tasks. Thus, as noted by other reviewers, we believe Table 2 to be a strong experimental result: not in the sense that it shows any advantage of our method, but because it is a new result clearly identifying the limitations of linear transformers.

Having said that, we would like to point out that our model outperforms other linear models like RWKV on multiple benchmarks (HellaSwag, PIQA) with much fewer training tokens, which we highlight as one of the strengths of our proposed method.

Thank you for your detailed review of our paper. We wanted to provide a holistic view of the linear uptraining process and highlight both its strengths and limitations. We appreciate your acknowledgment of the interesting aspects of our work and the potential it shows for uptraining current open-source LMs into linear models.

---

In terms of technical novelty, the contribution of this paper is rather limited: the architecture for uptraining basically follows RetNet (including incorporation of RoPE and decaying factor).

Although our kernel closely follows the design in RetNet, the way we train the model is very different. We demonstrate that by uptraining pretrained models, we can achieve RetNet-like performance much more quickly. This is reflected in our results: with the same number of tokens (100B), our model significantly outperforms RetNet (e.g. HellaSwag 60.7 vs 77.1). Furthermore, RetNet experiments were limited to 2048 context training; we show the challenges of scaling the approach further.

In addition, we argue that the value of our paper lies not only in the introduction of new architectural elements but also in the thorough experimental evaluation we conducted. For instance, we explore the performance of transformers and linear models (not just our SUPRA models but also others) on long-context benchmarks and on benchmarks like MMLU. This is something that previous papers such as RetNet or T2R did not address.

I wonder whether the kernel function of relu (following T2R) is really necessary, compared with using an identify function as in RetNet.

We chose the ReLU kernel function because it allows non-negative attention factors which is desirable given that softmax attention is also non-negative.

It would be better to compare with the following baseline

Thank you for bringing this up! We were aware of the Hedgehog paper and in fact cited it in our related work. However, there is no released code beyond the pseudocode attached to the paper. This made it time-consuming to implement and test, but we agree that it is an interesting baseline and will try to include it in our camera ready version.

Unfortunately, the Hedgehog paper also does not publish evaluations on the common benchmarks like HellaSwag, ARC, and PIQA which does not allow a direct comparison of our respective results.