Linearizing Large Language Models
Authors: Jean Mercat
Igor Vasiljevic
Sedrick Scott Keh
Kushal Arora
Achal Dave
Adrien Gaidon
Thomas Kollar
Venue: COLM (2024)
Keywords: linear attention, efficient attention, RNN
TL;DR: converting LLMs into RNNs through minimal up-training

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.

Published as a conference paper at COLM 2024
Linearizing Large Language Models
Jean Mercat" Igor Vasiljevic*
Sedrick Keh' Adrien Gaidon
Kushal Arora
Achal Dave
Thomas Kollar
Toyota Research Institute
[jean, [jean.mercat, mercat, igor. vasiljevic, sedrick.keh/ wwwkkeeeeerrrreee cotri.globall lkssaal.aora achal.dave, adrien. gaidon, thomas.
2000
90
1500
37.5
80
70
500
12.5
60
50
25
training (B)
Memory (GB)
HellsSwag
RWKV (RNN)
SUPRA (RNN)
Mistral-7B
Figure 1: We reuse pre-trained LLMs (gray) and convert them to RNNs with minimal uptraining (SUPRA), outperforming linear attention models (RWKV) natural language tasks like HellaSwag, attaining the same memory advantages (center), but also inheriting the limitations of RNNs tasks like MMLU (right).
Abstract
Linear transformers have emerged as a subquadratic-tine alternative softmax attention and have garnered significant interest due their fixed-size recurrent state that lowers inference cost. However, their orig- inal formulation suffers from poor scaling and underperforms compute- matched transformers. Recent linear 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 mod- els 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 pre- training linear transformers, we propose Scalable UPtraining for Recur- rent Attention (SUPRA) We present 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 re- quiring 5% the training cost. We find that our linearization technique leads to competitive performance tify persistent in-context learning and long-conted modeling shortfalls for even the largest linear models. Our code and models can be found at https://github com/TRIIML//inear _neeeeeeerro__mmmmmm
standard benchmarks, but we iden-
Equul contribution.
We borrow the term "uptraining" from Ainslie et al. (2033) to refer to continued training with a modified architecture, as opposed to fine-tur fine-tuning, which usually refers to continued training on different dataset.

---

Published as a conference paper at COLM 2024
1 Introduction
Over the last few years, Transformers (Vaswani et al. 2017 have displaced Recurrent Neu- ral Networks (RNNss 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 has led to renewed interest in recurrence the fixed-size hidden state remains an attrac- tive modeling proposition to reduce the cost of inference for language and multimodal models.
nature of transformers
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 RNNss 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 under- standing benchmarks. number of novel RNN architectures have attempted 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 KecurrentGemma (Griffin Team et al. 2024). These models are pre-trained on the same pre-training datasets as transtormers and show promising results.
State-space models (Gu al. 2021) (SSMs) are another recurrent alternative to softmax trans- formers, 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 pre- training 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 introd ucing an approximating MLP attention module. ZZhng et al. (2024) improved on this method by adding loss to match softmax attention to approximate more closely the base
transformer.
While approximating attention is an intriguing approach re-using pre-trained transform- ers, it leads to instability and poor performance when uptraining large-scale models. We instead take different approach: rather than approximate softmax attention, we replace 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, propri- etary datasets for trillions of tokens (e.g. Mistral (Jiang et al. (023) and Llama2 (Touvron eet al., 2023 ). Fine-tuning these models on publicly available data for a small fraction of pre-training tokens (see Figure(] [], 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), linearization strategy to uptrain state-of-the-a 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 up- trained RNNs to transformers, revealing a persistent gap for in-context learning and long-cone long-context tasks.

---

Published as a conference paper at COLM 2024
recurrent inference
s(t)
s((+1)
Groupham
ReLU
ReLU
pretrained transformer
FC
FC
FC
LayerNorm
equivalent
uptrained linear cell
LayerNorm
GroupNorm
MHA
ReLU
ReLU
FC
FC
FC
Figure 2: Our linearization strategy: we replace the softmax normalization with GroupNorm (GN) and introduce small MLP to project the queries and keys, converting pre-trained attention block (left) to linear attention (right). The model can be be trained in parallel as a transformer and used recurrently at inference time with mathematically equivalent reformulation.
Methodology
In this section we review linear transformers, and describe the linearization technique of Kasai et al. (2021) it lays the groundwork for our approach. Finally, we present SUPRA, our method for uptraining large transformers into RNNs.
2.1 Background: Linear Attention
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 more generic similarity function sim (q, k) between the queries q and keys k:
Σ sim q: k))vi
(1)
Σ
sim kj)
Standard softmax attention is special case, using sim (q, k) = exp
9
The authors explore several alternative functions for sim q, k), including a linear kernel. Their main architecture uses the similarity function sim
(( p(k) with fixed exponential linear unit kernel (xx) = elu(x) + 1. They show the computational benefits of linear attention and, more importantly for this work, they demonstrate how such 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; and a normalization factor z; each time step. Katharopoulos al. (2020) call these terms the attention memory and normalized memory. This RNN formulation is mathematically

---

Published as a conference paper at COLM 2024
equivalent to linear attention, allowing the user to choose the most efficient one for given task and hardware. Consider stream of tokens we want to generate X X1, X2, x3,... At inference time, use the following update rule, where subscripts denote timestep in the recurrence (calling k; = Wxx; etc):
So = Zo
(2)
1 + v
(3)
2 = Zi + ((kk)
(4)
v'; Φ 9:)'ss
(5)
((q:)
The state s; acts 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.
2.2 Finetuning Transformer into an RNN
Kasai et al. (2021) introduced linear transformer uptraining procedure that converts 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:
sim y) +x() ply).
(6)
Instead of choosing & as simple non-linearity, the authors use trainable layer:
Φ (x) = relu(Wx + b).
(7)
The weights are shared between keys and queries for a given attention head. By using Φ and rearranging the operations, attention can be written as:
((q:) - (k))v
(8)
((q) 7 Σ (kk))
This allows the recurrent inference described in Section 2.1 However, this formulation has a number of drawbacks. First, requires significant re-training of the model, using approximately 20% of pre-training tokens for conversion, while suffering 5 10% drop in performance on language benchmarks. Furthermore, this approach was tested on relatively small models (s= 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.
2.3 SUPRA: Scalable UPtraining for Recurrent Attention
Rather than pre-training linear models from scratch, we choose to instead uptrain state-of-ther art transformers. Leveraging models that take advantage of high-quality (but proprietary) pre-training datasets, linearize them using modest fraction of pre-training data (see Figure 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.) (022a)). In Section 3.3 we show that uptraining a 1B model following the procedure in T2R causes large drop in performance. We instead follow Retentive Networks (Sun et al.) 2023) and replace the normalization with GroupNorm operation.

---

Published as a conference paper at COLM 2024
Next note that linear attention suffers more with absolute positional encoding than softmax attention, and modern relative positional encoding scheme like RoPE (Su et al. 2021) is crucial for competitive performance. Rather than training 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, (xx) relu(Wx b) and use the rotary positional embedding (RoPE (Su al. 2021 such that the similarity function becomes
sim kj) = RoPE(((q)) (oPE((((j))
(9)
We normalize the output with GroupNorm (Wu & He 2018) instead of dividing by the of sim qi ki) (as in T2R). We use fixed decay vector Y (0,1)*, with / heads, as in |unn et al. (2023). This leads to the following attention formulation (see Figure for graphical representation):
v GroupNorm
Σ (im (g.k))y))
(10)
These new parameters are trained jointly with the rest of the network; test time, we use the recurrent formulation for inference.
Experiments
We uptrain a variety models from the 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 uptraining, but our recipe is general to any transformer model.
range into RNNs (Llama2 (Touvron al.
base models for
Model
Size
Tokens
HellaSwag
PIQA
WG
ARC-E
ARC-C
MMLU
Average
StableLM2
1.6B
2000
69.0
76.7
63.6
68.6
38.9
38.4
59.2
StableLM
3B
1000
73.8
79.3
65.8
72.1
40.0
44.2
62.5
Gemma
2B
2000
71.4
78.6
64.4
74.0
41.5
41.2
61.9
Mamba
1.4B
600
59.0
73.9
25.2
53.0
RWKV-5
1.58
1100
53.1
71.6
32.7
26.2
50.8
Mamba
2.8B
600
66.2
26.3
56.3
Llama2
7B
2000
76.0
79.1
45.9
65.4
Llama2
7B
+20
76.3
78.8
43.8
65.3
Gemma
7B
6000
80.7
81.9
53.2
62.9
72.2
Mistral
7B
8000(?)
81.0
82.1
80.9
53.8
62.4
72.4
RetNet
6.7B
200
60.7
75.4
58.1
RWKV-5
7B
1100
70.9
77.2
67.4
71.8
43.6
31.0
60.3
RWKV-5-1.7T
7B
1700
73.0
78.6
72.9
75.8
45.6
34.9
63.5
Mamba (ours)
7B
1200
77.9
81.0
71.8
77.5
46.7
33.3
64.7
Llama2-SUn SUPRA
7B
+20
71.8
78.6
65.8
71.1
39.5
24.9
58.6
Mistral-SUPRA
7B
+20
74.8
80.1
67.44
74.6
42.3
28.0
61.2
Mistral-SUPRA
7B
+100
77.1
80.4
70.3
75.9
45.8
34.2
64.0
Table 1: Linear models (RNNs and SSMs) highlighted in gray, 5-shot results are used for MMLU. Norm results are used for PIQA, HellaSwag, ARC-C. RetNet results taken from RetNet paper.

---

Published as a conference paper at COLM 2024
We compare our procedure to 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 Wef (Penedo et al., 2023) dataset from scratch for 1.2T tokens, to serve strong baseline for a pre-trained recurrent model
We use a fork of LpenLM (Gururangan et al. 2023) for all training and fine-tuning. Please see Section for hyperparameters and turther details on reproducibility.
Language Modeling. In Table 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 for discussion below). Our technique outperforms RWKV-5 with minimal uptraining and is competitive with our 7B Mamba trained from scratch on 1.2T tokens.
Model
Size
Train Context
Qasper (2-shot)
NarrativeQA (0-shot)
2048
4096
8192
16384
2048
4096
8192
16384
Llamal-7B
7B
2048
24.43
7.23
5.08
4.88
21.44
1.03
0.0
0.0
Llama2-7B
7B
4096
23.26
27.26
6,20
5.49
21.32
22.61
0.0
0.0
Llama2-7B*
7B
4096
23.26
27.26
31.46
25.52
21.32
22.61
23.0
14.27
Mistral-7B
7B
8196
21.53
25.50
33.61
6.88
24.94
26.90
25,93
0.63
RecurrentGemma-2B
2.7B
8192
22.44
13.16
13.42
12.66
19.80
11.59
12.93
12.95
RWKV-5-1.7T
7B
2048
22.28
23,87
22.30
20.35
17.77
18.65
17.81
16.00
Mamba (ours)
7B
2048
19.68
5.58
5.90
6.32
19.70
0.28
0.0
0.0
Mistral-SUPRA
7B
2048
19.44
17.13
17.11
17.22
18.99
17.76
17.75
17.74
Table 2:
Long context evaluations. Performance at various context size cutoffs for Qasper
(2-shot) and NarrativeQA (0-shot). denotes linear RoPE scaling with YaRN (Peng et al., 2023b
Long Context. Recurrent models were thought to perform well on long-context tasks be- cause of their ability to preserve performance beyond their training sequence size. However, their downstream performance on long-context 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, consider downstream natural language tasks from the SCROLLS benchmark (Shaham let al. 2022a). Specifically, in Table we present two tasks - Qasper (Dasigi et al. (021)) and NarrativeQA (Kočiský et al., 2018 from the set of tasks evaluated in the Llama2-Long report (Xiong et al. 20233. 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 de- grade beyond that. Interestingly, applying the YaRN trick (Peng et al. 2023b) enables transformers to scale beyond their training context quite well. RWKV shows 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
2The Mistral-SUPRA and Mamba-7B models are released along with the

---

Published as a conference paper at COLM 2024
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.
Decay factors. The default decay factors proposed Qin et (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(?)20x) 2048 3.35e-4). This can be related to smooth version 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 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 performance drop on short-context benchmarks and no substantial improvement long-context evaluation.
2.3
MUSPra
Mistral
2.2
Cress Entropy
2.0
1.9
500
1000
1500
2000
Taken Position
Figure 3: Plot of the cross entropy results for the base Mistral model and Mistral-SUPRA model for each token positions the sequence. subset of the C4 dataset was used. The transformer model benefits from more context while linear model plateaus.
Model
Size
Tokens
HellaSwag
ARC-E
ARC-C
Mamba
1B
100B
58.3
62.7
29.1
SUPRA (from scratch)
1B
100B
54.7
59.6
27.8
T2R (from scratch)
1B
100B
55.2
61.0
28.4
Transformer
1B
100B
55.9
60.4
29.3
Transformer*
1.6T
62.1
66.2
34.3
Fine-tune only new weights
1B
(1.6T)+10B
33.2
37.8
22.6
2-step fine-tune
1B
(1.6T)+10--10B
56.1
62.2
30.4
T2R
1B
(1.6T)+10B
40.6
53.8
27.9
SUPRA
(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+ein
1B
(1.6T)+10B
49.9
57.4
30.6
SUPRA-weaker base
1B
(100B)+10B
51.7
57.5
27.4
Table 3: Ablating different choices for linear uptraining: note the importance of normaliza- tion. For the second half of the table, we uptrain transformer trained on 1.6T tokens on 10B further tokens. *This model was trained on a different mix of data.

---

Published as a conference paper at COLM 2024
Ablations. 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 experi- mented 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 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 non-linear activation than without.
Finally, we compare the results of SUPRA uptrainings from two pre-trained softmax models. appears that pre-training a linear model for a 100B token yields better results than fine- tuning softmax model that was trained with the same budget. These results also shows, along with the comparison of LLama2-SUPRE and Mistral-SUPRA in Table [ that SUPRA benefits significantly from stronger pre-trained transformer. Thus, given limited training budget, using SUPRA from a strong pre-trained model is the best option.
Discussion
Comparison to pre-training SSMs/RNNs. With only 20B tokens of training, which repre- sents - - of RWKV and RetNet training cost, we obtain a model that outperforms both on HellaSwag and that is competitive 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.
other benchmarks (see Table Given the existing
Comparison to Transformers on Short-Context Tasks, Our approach does not explicitly approximate attention from the base transformer model (see Appendix AA, we do see 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 pre- training 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 vali- dation set perplexity transformer models over long context (e.g. (unn et al. (2023)) but did not evaluate linear models on natural language long-context evaluations like SCROLLS (Sha- ham 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 (with- out modification) do not. However, Table also demonstrates that simple linear scaling of the rotary positional embedding (Peng et al. emozilla| 2023) can allow for context

---

Published as a conference paper at COLM 2024
scaling beyond the window used for training given transformer model, effectively nulli- fying 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-contex. 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 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.
5 Related Work
Linear Transformers. The linear transformers introduced in Katharopoulos et al. (2020) lagged behind vanilla transformers in downstream performance, and subsequent architec- tures such as Tormerormer (Qin et al. 2022a) and RetNet (Sun et al. 2023) narrow the gap, but do not demonstrate competitive results with modern transtormers at scale. RWKV (Peng et al. 2023a), linear transformer that takes inspiration from LSTM (Hochreiter & Schmid huber 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 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 (Choro manski et al. 2020), Cosformer (Qin et al. 2022b), LUNA (Ma et al. 2021), RFA (Peng et al. 2021), Attention-tree Transformer (Zhai et al. (2221). All of these approaches sacrifice per- formances 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 transtormers. 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 architecture (throughput, memory, performance, etc.) by using gated attention and other such optimization tricks.
(Yang et al. 2024) aim to further improve the linear
State-space models (SSMs) such H3 (Dao et al. 2022), Hyena (Poli ltt al.) (202)), 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.
State Space Models.
Uptraining Linear Transformers. Hedgehog (Zhan et al. 2024) builds on the work of Kasai et al. 22021), 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 approximating the softmax attention matrices with linear alternatives. In this work, we do not aim approximate softmax attention, we replace it with linear alternative (see ablation above and appendix A). Their method is only tested for smaller scale models and with parameter-efficient fine- ain-tuning for larger models, but presents challenges for scaling for two reasons: (1) their training strategy involves comparing full attention matrices which
the third. Moreover, they aim at

---

Published as a conference paper at COLM 2024
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), while our normalization setup leads to stable uptraining of large models. In addition, Mao (2022) investigates various update rule configurations for uptraining and propose decaying fast weights. Concurrent work such as Chen et al. (2024) and Choi (2024) also try to convert vanilla transformers into linear models, albeit with different uptraining methodologies and varying degrees of success. Notably, Diliang (Chen) it al., 2024) employs a novel Frequency Domain Kernelization technique to conduct the lineanzation, allowing them to uptrain a 7B model on 40B tokens.
Conclusion
We introduced SUPRA, a technique for converting large-scale pre-trained softmax trans- formers into recurrent neural networks, enabling the study of the strengths and limitations of recurrent models at scale with minimal compute cost. Compared to pre- 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 explicitly training for in-context learn- ing (Akyürek et al. 2024). We leave explorations of specialized and instruct data in the context of linear transtormers to future work. More sophisticated gating mechanisms as in in Peng et al. (023a) and De et al. (2024) are promising alternatives to our simple lineariza- tion. Using our uptraining method would greatly reduce the necessary time and cost of such experimentation.
10

---

Published as a conference paper at COLM 2024
7 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 fast Triton (Tillet et al. 2019) kernel for linear attention computation. Evaluations are done with the Eleuther evaluation harness (Gao
et a.. 2023).
Data 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-200 (Black et al. 2022) tokenizer. We tokenize with sequence packing and use default sequence length of 2048.
Hyperparameters We 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 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 3 1e-5 for our 1B uptrainings and for trainings from scratch. We use the Adam optimizer with B1 0.9 and B2 = 0.95. We trained our models with mini-batches totaling 2M tokens each. Our default RoPE frequency uses the Llama value of 10 For longer sequence lengths, we use a RoPE frequency of 106.
Training Depending on the model size and the availability, we use from 4 to 32 nodes of GPUs Nvidia H100 with Pytorch FSDP. We use 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
def aonear_attn. func (q, k, v, qk_scale: float, use_decay: bool = True, normalize: bool False) -> torch. Tensor: в**
Args:
q: queries, shape (batch_size, num num_heads, seq_len, dim_qk)
k: keys, shape (batch_size, num_heads, seq_len, dim_qk)
v: values, shape (batch_size, num_heads, seq_len, dim_v) qk_scale: scale factor for queries and keys
h = q. shape[1]
if _ _ecay:
11
slope_ (h, q. device, q. dtype)
12
else:
13
s = no no.slope_tensore (h, q. device, q dtype)
14
15
output lightning_attn.ops.
qk_scale, v,
16
normalize:
17
norm = torch. clamp_mint
18
torch. einsum
"nhld, nhld->nh1", q, k. cumsum (2) qk_scale
) 1e-6
21
)
22
return output norm. uns.ueeze (-1)
23
else:
return output

---

Published as a conference paper at COLM 2024
References
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. ar preprint ar Xxxe:2355.13245, 2023.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitectures and algorithms. ar ar preprint arXix:2401.1299 2024.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arvi preprint arXit:1607.064500 2016.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document trans- former. Xrv preprint ar Xit::004.005100 2020.
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Gpt-neox-20b: An open-source autoregressive language model, 2022. URL 06745
https://arxiv.org/abbs/
Hanting Chen, Zhicheng Liu, Xutao Wang, Yuchuan Tian, and Yunhe Wang. Di- jiang: Efficient large language models through compact kemelization. arXiv preprint arXio:2403.19928, 2024.
Sehyun Choi. Cross-architecturre transfer learning for linear-cost inference transformers, 2024. UtL/tttpss//raxt.orr/////////// 02684
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. ar preprint #rXiv:2009.14777 2020.
Tri Dao, Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint ar Xix::2212.14000 14052, 2022.
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset of information-seeking questions and answers anchored in research papers.
In Proceedings of the 2021 Conference of the North American Chapter of the Association for
Computational Linguistics: Human Language Technologies, pp. 4599-4610, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v112000 mai.....m.a...mm...35355.5 URL https://aclanthhl..og org/ maaccl-main... 3iiain3.65
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru
Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. ar Xiv preprint ar Xit::2402.19427 2024.
emozilla.
Dynamically scaled rope further increases strength of retaining walls, http:a/lll..aa//cmoee///////rrr/ Reddit post.
2023. URL. httpps:/www. reddit.
Leo Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL https://zenodo.. htt.r//cnoo.oorreeerre////////50
12

---

Published as a conference paper at COLM 2024
Alexsandar Botev Griffin Team, Soham De, Samuel L Smith, Anushan Fernando, George- Christian Muraru, Ruba Haroun, and Leonard Berrada et al. Recurrentgemma. arXiv preprint ar Xio::444..07839 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arv preprint ar xrXi::2312.0075,, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arx preprint arXiv:2111.00396 2021.
Suchin Gururangan, Mitchell Wortsman, Samir Yitzhak Gadre, Achal Dave, Maciej Kilian, Weijia Shi, Jean Mercat, Georgios Smyrnis, Gabriel Ilharco, Matt Jordan, Reinhard Heckel, Alex Dimakis, Ali Farhadi, Vaishaal Shankar, and Ludwig Schmidt. OpenL a minimal but performative language modeling (lm) repository, 2023. URL https://gi thub. com/ mLLounndations/openn./// GitHub repository.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, (8):1735-1780, 1997.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint ar arXiv:2310.06825, 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. ar Xiv preprint ar Xiiv:2001.08861. 2020.
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into rnns. ar Xiv preprint ar XXix:2103.13076 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The NarrativeQA reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018. doi: 10.1162/
tacl a a00020. /tt/:tttss://caanthhhllo http:/antthoooooooog00001/11111 org/Q18- 1023
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint ar arXiv:2310.01188, 2023.
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34:2441-2453 2021.
Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. In
Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 10236-10242, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/ v1 21/2022.emnlp-main 697. URL htt.hhsss//aaaaatoollooooo hwwpoolooyooggg///// wwwporggg/222// emw.p-ainnn....i........ 697
Jean Mercat. Higher order linear transformer. arXiv preprint ar arXio:2010.14816 2020.
Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint 0rX14:2404.07143, 2024.
13

---

Published as a conference paper at COLM 2024
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay, The refinedweb dataset for falcon lIm: outperforming curated corpora with web data, and web data only. ar Xiv preprint ar Xii::2366.01116, 2023.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing rnns for the transformer era. ar Xiv preprint ar Xiw:2305.13045 2023a.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arxiv preprint ar Xr1o:223090007 000000 2023b.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. ar Xiv preprint ar Xio:2103.00114, 2021.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolu- tional language models. Xiv preprint ar Xiv:23002.10866, 2023.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling trans- former inference. Proceedings of Machine Learning and Systems, 5, 2023.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. ar Xiv preprint xrXiv:2210.10340, 2022a.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. arviv preprint 2rxiv:2002.08771, 2022b.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling, 2023. URL http::///aaaxi..rg/bbsbss//
04823
Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. Xiv preprint ar wxxi:24401.04633 2024a.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear rnns with state expansion, 2024b. URL htthttsp////rxivoorr abs/2404.077 07904
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 12007-12021, Abu Dhabi, United Arab Emirates, December 2022a. Association for Computational Linguistics. URL htttss::///cclantoloogy 823]
org/2022.
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, et al. Scrolls: Standardized comparison over long language sequences. ar Xiv preprint ar Xiv::2200 0035333, 2022b.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXio preprint Xiv:2101 .44.64104.09644, 2021.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. ar Xie preprint arXit:2307.088221 2023.
14

---

Published as a conference paper at COLM 2024
Philippe Tillet, H. T. Kung, and David Cox. Triton: intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN Interna- fional Workshop Machine Learning and Programming Languages, MAPL 2019, pp. 10-19, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450367196 doi: (4.14453315508.33222
URL https: org.0/..rg/11..0 1145/33151338 3329973
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Sounyy Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. ar preprint arXiv:2307 XXX0:2207.00228, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp- 3-19, 2018.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective long-context scaling of foundation models, 2023.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training, 2024. URL http:t/s////rri
abs/2312.06635
Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXivv2105.14103, 2021.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry. ar arxiv preprint arXio:2402.04347, 2024.
15

---

Published as a conference paper at COLM 2024
A Attention Approximation
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 meaningful comparison difficult. We represent non-normalized comparisons in Figure 4 It represents the cosine similarities and singular value distances between the atfention 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 color scale. In Figure 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 normalized. In Figure 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.
SID PATT
Figure 4: Representation of the cosine similarity and the distance between the singular values of the softmax atlention matrices compared the SUPRA attention matrices.
60
Figure 5: Representation of the cosine similarity and the distance between the singular values of the normalized softmax attention matrices compared to the normalized SUPRA attention matrices.
16


Reviews:
[{"review_id": "eI6lcRWppj", "reviewer": "Reviewer_JkbR", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n---\n\n> 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).\n\nAlthough 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.\n\nIn 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.\n\n\n> I wonder whether the kernel function of relu (following T2R) is really necessary, compared with using an identify function as in RetNet.\n\nWe chose the ReLU kernel function because it allows non-negative attention factors which is desirable given that softmax attention is also non-negative.\n\n> It would be better to compare with the following baseline\n\nThank 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.\n\nUnfortunately, 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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. \n* it would be better to compare with the following baseline [1]\n\n[1] The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry", "extra_scores": {}}, {"review_id": "Ds4W3dGe8e", "reviewer": "Reviewer_TZps", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "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.  \n\n---\n\n> Kasai et al. is actually doing the same thing [...] The experiments conducted in this paper are not newly enabled because of the method proposed.\n\nWhile 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).\n\nIn 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.\n\n> The ablation test shown in Table 2 does not seem to support that the proposed method has any specific strength.\n\nOur 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.\n\nHaving 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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.\n\nMoreover, the ablation test shown in Table 2 does not seem to support that the proposed method has any specific strength.\n\nOverall, the narration of this paper is messy. It indicates that the authors lack a well-polished story.\n\nOne 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.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "V5qlNxbyEG", "reviewer": "Reviewer_JrUb", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n**Kernel Choice**:\nThe 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!\n\n**Regarding Table 2**:\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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. \n- 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).  \n- 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. \n\n**References**\n- 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.\n- 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.", "extra_scores": {}}, {"review_id": "XbRqsIFAmM", "reviewer": "Reviewer_XQ7B", "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.  \n\nThis 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. \n\nMore interestingly, this work also demonstrates some drawbacks of linear models, for example the incapability in long-text modeling and incontext learning.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n> I found the paper lacking an important baseline: continually train the original model [...]\n\nGood 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.\n|Model|New tokens|HS|PIQA|WG|ARC-E|ARC-C|MMLU|Average|\n|--|--|--|--|--|--|--|--|--|\n|Llama2 (base)|-|76.0|79.1|69.1|76.3|46.3|45.9|65.4|\n|Llama2|20B|76.3|78.8|70.1|76.8|46.2|43.8|65.3|\n|Llama2-SUPRA|20B|71.8|78.6|65.8|71.1|39.5|24.9|58.6|\n\nAll 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). \n\nThis 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.\n\n> The perplexity values should be reported for all experiments. \n\nPrevious 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.\n\nBelow 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). \n\n|Model|8|32|128|256|512|1024|2048|\n|-|-|-|-|-|-|-|-|\n|Llama2 (base)|2.76|2.70|2.04|1.94|1.85|1.82|1.78|\n|Llama2-Continued-20b|2.76|2.69|2.05|1.94|1.84|1.82|1.79|\n|Llama2-SUPRA|2.86|2.79|2.18|2.07|2.02|2.01|2.00|\n\nThis 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.\n\nIn the final version, we will include perplexity values for all our other experiments and models.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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.\n\n2. The perplexity values should be reported for all experiments.  Perplexity is still one of the most important metrics.\n\n3. 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:\n\n1. 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.\n\n2. In table3, why do you only report the results of HellaSwag and ARC? Do you have a specific reason? What about other metrics?\n\n3. Could you release an an anonymous huggingface repo for the 7B mamba model?  \n\n4. It would be interesting to use the attention distillation loss proposed in Hedgehog for learning the feature map?\n\n5. It would be interesting to try data-dependent decays proposed in Gated Linear Attention. https://github.com/sustcsonglin/flash-linear-attention\n\n\n \nMissing reference: (although some are released after the submission deadline, it is beneficial to discuss in the next iteration): \n\n- linear attention with data-dependent decays [1, 2, 5].\n- Fine-tuning attention-based LMs into linear attention [1, 3, 4].\n\n[1] Fine-Tuning Pre-trained Transformers into Decaying Fast Weights. https://aclanthology.org/2022.emnlp-main.697.pdf\n\n[2] Gated Linear Attention Transformers with Hardware-Efficient Training. https://arxiv.org/abs/2312.06635\n\n[3] Cross-Architecture Transfer Learning for Linear-Cost Inference Transformers. https://arxiv.org/abs/2404.02684\n\n[4] DiJiang: Efficient Large Language Models through Compact Kernelization. https://arxiv.org/abs/2403.19928\n\n[5] HGRN2: Gated Linear RNNs with State Expansion https://arxiv.org/abs/2404.07904", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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).
Decision: Accept