Linearizing Large Language Models

Summary

This paper follows the previous line of works in finetuning softmax attention LMs to linear attention ones. The benefit is (1) to reduce the cost of training from scratch: only need 10% training tokens to achieve similar performance, as shown in this paper; (2) to reduce the inference cost by transforming to the recurrent mode during inference. This paper adopts many advanced linear attention techniques to enhance the performance, for example using decays and output normalization. The authors give an important observation that output normalization is crucial for stable finetuning, while directly adopting the methods from T2R will lead to inferior results. More interestingly, this work also demonstrates some drawbacks of linear models, for example the incapability in long-text modeling and incontext learning.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

- Although the techniques are not new (reason to reject 1), the large-scale empirical verification of these techniques is important to the community. This paper also presents some interesting empirical results, revealing the incapability of linear attention/RNNs in long text modeling and in-context learning.

Reasons to reject

1. No new techniques are proposed. All methods are proposed before, including incorporating decaying factors, applying output normalization and trainable MLP feature map, etc. The main contribution to me is re-emphasizing the importance of output normalization in the large-scale funetuning setting. 2. The perplexity values should be reported for all experiments. Perplexity is still one of the most important metrics. 3. I found the paper lacking an important baseline: continually train the original model on the same data to see if MMLU declines a lot. If so, this is not the issue of linear Transformers, but the issue of the training data itself. This is a very important baseline. I will increase the score if the author can answer this question in depth.

Questions to authors

Questions and Suggestions: 1. In table1, it seems that continual training for more tokens could bridge the gap between linear and softmax attention. It would be interesting to continually train on a small model (e.g., 1B TinyLLaMa) for more tokens to see, to what extent, linear attention could bridge the performance gap to softmax attention. 2. In table3, why do you only report the results of HellaSwag and ARC? Do you have a specific reason? What about other metrics? 3. Could you release an an anonymous huggingface repo for the 7B mamba model? 4. It would be interesting to use the attention distillation loss proposed in Hedgehog for learning the feature map? 5. It would be interesting to try data-dependent decays proposed in Gated Linear Attention. https://github.com/sustcsonglin/flash-linear-attention Missing reference: (although some are released after the submission deadline, it is beneficial to discuss in the next iteration): - linear attention with data-dependent decays [1, 2, 5]. - Fine-tuning attention-based LMs into linear attention [1, 3, 4]. [1] Fine-Tuning Pre-trained Transformers into Decaying Fast Weights. https://aclanthology.org/2022.emnlp-main.697.pdf [2] Gated Linear Attention Transformers with Hardware-Efficient Training. https://arxiv.org/abs/2312.06635 [3] Cross-Architecture Transfer Learning for Linear-Cost Inference Transformers. https://arxiv.org/abs/2404.02684 [4] DiJiang: Efficient Large Language Models through Compact Kernelization. https://arxiv.org/abs/2403.19928 [5] HGRN2: Gated Linear RNNs with State Expansion https://arxiv.org/abs/2404.07904

Summary

This paper proposes Scalable UPtraining for Recurrent Attention (SUPRA), a method to convert existing pretrained transformers into a linear transformer that is much more efficient in terms of data and inference. The authors conducted extensive analyses to validate their proposed methods. The results show that SUPRA can maintain much of the performance of the original transformers, though the performance is only competitive with existing RNNs and below original transformers.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

- This is a very interesting paper exploring the intersection between RNNs and transformers. Linearizing transformers seems to be a pretty interesting and simple approach to achieve a good compromise between efficiency and performance. While the linearized model is not able to outperform the orignal transformers and existing RNNs for now, it does offer an alternative to existing approaches. Personally, I think the community should encourage such exlorations in addition to SOTA chasing. - This paper raises some new interesting questions, especially about how linear transformers are competitive with transformers in some tasks but lag behind in others (in-context learning, instruction-tuning). - The authors have made good efforts to ensure the reproducibility of this paper.

Reasons to reject

- Kernel choice. The paper only looks at a very specific linear kernel but I would appreciate more justifications. The chosen kernel has a non-negligible impact on the final performance (Tsai et al., 2019; Arora et al., 2024 - the latter is too new so it's not a problem for not discussing it). The proposed method has made some incremental improvement to T2R but more in-depth analysis of the impact of kernel can greatly improve this paper. **References** - Tsai, Y. H. H., Bai, S., Yamada, M., Morency, L. P., & Salakhutdinov, R. (2019). Transformer dissection: a unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv:1908.11775. - Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., ... & Ré, C. (2024). Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668.**

Questions to authors

- In table 2, it is mentioned that *"†These results were evaluated on subset of the first samples and will be updated"*. I don't think this is a good practice that should be encouraged. It is the authors' responsibility to make sure that the results are complete for review. Otherwise, incomplete evaluations might cause misleading conclusions.

Summary

This paper explores a way to fine-tune pre-trained Transformers into RNNs, which benefit from efficient inference. It is based on the work of Kasai et al. (EMNLP 2021), with improved model components and new experiments applied to a larger scale, and to compare with some recent pre-trained RNN models.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

Linear Transformers is an intriguing topic. This paper is timely with a solid contribution of some new experiment results.

Reasons to reject

There is not much highlight in the proposed method. It tweaks the model of Kasai et al. (EMNLP 2021) by switching to components that empirically turn out to be better in recent years, which I believe would improve performance but there is no conceptually new idea. In the Introduction, the authors compare this work with Kasai et al. and claim that "We take a different approach: rather than approximate softmax attention, we instead replace it with a linear kernel and a normalization strategy to fine-tune the most performant LLMs into RNNs. We take advantage of models trained on high-quality, proprietary datasets..." But Kasai et al. is actually doing the same thing -- it didn't use the GroupNorm component as in this paper and didn't apply to pre-trained models of 7B scale, but that doesn't mean it cannot -- the experiments conducted in this paper are not newly enabled because of the method proposed in this paper. Moreover, the ablation test shown in Table 2 does not seem to support that the proposed method has any specific strength. Overall, the narration of this paper is messy. It indicates that the authors lack a well-polished story. One interesting result from the experiments, is that even at the very large 7B scale, fine-tuning a pre-trained Transformer into an RNN can immediately decrease the performance of some tasks -- more detailed investigation into this issue might be a good direction to explore in the future.

Summary

This paper studies an interesting question of finetuning a standard attention-based LLM into a sub-quadratic model (i.e., linear attention models in this work). This paper tried to address several technical problems: 1) which sub-quadratic architecture is more compatible with finetuning, 2) choice of decay, kernel functions and normalization, 3) whether to fine-tune in a two-stage way. Overall, this paper reveals the potential of finetuing open-source LMs into their sub-quadratic versions as an alterative to training one from scratch.

Rating

6

Confidence

5

Ethics flag

1

Reasons to accept

An interesting work showing the potential of uptraining current open-source LMs (e.g., llama and mistral) into linear models. It also reveals several limitations of uptrained models, which are useful for future research: they are harder to benefit from longer-context training, and still suffer from poor few-shot learning performance.

Reasons to reject

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).

Questions to authors

* I wonder whether the kernel function of relu (following T2R) is really necessary, compared with using an identify function as in RetNet. * it would be better to compare with the following baseline [1] [1] The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry

Additional kernel experiments

We would like to complete our answer with the additional kernel experiments that we have done. Please refer to [our answer to reviewer JkbR](https://openreview.net/forum?id=soGxskHGox&noteId=g6AZlJ3S2v) that has expressed a similar interest in a more thorough ablation of the kernel choice.