Linearizing Large Language Models

Authors: jean  mercat
igor  vasiljevic
sedrick  keh
kushal  arora
achal  dave
adrien  gaidon
thomas  kollar
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-13 22:07:14

Abstract: Linear transformers have emerged as a subquadratic-time alternative to softmax
attention and have garnered significant interest due to their fixed recurrent state.
However, they suffer from poor scaling and under-perform compute-matched
transformers. Prior models such as RWKV and Mamba have attempted to address
these shortcomings by proposing novel time-mixing and gating architectures,
but pre-training large language models requires significant data and compute
investments. In this paper, we propose Scalable UPtraining for Recurrent Attention
(SUPRA), an alternative to pre-training linear transformers. We present a method
to uptrain existing large pre-trained transformers into Recurrent Neural Networks
(RNNs) with a modest compute budget. This allows us to leverage the strong pre-
training data and performance of existing transformer LLMs, while requiring 5%
of the training cost. We find that our linearization technique leads to competitive
performance on standard benchmarks, but we identify a persistent in-context
learning shortfall for even the largest linear models.

Main paper: 
 introduction: 
 Over the last few years, Transformers (Vaswani et al., 2017) have displaced Recurrent Neural Networks (RNNs) in sequence modeling tasks, owing to their highly parallel training efficiency and unmatched scaling performance (Kaplan et al., 2020) . However, this training efficiency comes at the cost of inference cost that scales linearly with the number of tokens, compared to the fixed-cost inference of RNNs. The memory-intensive nature of transformers has led to renewed interest in recurrence-the fixed-size hidden state remains an attractive modeling proposition to reduce the cost of inference for language and multimodal models.Several recent works, starting with Linear Transformers (Katharopoulos et al., 2020) , have observed a relationship between a linearized form of attention and recurrence, leading to a duality between transformers and RNNs: models can be trained with sequence parallelism (i.e. as transformers, avoiding backpropagation through time), but can operate as RNNs at inference time. Although this architecture allows efficient training of RNNs, softmax transformers continue to outperform linear transformers across natural language understanding benchmarks. A number of novel RNN architectures have attempted to bridge this performance gap. These include RWKV (Peng et al., 2023a) , Retentive Networks (Sun et al., 2023) , TransNormer (Qin et al., 2022a) , and more recently, Griffin (De et al., 2024 ) and RecurrentGemma (Griffin Team et al., 2024) . These models are pre-trained on the same pre-training datasets as transformers and show promising results.State-space models (Gu et al., 2021) (SSMs) are another recurrent alternative to softmax transformers, combining RNNs and convolutional networks to efficiently model long sequences. The Mamba (Gu & Dao, 2023) architecture is a SSM that shows impressive performance at smaller scales, matching or exceeding the performance of softmax transformers on a number of natural language understanding (NLU) benchmarks. However, a gap remains for long-context NLU tasks, showing a persistent advantage of softmax attention.Architecture search at the scale of large language models is expensive. Rather than pretraining linear models, another approach is to convert an existing transformer into an RNN; Kasai et al. (2021) proposed to uptrain encoder-decoder transformers into RNNs by introducing an approximating MLP attention module. Zhang et al. (2024) improved on this method by adding a loss to match softmax attention to approximate more closely the base transformer.While approximating attention is an intriguing approach to re-using pre-trained transformers, it leads to instability and poor performance when uptraining large-scale models. We instead take a different approach: rather than approximate softmax attention, we replace it with a linear kernel and a normalization strategy to uptrain the most performant LLMs into RNNs (see Figure 2 ). We take advantage of models trained on high-quality, proprietary datasets for trillions of tokens (e.g. Mistral (Jiang et al., 2023) and Llama2 (Touvron et al., 2023) ). Fine-tuning these models on publicly available data for a small fraction of pre-training tokens (see Figure 1 ), we obtain linear models that are competitive with the best linear transformers for a fraction of the compute. We call our approach Scalable UPtraining for Recurrent Attention (SUPRA).Our contributions are as follows:• We propose Scalable UPtraining for Recurrent Attention (SUPRA), a linearization strategy to uptrain state-of-the-art LLMs into performant RNNs.• We show that this simple uptraining technique is competitive with the strongest pre-trained recurrent LLMs.• We investigate the limitations of recurrent LLMs, comparing pre-trained and uptrained RNNs to transformers, revealing a persistent gap for in-context learning and long-context tasks.Figure 2 : Our linearization strategy: we replace the softmax normalization with GroupNorm (GN) and introduce a small MLP to project the queries and keys, converting a pre-trained attention block (left) to a linear attention (right). The model can be be trained in parallel as a transformer and used recurrently at inference time with a mathematically equivalent reformulation. 
methodology: 
 In this section we review linear transformers, and describe the linearization technique of Kasai et al. (2021) as it lays the groundwork for our approach. Finally, we present SUPRA, our method for uptraining large transformers into RNNs.Linear Transformers (Katharopoulos et al., 2020 ) establish a connection between transformers and RNNs, generalizing the definition of attention by replacing the softmax dot-product attention v ′ with a more generic similarity function sim(q, k) between the queries q and keys k:EQUATION)Standard softmax attention is a special case, using sim(q, k) = exp q T k √ d.The authors explore several alternative functions for sim(q, k), including a linear kernel.Their main architecture uses the similarity function sim(q, k) = ϕ(q) • ϕ(k) with a fixed exponential linear unit kernel ϕ(x) = elu(x) + 1. They show the computational benefits of linear attention and, more importantly for this work, they demonstrate how such a model can be expressed as an RNN in the case of attention with causal masking.Recurrent Inference Linear attention can be expressed as an RNN that updates a state s i and a normalization factor z i at each time step. Katharopoulos et al. (2020) call these terms the attention memory and normalized memory. This RNN formulation is mathematically equivalent to linear attention, allowing the user to choose the most efficient one for a given task and hardware. Consider a stream of tokens we want to generate X = [x 1 , x 2 , x 3 , . . . ]. At inference time, we use the following update rule, where subscripts denote timestep in the recurrence (calling k i = W K x i , etc):EQUATION)s i = s i-1 + ϕ(k i )v T i(3)EQUATIONThe state s i acts as a constant-size KV cache. Instead of appending new values to the cache, the state is updated. This allows for inference cost that is constant in the number of generated tokens. Kasai et al. (2021) introduced a linear transformer uptraining procedure that converts a pre-trained softmax transformer into an RNN by approximating the attention computation with multi-layer perceptrons (MLPs). The method (T2R) starts with a softmax attention model, and linearizes the softmax operation. Recall the kernel linear attention similarity function:EQUATION)Instead of choosing ϕ as a simple non-linearity, the authors use a trainable layer:EQUATIONThe weights are shared between keys and queries for a given attention head. By using ϕ and rearranging the operations, attention can be written as:EQUATION)This allows the recurrent inference described in Section 2.1. However, this formulation has a number of drawbacks. First, it requires a significant re-training of the model, using approximately 20% of pre-training tokens for conversion, while suffering a 5 -10% drop in performance on language benchmarks. Furthermore, this approach was tested on relatively small models (≈ 100M scale). Because it mimics the attention formulation closely, it suffers from stability issues at larger scales. To address these issues, we modify the approach to adapt it to large-scale model uptraining.Rather than pre-training linear models from scratch, we choose to instead uptrain state-of-theart transformers. Leveraging models that take advantage of high-quality (but proprietary) pre-training datasets, we linearize them using a modest fraction of pre-training data (see Figure 1 ). We build on T2R, identifying two major issues and proposing SUPRA, an approach to fine-tuning very large transformers into RNNs.We first follow the literature in linear transformers and identify the normalization factor in linear attention as unstable (e.g. TransNormer (Qin et al., 2022a) ). In Section 3.3 we show that uptraining a 1B model following the procedure in T2R causes a large drop in performance. We instead follow Retentive Networks (Sun et al., 2023) and replace the normalization with a GroupNorm operation.Next we note that linear attention suffers more with absolute positional encoding than softmax attention, and a modern relative positional encoding scheme like RoPE (Su et al., 2021) is crucial for competitive performance. Rather than training a linear transformer from scratch incorporating these findings (RetNet, TransNormer) we use MLP kernels to convert large language models into RNNs.Starting with the pre-trained model, we add weights shared between keys and queries W, ϕ(x) = relu(Wx + b) and use the rotary positional embedding (RoPE (Su et al., 2021) ) such that the similarity function becomesEQUATIONWe normalize the output with a GroupNorm (Wu & He, 2018) instead of dividing by the sum of sim(q i , k j ) (as in T2R). We use a fixed decay vector γ ∈ (0, 1) h , with h heads, as in Sun et al. (2023) . This leads to the following attention formulation (see Figure 2 for a graphical representation):EQUATION)These new parameters are trained jointly with the rest of the network; at test time, we use the recurrent formulation for inference. 
experiments: 
 We uptrain a variety of models from the 1B to 7B range into RNNs (Llama2 (Touvron et al., 2023) and Mistral (Jiang et al., 2023)), and evaluate our models in two settings: standard language understanding benchmarks and long-context evaluations. We compare the results of different architectural choices and training strategies, and then show the limitations of linear models on various benchmarks, describing the persistent gap between vanilla attention and recurrence. We choose Llama2-7B and Mistral-7B as our base models for uptraining, but our recipe is general to any transformer model. We compare our procedure to a variety of pre-trained recurrent models. Given that the largest available state-space models are at the 2.8B scale, we also train a Mamba model on the RefinedWeb (Penedo et al., 2023) dataset from scratch for 1.2T tokens, to serve as a strong baseline for a pre-trained recurrent model2 .We use a fork of OpenLM (Gururangan et al., 2023) for all training and fine-tuning. Please see Section 7 for hyperparameters and further details on reproducibility.Language Modeling. In Table 1 we report results on standard NLU evaluations using the Eleuther evaluation harness (Gao et al., 2023) . We primarily compare to transformers and linear models at the 7B scale, and we train a Mamba model at 7B for comparison with RWKV-5. As our model is initialized from strong pre-trained transformers (Llama2 and Mistral-7B), it preserves performance on most benchmarks (except MMLU; see Section 4 for a discussion below). Our technique outperforms RWKV-5 with minimal uptraining and is competitive with our 7B Mamba trained from scratch on 1.2T tokens. Long Context. Recurrent models were thought to perform well on long-context tasks because of their ability to preserve performance beyond their training sequence size. However, their downstream performance on long-context tasks has not been well-documented. Prior studies either do not conduct long-context evaluations (Katharopoulos et al., 2020; Kasai et al., 2021) , evaluate only on perplexity (Sun et al., 2023; De et al., 2024; Gu & Dao, 2023) , or evaluate on datasets which require task-specific training (Peng et al., 2023a) . Instead, we consider downstream natural language tasks from the SCROLLS benchmark (Shaham et al., 2022a) . Specifically, in Table 2 we present two tasks -Qasper (Dasigi et al., 2021) and NarrativeQA (Kočisk ý et al., 2018) -from the set of tasks evaluated in the Llama2-Long report (Xiong et al., 2023) . We evaluate both tasks with an input context cut-off at different lengths. A strong long-context model should perform better given more context. However, the training context lengths for these models do not go beyond 8k tokens. Transformer models show the strongest results up to the context length they were trained for but degrade beyond that. Interestingly, applying the YaRN trick (Peng et al., 2023b) enables transformers to scale beyond their training context quite well. RWKV shows a strong ability to handle much longer context than its training. Our Mamba model on the contrary is not able to generalize beyond its training context length. Surprisingly, the RecurrentGemma model (Griffin Team et al., 2024) shows degrading performances even within its training context length. Finally, our Mistral-SUPRA model preserves some performance at larger context lengths but we believe it to result from the decay along the context length that shortens the effective context. This is discussed in more details below. We find a significant gap in performance between transformers and available linear models, including models uptrained from strong long-context transformers. We speculate that more sophisticated recurrent state update rules may be required to perform well at this task. Ideas such as gating strategies (De et al., 2024) , higher order linear attention (Mercat, 2020) , or associative binding (Munkhdalai et al., 2024 ) could be explored.The default decay factors proposed in Qin et al. (2024a) gives better results than no decay on short context benchmarks but at a long range, the decay cancels out the influence of the context (max(γ) 2048 = 3.35e -4 ). This can be related to a smooth version of window attention (Beltagy et al., 2020) . However, the figure 3 shows that the cross entropy of the linear model plateaus at longer range while the original model keeps improving. This is also confirmed on downstream evaluation, in table 2. As more context is given to the model, the long-range evaluation performance plateau. When using the decay values proposed in Sun et al. (2023) , that allow longer range attention, we observe a performance drop on short-context benchmarks and no substantial improvement on long-context evaluation. Ablations. Table 3 compares transformers pre-trained on 100B tokens to Mamba (Gu & Dao, 2023) , T2R (Kasai et al., 2021) , and our approach. At this scale, with 100B tokens of training, the Mamba model performs best and other models show similar performance.The second half of Table 3 shows results for uptraining from a pre-trained transformer. TheT2R (Kasai et al., 2021) uptraining was unstable, yielding poor results compapred to SUPRA. This confirms that normalization is key for maintaining performance of the base LLM when uptraining.To test the hypothesis that linear attention approximates the softmax attention, we experimented with a 2-step approach. The first step trains only the new parameters such that the model could learn to approximate the softmax. The second step fine-tunes all the weights. The results show no benefit from the two steps approach and indicates that the softmax is not approximated. See Appendix A for a different approach to compare softmax attention and linear attention.We also experiment with different kernel choices. The "no embed" models remove the additional fully connected layer that we added. Removing this layer significantly degrades the results. The "1+elu" models replace the ReLU activation with x-> 1 + elu(x) that is used in Katharopoulos et al. (2020) and that is simply a smooth version of the ReLU function. The "no activation" models do not use any activation function. It seems that the bias added by the linear layer of the kernel is very important to produce good results but the activation itself has little impact on the performance. Results are slightly better with a non-linear activation than without.Finally, we compare the results of SUPRA uptrainings from two pre-trained softmax models. It appears that pre-training a linear model for a 100B token yields better results than finetuning a softmax model that was trained with the same budget. These results also shows, along with the comparison of LLama2-SUPRA and Mistral-SUPRA in Table 1 , that SUPRA benefits significantly from a stronger pre-trained transformer. Thus, given a limited training budget, using SUPRA from a strong pre-trained model is the best option. 1 ). Given the existing performance gap between the strongest transformer models and the most performant linear models, SUPRA is a simple recipe for conversion, allowing the study of strong RNNs with limited uptraining. 
discussion: 
 Comparison to Transformers on Short-Context Tasks. Our approach does not explicitly approximate attention from the base transformer model (see Appendix A), we do see a modest drop in performance across all benchmarks compared to softmax transformers. This could be partially explained by the lower quality of our data compared to the pretraining mix used to train models like Mistral-7B. It is also likely that linear transformers are inherently less expressive. However, the performance drop is relatively modest on most benchmarks, and significantly smaller than the drop from T2R uptraining, which shows the relevance of our approach.Long Context Comparisons. Prior work on linear attention showcased similar or better validation set perplexity to transformer models over long context (e.g. Sun et al. (2023) ) but did not evaluate linear models on natural language long-context evaluations like SCROLLS (Shaham et al., 2022b) . The results in Table 2 show that recurrent models generally maintain performance beyond their training context (except for Mamba-7b) while transformers (without modification) do not. However, Table 2 also demonstrates that simple linear scaling of the rotary positional embedding (Peng et al., 2023b; emozilla, 2023) can allow for context scaling beyond the window used for training a given transformer model, effectively nullifying the performance edge of these linear models. Furthermore, transformers generally outperform linear models at their maximum training context length. Further research is needed into extending linear models to long-context inference to take full advantage of the lower inference cost relative to vanilla transformers.Limitations. Since our method relies on initializing with strong pre-trained transformers, our models inherit any of the biases and weaknesses of their base models. Additionally, models that are already instruct-tuned do not linearize as well as base models. Our models suffer from poor performance on MMLU which requires in-context learning (5-shot), a weakness of linear models (Aky ürek et al., 2024) . We leave the investigation of these weaknesses of linear models to future work and hope that our proposed uptraining approach can help facilitate and accelerate the research in this area. 
related work: 
 Linear Transformers. The linear transformers introduced in Katharopoulos et al. ( 2020) lagged behind vanilla transformers in downstream performance, and subsequent architectures such as TransNormer (Qin et al., 2022a) and RetNet (Sun et al., 2023) narrow the gap, but do not demonstrate competitive results with modern transformers at scale. RWKV (Peng et al., 2023a) , a linear transformer that takes inspiration from LSTM (Hochreiter & Schmidhuber, 1997) , is competitive with compute-matched transformer-based models, but lags behind on a number of NLU benchmarks. Griffin (De et al., 2024 ) is a concurrent model that takes a hybrid approach, combining a sliding window with linear attention shows impressive performance relative to vanilla transformers, but is trained on a high-quality proprietary dataset.Another thread in the literature focuses on efficient attention alternatives (Performers (Choromanski et al., 2020), Cosformer (Qin et al., 2022b) , LUNA (Ma et al., 2021) , RFA (Peng et al., 2021) , Attention-free Transformer (Zhai et al., 2021) ). All of these approaches sacrifice performances for efficiency. Efficiency improvements for vanilla transformers have narrowed the capabilities gap between vanilla and linear transformers. The KV cache (Pope et al., 2023) greatly narrows the inference efficiency gap between linear and vanilla transformers. RingAttention (Liu et al., 2023) allows for very long context scaling of vanilla attention without approximation. Other methods such as HGRN (Qin et al., 2023) , HGRN-2 (Qin et al., 2024b) , and FlashLinearAttention (Yang et al., 2024) aim to further improve the linear architecture (throughput, memory, performance, etc.) by using gated attention and other such optimization tricks.State Space Models. State-space models (SSMs) such as H3 (Dao et al., 2022) , Hyena (Poli et al., 2023) , and Mamba (Gu & Dao, 2023) are recent alternatives to vanilla transformers, combining the strengths of convolutional and recurrent models with efficient hardware implementations. Instead of parallelizing training over the sequence, they produce an efficient way to train the sequential RNN. While these models are competitive with vanilla transformers on some tasks, we show that SSMs share the limitations of linear transformers on several in-context learning and long-context tasks.Uptraining Linear Transformers. Hedgehog (Zhang et al., 2024) builds on the work of Kasai et al. (2021) , identifying three different ways of training linear transformers -from scratch, uptraining quadratic transformers for a specific task, and uptraining generally.The authors focus on the first two, and we focus on the third. Moreover, they aim at approximating the softmax attention matrices with linear alternatives. In this work, we do not aim to approximate softmax attention, we replace it with a linear alternative (see ablation above and appendix A). Their method is only tested for smaller scale models and with parameter-efficient fine-tuning for larger models, but presents challenges for scaling for two reasons: (1) their training strategy involves comparing full attention matrices which is computationally expensive, and not feasible for full fine-tuning of large models with long sequences and (2) their method also inherits the gradient instabilities of linear transformers studied in Sun et al. (2023)  
conclusion: 
 We introduced SUPRA, a technique for converting large-scale pre-trained softmax transformers into recurrent neural networks, enabling the study of the strengths and limitations of recurrent models at scale with minimal compute cost. Compared to pre-training linear models from scratch, the SUPRA strategy produces competitive models comparable to the best available recurrent LLMs (RWKV and Mamba) at the 7B scale.We identify the strengths of linear models on standard NLU benchmarks but also the enduring limitations on in-context (i.e. MMLU) and long-context (NarrativeQA, Qasper) tasks, showing that linearized models do not inherit these capabilities from the base softmax transformers.One possible path to rectifying these limitations is explicitly training for in-context learning (Aky ürek et al., 2024) . We leave explorations of specialized and instruct data in the context of linear transformers to future work. More sophisticated gating mechanisms as in in Peng et al. (2023a) and De et al. (2024) are promising alternatives to our simple linearization. Using our uptraining method would greatly reduce the necessary time and cost of such experimentation. 
reproducibility: 
 Codebase We train our linear models using our fork of OpenLM (Gururangan et al., 2023) that we modify to include a linear attention function (printed below). We use Lightning Attention 2 (Qin et al., 2024a ) that offers a fast Triton (Tillet et al., 2019) kernel for linear attention computation. Evaluations are done with the Eleuther evaluation harness (Gao et al., 2023) . 
: 
 Figure 1 : We reuse pre-trained LLMs (gray) and convert them to RNNs with minimal uptraining (SUPRA), outperforming linear attention models (RWKV) on natural language tasks like HellaSwag, attaining the same memory advantages (center), but also inheriting the limitations of RNNs on tasks like MMLU (right).Linear transformers have emerged as a subquadratic-time alternative to softmax attention and have garnered significant interest due to their fixed-size recurrent state that lowers inference cost. However, their original formulation suffers from poor scaling and underperforms computematched transformers. Recent linear models such as RWKV and Mamba have attempted to address these shortcomings by proposing novel timemixing and gating architectures, but pre-training large language models requires significant data and compute investments. Thus, the search for subquadratic architectures is limited by the availability of compute and quality pre-training datasets. As a cost-effective alternative to pretraining linear transformers, we propose Scalable UPtraining for Recurrent Attention (SUPRA). 1 We present a method to uptrain existing large pre-trained transformers into Recurrent Neural Networks (RNNs) with a modest compute budget. This allows us to leverage the strong pretraining data and performance of existing transformer LLMs, while requiring 5% of the training cost. We find that our linearization technique leads to competitive performance on standard benchmarks, but we identify persistent in-context learning and long-context modeling shortfalls for even the largest linear models. Our code and models can be found at https://github.com/TRI-ML/linear open lm. 
 
 Appendix: We train and uptrain models on RefinedWeb (Penedo et al., 2023) (with 2 epochs for our Mamba training), which we tokenize with the pre-trained model's tokenizers. When training from scratch, we used the GPT-NeoX-20B (Black et al., 2022) tokenizer. We tokenize with sequence packing and use a default sequence length of 2048.We use square matrices with biases for the linear layers in the kernel ϕ functions to keep the same feature dimension in the queries and keys. We use the same kernel, with the same weights for both queries and keys and apply a ReLU activation. We use 1000 steps of linear learning rate warmup and a cosine learning rate decay from 3e -5 to 1e -5 for our 7B uptrainings and from 3e -4 to 1e -5 for our 1B uptrainings and for trainings from scratch. We use the Adam optimizer with β 1 = 0.9 and β 2 = 0.95. We trained our models with mini-batches totaling 2M tokens each. Our default RoPE frequency uses the Llama value of 10 4 . For longer sequence lengths, we use a RoPE frequency of 10 6 .Training Depending on the model size and the availability, we use from 4 to 32 nodes of 8 GPUs Nvidia H100 with Pytorch FSDP. We use a mixed precision strategy from OpenLM that automatically selects between bfloat 16 and float 32 for different operations. A linear 7B parameter model uptraining throughput is around 4300 tokens per second per GPU.Models Our results can be reproduced by following the same training recipe or using the model weights that we release: Mistral-SUPRA and Mamba-7b.1 def linear_attn_func (q , k , v , qk_scale : float , use_decay : bool = True , normalize : bool = False ) -> torch . Tensor : 11 s = slope_tensor (h , q. device , q. dtype ) 12 else :13 s = no_slope_tensor (h , q. device , q. dtype ) In this section we investigate whether our up-training procedure leads to linear attention that approximates the softmax attention from the base model, as might be expected.There are many possible ways to compare attention matrices. Moreover, some architecture changes such as attention decay and lack of normalization in the linear attention make a meaningful comparison difficult. We represent non-normalized comparisons in Figure 4 . It represents the cosine similarities and singular value distances between the attention matrices at every layer and for every head of the Mistral model compared with our Mistral-SUPRA. Each pixel of these images is a scalar similarity measure between two matrices represented by a color scale. In Figure 4 , we see large differences between the matrices.Since we removed the attention matrix normalization and replaced it with a LayerNorm Ba et al. ( 2016), we want to compare normalized attention matrices instead. We divide each line of the matrix by the absolute value of the sum of its elements such that the softmax attention matrix is unaffected and the linear attention matrix is normalized. In Figure 5 , we see significantly higher between most matrices with some exceptions. These observations indicate that the linear attention matrices derived from SUPRA are not an approximation of the softmax matrices. The Mistral-SUPRA and Mamba-7B models are released along with the code. 


Decision: This paper proposes Scalable UPtraining for Recurrent Attention (SUPRA), a method to uptrain existing transformer-based LLMs into RNNs (linear attention) with limited compute. The novelty (3/4 reviewers critique is too be *too limited*) isn't in the architecture (that's mostly taken from previous work), but rather in the empirical work on how to *train* such an architecture (with some updated components) and identifying the limitations of the resulting model. Perhaps it is the observations of limitations of the new model in terms of in-context learning that is the most interesting/original of the work, but it might not be enough to justify acceptance.

Pros:
- Practically useful observations on how to uptrain an attention-based LLM into an RNN.
- Interesting discussion on the limitations of such a model in terms of in-context learning.

Cons:
- Lack of novelty (3/4 reviewers) since mostly reuses an architecture proposed in previous work.
- Messy narrative (TZps) and some doubts about the reported results and baselines (TZps, XQ7B).

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: * 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
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: - 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: 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

Comments: title: Additional kernel experiments
Kernel experiments
additional answers
comment: Thank you for your interest and for your answer. We agree that a 0.6 difference is not insignificant and our further results below do suggest that GLA is indeed a stronger approach than the fixed decay used in SUPRA. It is worth studying further for long context and at larger scale but we believe this to be a broader study. 

|Model|Size|Tokens|HellaSwag|ARC-E|ARC-C|
|-|-|-|-|-|-|
|SUPRA|1B|(1.6T)+10B|57.0|62.4|31.6|
|SUPRA GLA|1B|(1.6T)+10B|57.6|64.3|33.1|

Additionally, we would like to share our perplexity evaluation as a function of context length to show why we believe that GLA improves the model overall but does not solve the long-context modeling capability issue of linear models. While showing a lower loss with GLA, linear model performance (including GLA) tend to plateau as the context is extended. 

|Model|8|32|128|256|512|1024|2048|
|-|-|-|-|-|-|-|-|
|SUPRA GLA 1b c4|3.75|3.15|2.91|2.68|2.55|2.71|2.58|
|SUPRA 1b c4|3.75|3.18|2.94|2.71|2.59|2.76|2.68|
|Softmax 1b c4|3.69|3.06|2.8|2.48|2.36|2.56|2.38|
Thank you for these results! I have two comments:

- 1B SUPRA vs. 1B GLA: Do you have any results beyond the HellaSwag dataset? I don’t consider a 0.6 difference to be slight.  It might be worth investigating whether training on more tokens offers an improvement over SUPRA's method with data-independent decays.

- Regarding Computational Resources: I understand the limitation regarding computational resources. However, since you were able to train the 7B Mamba model on 1T tokens as a baseline, it might be more fruitful to explore pushing this approach to its limits
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.
Thank you for your answer. We agree that adding more experiments with different kernel activation functions completes Table 3 nicely, strengthens the paper, and should answer your question. We uptrained combinations of kernel variants for 10B tokens from a pre-trained 1B parameters model:

The "no embed" models remove the additional fully connected layer that we added. Removing this layer significantly degrades the results. The "1+elu" models replace the ReLU activation with 
$f(x)=1+elu(x)$, as used in  [Transformers are RNNs](https://arxiv.org/abs/2006.16236), which is simply a smooth version of the ReLU function. The "no activation" models do not use any activation function. The bias added by the fully connected layer of the kernel is important for producing good results. However, the activation itself has little impact on performance. That being said, the results are slightly better with a non-linear activation than without.

|Model|Size|Tokens|HellaSwag|ARC-E|ARC-C|
|-|-|-|-|-|-|
|SUPRA|1B|(1.6T)+10B|57.0|62.4|31.6|
|SUPRA-no activation|1B|(1.6T)+10B|55.9|61.6|29.8|
|SUPRA-no embed-no activation|1B|(1.6T)+10B|36.5|48.5|23.7|
|SUPRA-1+elu|1B|(1.6T)+10B|56.5|64.0|31.1|
|SUPRA-no embed-1+elu|1B|(1.6T)+10B|49.9|57.4|30.6|
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.
Thanks for the reply.

I understand that the main contributions lie in the new training procedure, and agree that the results on MMLU and long-context benchmarks are valuable. 

Re relu vs identity, I'm wondering whether having non-negativity is really necessary from the empirical perspective (because linear attentions are not strictly mimicking softmax attention anyway.