Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency
Authors: Naoki Nishikawa
Rei Higuchi
Taiji Suzuki
Venue: NeurIPS 2025 poster ()
Keywords: Degrees of Freedom
Linear Attention
Primary area: theory

Abstract: Linear attention has attracted interest as a computationally efficient approximation to softmax attention, especially for long sequences. Recent studies has explored distilling softmax attention in pre-trained Transformers into linear attention. However, a critical challenge remains: *how to choose the feature dimension that governs the approximation quality*. Existing methods fix this dimension uniformly across all attention layers, overlooking the diverse roles and complexities of them. In this paper, we propose a principled method to automatically determine the feature dimension in linear attention using the concept of statistical *degrees of freedom*, which represent the effective dimensionality of the inputs. We provide a theoretical bound on the approximation error and show that the dimension chosen by our method achieves smaller errors under a fixed computational budget. Furthermore, we introduce an efficient layerwise training strategy to learn nonlinear features tailored to each layer. Experiments on multiple pre-trained transformers demonstrate that our method improves the performance of distilled models compared to baselines without increasing the inference cost. Our findings also provide insight into how the complexity of the attention mechanism evolves across layers.

Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency
Naoki Nishikawa The University of Tokyo, RIKEN AIP wcw. u-tokyo waooky..cc.. jc.
Schiikwwa-aakk2220 ecc.
Rei Higuchi
Taiji Suzuki
The University of Tokyo, RIKEN AIP rriih4eg, ecc. u-ckyo-ooko..co ac. jc.jp
The University of Tokyo. RIKEN AIP
taijißmist. i. la-tokyo.ac.
Abstract
Linear attention has attracted interest as computationally efficient approximation to softmax attention, especially for long sequences. Recent studies have explored distilling softmax attention in pre-trained Transformers into linear attention. How- ever, critical challenge remains: how to choose the feature dimension that governs the approximation quality. Existing methods fix this dimension uniformly across all attention layers, overlooking the diverse roles and complexities of them. this paper, we propose principled method to automatically determine the feature dimension in linear attention using the concept of statistical degrees of freedom, which represent the effective dimensionality of the inputs. We provide a theoretical bound on the approximation error and show that the dimension chosen by our method achieves smaller errors under fixed computational budget. Furthermore, we introduce an efficient layerwise training strategy to learn nonlinear features tailored to each layer. Experiments multiple pre-trained transformers demon- strate that our method improves the performance of distilled models compared to baselines without increasing the inference cost. Our findings also provide insight into how the complexity of the attention mechanism evolves across layers.
Introduction
Transformers have become the standard for sequence modeling across diverse domains such as natural language processing (Vaswani al. 2017). computer vision (Dosovitskiy 2020), and speech processing (Dong et al., 2018). A key factor in their success is the attention mechanism, which effectively aggregates the information from input tokens. However, the standard softmax attention requires computing pairwise interactions between all tokens in a sequence, resulting in quadratic time and memory complexity with respect to sequence length. This scalability issue poses significant challenges for large-scale applications.
To address this limitation, numerous efforts have been made to design more efficient alternatives. One prominent approach is linear attention, which approximates softmax attention by replacing the kernel between queries and keys with inner product of finite dimensional features. This reduces both time and memory complexity to linear in the sequence length, enabling scalable inference. The idea was initially proposed by |Katharopoulos al. (2020), and has since been extended in subsequent works (Peng et al. 2021 Choromanski et al. 202 Qin et al. 2022). Recent architectures based on state space models (SSMs), such Mamba Guu and Dao 2024).
also closely related to linear
attention (Dao and Gu 2024 Wang et al. 2024 Sieber et al. 2024 Han et al. 2024).
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Fead Forward Network
(q.
Ra(n. R,)
Feed Forward Network
2- K (q.
Σッ Kz(9, ky.)
Thane
Softmax Attention
[(a,)) exp
(g(,,)) (x) *2(y) 0, 32 g.
Linear Attention
Distil
Σz: Gram Matrix
M2 tr En(E2 + AZ)
OLLM
Fead Forward Network
Σ
K k(q.k)) Ly= ((,,k))
K kg)
Feed Forward Network
Trans forteer
スリー (i ky)
Softmax Attention
K(a,y) exp
(z, tx (x) (
Linear Attention
Distill
O(
Σ, Gram Matrix
vr Ex + X1)
O(LM
Figure 1: An overview of our method. We treat exp x2 (xx/////ve g/Vd), which appears in softmax attention, as kernel K (x, y), and perform distillation by approximating it with the inner product of nonlinear features. The required feature dimension to achieve certain error varies depending on the input distribution, and this can be calculated using degrees of freedom derived from the Gram matrix. When the error level À is specified, our method can automatically select the feature dimensions for linear attention, resulting in different feature dimensions for each layer.
Since linear attention is an approximation of softmax attention, one can directly distill the pre-trained softmax attention into linear attention. Even though prior studies (Chen al. 2024 Wang et al. 2024) have achieved some success in this approach, a critical challenge remains unresolved:
How should we choose the feature dimension that governs the approximation quality?
Most existing works pre-determine the feature dimensions and fix them across all layers, which ignores the diverse functional roles and input distributions of different attention layers, As observed in prior studies (Arora et al. 2018; Ravichandran et al., 2019 Suzuki et al.) 2020 Massaroli et al., 2024), the complexity of the roles played by each layer varies greatly. It is thus reasonable to expect that selecting the feature dimension adaptively, accordance with each layer's s complexity, can improve the performance of the distilled model.
In this paper, we propose a principled approach to determining the feature dimension linear attention. We begin by theoretically deriving a bound on the approximation error and show that the number of features required to approximate the original attention kernel is governed by the statistical degrees of freedom (DoF). Based this insight, we develop method to automatically select the feature dimension for each layer by estimating its DoF. Because the DoF captures the effective dimensionality of the input distribution, our method enables a data-adaptive and layer-specific allocation of feature dimensions.
Although our theoretical analysis identifies the optimal feature distribution, this distribution is generally intractable to compute. To obtain linear attention that achieves optimal approximation accuracy, we introduce training strategy for learning nonlinear feature maps. Instead of relying on costly end-to-end training with the pre-training objective, we propose to train each layer independently. This layerwise training approach significantly reduces computational cost while allowing each attention layer to learn features tailored to its input distribution.
We validate our approach through extensive experiments on GPT-2 and Pythia-IB. Our results show that (1) the optimal feature dimension varies substantially across layers, (2) our method improves the accuracy of distilled models over -ied-dimension baselines, and (3) our layerwise feature learning yields competitive performance with significantly reduced training cost.
Our contributions can be summarized as follows:
We propose principled method for automatically selecting the feature dimension of linear attention based on the statistical degrees of freedom, which reflect the complexity of each attention layer. We provide theoretical guarantees on the approximation error under this selection scheme.
We introduce an efficient layerwise training strategy to learn nonlinear features tailored each layer, which substantially reduces the computational cost compared to end-to-end training.
We empirically demonstrate that selecting feature dimensions by the proposed method improves performance of the distilled model and yields results comparable to the original models. The experimental results also offer insight into how attention complexity varies across layers.

---

Other related works. Some studies aim to distill the attention mechanism in pre-trained Trans- formers into models based on linear attention. Chen al. (2024) propose a method to distill softmax attention into linear attention by using quasi Monte Carlo method, which is more accurate than standard Monte Carlo sampling- [Wang al. (2024) explore distillation from transformers to Mamba taking into account the similarities between Mamba and linear attention. Furthermore, [Bick] itt al. (2024) and Ralambomihanta et al. (024) propose methods for distilling Transformers into SSM-based models, and Kasai et al. (2021) develop method to distill Transformers into RNN-based models. all of them, methods for selecting feature dimensions based on the complexity of the kernels have not been investigated.
Among numerous studies on the distillation of sequence models, including Transformers (e.g.. Wang et al. (2020b). Yang et al.] (202T), Massaroli et al. (2024) and (Sakamoto and Sato (2024) are particularly relevant to our research. These studies focus on distilling SSMs into smaller SSMs, and address the selection of state dimensions in the distilled models. However, their methods are highly dependent on the properties of SSMs and cannot be directly applied to the distillation from softmax attention to linear attention.
Wang et al. (2020a) proposes a method make softmax attention more efficient, and focus on the Tact that attention matrices tend to be low-rank. They propose compressing the keys and the values from d to L' x (L' < L), where is the sequence length and d the head size. Our method is similar in that it focuses on the effective dimension calculated from the Gram matrix of the attention kernel. However, we use it for dimension selection in linear attention instead of compressing the keys and values. Furthermore, our method significantly differs from theirs in that the feature dimension is automatically selected using data and takes different values for each layer.
Notations. Let μ be probability measure on Rd, We define (2(1) as the space of functions f Rd + IR such that f p(d)) < co. We denote J fz) as (===l((())) and J f())g()))(dz)) as (f,9) Lي(ي)- Lo For probability measures f41 and μ12 on IRd we denote the product measure as f41 @ 122- For bounded linear operator A X + Y, we denote the operator norm as Alllo For E N, we define [m] == 21.
2 Attention Mechanism and its Linearization
In this section, we first introduce the regular attention mechanism, and then explain linear attention, which is approximation of the regular one. For simplicity, in this section, we explain only single- head cases, but the description can be easily extended to multi-head cases as well. Moreover, we focus on unidirectional cases, which are commonly used in language models.
IT.] be a sequence of vectors (called tokens), where I; E is the i-th vector. The attention mechanism computes a new sequence [У1 y....., .........l yr. by taking a weighted sum of the projected tokens as follows:
Attention mechanism. Let [1i, T2,........ ,,,,,..,,,,,,
K Klyi, kig) E K kq.))y))
Uj. K(x,y)== exp(x=>>>>) (/pDa///Va).
y/4
where h; = W 3y 9: == W x,, 24 == WV T,, which are called the key, query, and value, respectively, and K,9, WV E d are the learnable parameters. The kernel K(x, y) is referred to as the attention kernel.
The time and memory complexity of the attention mechanism is (II2d) and LLI)+ respectively. Since we need to compute the inner products between every pair of keys and queries, the cost increases quadratically with the sequence length L. Although the attention mechanism is powerful, this becomes a drawback when dealing with the long sequences typical language and speech.
Linear attention. To deal with large computational cost, in linear attention, the computation of ① is simplified by approximating K (z, y) with a inner product of finite dimensional feature maps. Specifically, we suppose that K has a form K (x, y) = E.-r[Ø(zz (æ; 2)0(y; 2)], where T is a
probability measure and 2(=;2): IRd + IR is a feature map. Then, we can approximate K using finite
i..d. samples 21...., 2M T, i.c., K K(x v) R 1 Em=n o(2; 2m ddy; 2m) = (2x) (v)), where
M
3

---

((x) = VM ((q, ky) with ((q:) [kk)).
Linear attention approximates the output of ① by replacing
This simplification drastically reduces the computational cost. Specifically, we can compute the outputs in the following form:
B. ((g)) f (k2 (q)) A ((a))
D (kg)) (g:) (g)))
dj
A
(kk)), B,
(kj)o,
Since A, and B; (i = I,... L) can be computed recursively over i, the computational complexity with respect to input length L is significantly reduced from quadratic to linear: the time and memory complexities become O (LMM)) d) and OLLM + Ld + Md), respectively.
There multiple options for the feature map O. The simplest approach involves using the decompo- sition exp
exp
1 2vd
exp
||
exp
||9 ||g-A|| 2v2
and subsequently applying Bochner's the-
2vd
orem to break down the final component, exp
l|1 into random Fourier features. While this 2v
method appears convenient, it suffers from learning instability caused by the appearance of negative values in sin and cos functions. To address this issue, [Choromansii et al. [2021 propose using the de- composition of exp
k/ vā) Ez~N((01aa) (o(g;)) @(k; 2)], where c; 2) = exp diya 2vd This feature map is referred to as the Positive Random Feature (PRF). Because its values are consis- tently positive, it facilitates more stable learning.
|
3 Optimal Distillation from Softmax Attention to Linear Attention
In this section, we develop a theoretical framework and practical algorithm to select feature dimensions for linear attention, followed by method to train these features efficiently.
3.1 Theoretical Background: Degrees of Freedom for Kernel Approximation
In this subsection, we lay the theoretical groundwork for selecting the optimal feature dimension in linear attention. Our analysis builds on kernel approximation theory and introduces the concept of statistical degrees of freedom (DoF), which quantifies the effective dimensionality required to approximate the attention kernel.
Finite dimensional approximation of kernel. To discuss generally, we consider continuous positive definite kernel K : R2 x IRd + R, and suppose that K admits the representation
K (x, v) Egwr (0(:: 2 2) =)],
where r is a probability measure on a measurable set 2, and Φb Rdd -> is a feature map. We further assume that the function + K(z, æ) is integrable with respect to and that the map (I, 2) + ((x; 2) is square-integrable with respect to p 8
The simplest way to approximate K is to use the empirical kernel K' defined as K'(z,v) = ΣM δoz; 2/m (() 2'mn); where 2/ ,f. 2'y are i..d. samples from According to the strong law of large numbers, K" converges to K almost surely as M + oo for any fixed x and Rahimi and Recht| (2007) show that, for any translation-invariant kernel K and compact set Z Rdd K" uniformly converges to K at a rate of O(log M))
Rather than focusing on the global error bound over a compact set, we strive to achieve a more accurate approximation by leveraging the intrinsic structure of the input vectors. In other words, we assume that the input vectors and are generated from a probability measure p on Rd and adjust the approximation scheme utilizing the information of p. To this end, we consider approximating K using i.i.d. samples 21, 2,...., drawn from a distribution with density w.r.t. the measure T:
M
R(x,)) v) = M m=1
Σ
blz; Fm )olyn zm) = 00x;; 2) ( z),
((2m)

---

= le........ zM]" The kernel K defines a finite- dimensional Reproducing Kernel Hilbert Space (RKHS) H. which is expected to be approximation of the RKHS H. defined by K. Bach (2017) analyze the approximation error between H. and H. and show that is characterized by the value Nº.x (A > 0) defined as
where 2();2) (M
1 =E2 ((2)
n. sup
(0(:2), (E) (()))++++)((())))))))))))))))))) t2(d))
where 2: L2(p) L2( an integral operator defined by (#f))z)= (K f)).)))) Specifi- cally, they provide the following result.
Proposition (Proposition in Bachh (207)). Suppose that 21,..., 2m i.i.d. samples from
the distribution with density 9. Then, for any E (0,1), M 2 5.Ng,1 log 16N.22 _²² it holds
1. Em= (2 -1 2m M
3l)mlll.e (l2))) I 2tr and
3m ((mm) 1/2
sup inf f-llf S1 11 a2
f
Σ
ol: 2m)
$41,
1.(0))
with probability at least 1
This proposition demonstrates that the necessary number of features M to achieve the approximation error À is determined by Ng,,, which depends on the density q. In other words, by effectively selecting the density q, is possible to achieve an approximation error of with smaller .M.
One of the essential part of the proof of Proposition iss to bound the error between the operators ) (corresponding to K) and its empirical approximation Σ L2(p) Le(p) (corresponding to K) defined as (z))(z)
(K (ar, :), f) La(p)- In the proof of Proposition Bach 2017) analyze the error
between Σ and Σ as follows.
Lemma Bach (017)) Let A (E+ 1/2 -E-EE++++ (I[E+AI) AI) 1/2, For any X, 0, it holds
6
Mt2 /
P|| P[| Axllop >
1+
12 log"
(1+
Mr Ny.
exp
N₂ (1 + }1)
This lemma provides the upper bound of the largest difference of eigenvalues between (E + - + A7) and (E + XT)- -/)-1/2 which controlled by Ng.
Error bound between K and R. Proposition shows that the elements of H. can be approximated by those of H in terms of the (p) norm. However, for the purpose of distilling softmax attention into linear attention, it is essential to have error bound between K and K and, ultimately, to ensure that the output of the attention mechanism is well-approximated. To this end, we establish the following theorem using Lemma
Theorem 3. Let 8 E (0, 1), O, and E (0,3). Suppose that Z following two items hold:
4N log
64N,A 8t
Then, the
(i) holds
K K | SA SA + +Pck
L2(08p)
with probability 1 28, where (()) c12) are constants depending on K.
(i) Let E 1/2] and h E T.) We define a map v Rd RR = (E -h. Then, it holds
- u ((x)
S <V2(12-0= - c22 -1]bll
13(0)
with probability 1 1-6 8, where (3) is constant depending on K.
5

---

The proof can be found in Appendix A We emphasize that this theorem can be applied to any positive definite kernel represented as (2). Item (i) provides a bound on the L2 @ p) norm of the difference between the kernels K and K, indicating that the approximation error is governed by Nº.x- This suggests that careful choice of the sampling density q efficiently reduce the approximation error. Item (ii) bounds the error in kernel-based integration against an arbitrary function v. Notably, the integral (a()) K (z; y y)dp(z) corresponds the attention-weighted guarantees that the required feature dimension M for accurately approximating attention outputs effectively controlled by Ng.A. Finally, we remark that the space {(E + A E (p)3 strictly contains I,2 (p) when
weighted sum in ( Thus, this result
implying that our bound applies functions, including those that may not be square-integrable under p.
broader class of value
Remark 4. We compare our theory with two previous studies on the approximation accuracy of linear attention. (i/Choromankki et al. (202T) analyze the approximation quality of PRF-based linear attention in terms of the sup-norm error. They showed that the required number of features M scales as Oid log d). In contrast, our Theorem bounds the L p)-norm capture the structure of the input distribution. As a result, the approximation error to depend on the DoF rather than the dimension d. (ii) Luo al. (202T) shows exponential dependence on the norm bound R of queries and keys. In our analysis, such dependence is absorbed in constants like CK and 1K We do not aim to reduce this exponential dependence, but instead focus on selecting the feature dimension per layer to minimize approximation error under a fixed computational budget.
Degrees of freedom and their optimality. Next, we aim to optimally reducing the required value of Nq.x for a fixed A. Specifically, we consider the density qA defined as follows:
(x) Lr EEE + XT)-I (0(52), (E XT) ((;;)( A
La(e)
By setting 4- QA. we can see that Nar.A N = tr E + AI) The value NW is called the degrees of freedom (DoF).
When use the i.i.d. samples 21,.... 2M drawn from distribution with the density qx. the required number of samples M in Proposition and Theorem 3 is proportional to N log N which is known to be the optimal (Bach (2077). Proposition 3). In particular, the degrees of freedom smaller than N1.x- This suggests that the approximation error can be reduced by choosing the distribution q depending the input distribution p compared to the case obtain the random samples from the distribution T.
always
The findings above can be summarized as follows:
The feature dimension M should be set proportional to the degrees of freedom N tr(E(E + 11) up to logarithmic factor.
The degrees of freedom Nr depend on the input distribution p. implying that the optimal feature dimension M varies across layers.
Optimal feature efficiency can be achieved by sampling 2M, from q dr, rather than from r. The density & also depends on the input distribution p.
3.2 Feature Dimension Selection via Degrees of Freedom
Given that the degrees of freedom N determines the minimal required number of features, we now return to the attention kernel K (a,, v) exp and present a practical procedure to estimate it from data and use allocate layerwise dimensions. The overview of our method is presented in Algorithm
Since N is defined the expectation over the data distribution, it cannot be computed in practice. Therefore, we approximate it using finite samples. Specifically, we define N₂ as an approximation
This correspondence becomes exact when K invertible. Letting p' = w denote the pushforward measure of p by W the ineettio-wwiighee sum of values is represented as J ((k) v)dp' with k W v(k)= W WV(WK)-1k

---

Algorithm Selecting the Feature Dimension
Input: Set X' of T sequences with length L, sample size J N, tolerance À 0, cost N, original Transformer with layers and H heads.
Fed the sequences in X' into the pre- trained Transformer and collect the queries Qs,h and keys Ka,h of all layers E [S] and heads E [H].
for 1 to S do
for = to H do Randomly sample Fi....;2J from Qa U Ka,h and compute the Gram matrix Σ, R
Compute ((,,)) tr Eo,hh (Ca, A1)- end for Obtain ( maX N N(a.h) end for
Set t- = C (5-
Return: Feature dimension M, round(t (N() for layers I....,S..
of N given by Ñ₂ = =r AT)-1 where E. R (J N) a gram matrix of K, i.e., K (z,, [L)e which is the approximation of E based finite samples ; 2. p. For the attention kernel, the inputs are the queries and keys of the attention layers. Therefore, compute the approximated degrees of freedom N₂ for each layer follows:
1. Prepare T sequences of length L, and compute queries and keys for all sequences and tokens, resulting in LT queries and LT keys, respectively.
2. From the collected 2LT queries and keys, randomly sample / elements 21...., ZJ.
3. Compute approximated degrees of freedom N₂ defined above.
The degrees of freedom obtained through this procedure vary across heads and layers, shown in Section 4 In typical Transformer implementations, the heads within the same layer are processed using shared tensor, so it is desirable that the feature dimensions be identical as well. Hence, define the degrees of freedom (7()) of the s-th layer as the maximum degrees freedom among the heads.
According to the results of Theorem (excluding the logarithm), we set the feature dimension of the s-th layer to M = ((1) for some > 0. Then, the computational cost of inference matches
that of linear attention with fixed feature dimension given by 1 -5 M -- E N (+))) In our experiments, we fixed in advance. Given computational cost C,, we determine the constant t such that the inference cost of the model matches the computational
when we the feature
dimension is uniformly set to C' across all layers, which leads t- C ( L N7(8)
3.3 Layerwise Training of Features for Kernel Approximation
Theorem guarantees that we can attain the optimal error between K and K by choosing the density q appropriately, However, it is difficult to obtain such density in practice. To obtain features 2; and qzx)) as) (i E [M] that approximate the attention kernel as accurately as our theoretical bound, we propose to learn them from the training data.
One possible strategy is train the features using the same loss function as the one used in the standard wee-trainnng of Transformers. That said, this approach is computationally expensive, as requires computing gradients across all layers simultaneously. Instead, we propose to train the features in layerwise manner. Specifically, we consider two types of loss functions: (i) 1 loss and (ii) softmax loss. describe the details, let us assume that the set of queries Qt (qt,1,--- (..1) and keys Kr (ke.1,...,,k kt, of I sequences, each of length L, are given.
7

---

Table 1: The dimension of features selected by method. The feature dimensions of each layer determined to match the average cost. For each row, the top three largest dimensions are highlighted in bold, and the three lowest dimensions are emphasized with underlines.
Model
Cost
2
3
4
10
11
12
13
15
16
GPT-2
64
130
182
35
44
42
65
92
33
28
34
46
39
Pythia-IB
128
277
138
275
132
204
167
115
142
231
64
96
73
40
52
34
This loss function tries to minimize the error between the attention kernel and its approximation K in terms of the L norm. That is, the features 2 = (Z1....; 22As) are trained using
ii L loss.
the loss function defined as {(2 1.r Σeerrnee(4)ree (K(q,,, ke, K 2.);
(ii) Softmax loss. We train the features make the attention matrix and its approximation closer in termss of the cross entropy loss. In other words. we train the features 2 = (Z1,..., =MM de- Pr (gt,1, K.,) log Pe (4a,, K4,2,00 where pe (11,, Ka) = R Rlgn, (.
fined as {(z) = - ELET
exp(y exp(s.,ik,,1)) Σg- ex.
pr (qr,,, Kri 2. = Σラ- (kaai,)))
In general, L loss is more efficient than Softmax loss because it does not require additional nonlinear operations such as exp or normalization. In Section we compare the performance of the two losses (i) and (ii) with the case where the features are directly learned using the pre-training task. Then, we demonstrate that learning with our proposed loss is as effective training using pre-training loss.
4 Experiments
To evaluate the effectiveness of our method, conduct experiments on two types of pre-trained Transformers: GPT-2 (Radford et al., 2019) and Pythia- (Biderman et al. (023). As non non-lineat linear features for the linear attention, we use PRF (Choromanski al.) [021). l). which we describe in Section 2
For dimension selection and training of features, we utilize the Wikipedia datase
4.1 Optimal Feature Dimensions Vary across Layers and Heads
First, apply Algorithm to the two pre-trained models and determine the feature dimensions for each layer. We set the cost C' be the same as the head size (64 for GPT-2 and 128 for Pythia-IB), following prior work that sets the feature dimension of each layer to match the head size. Furthermore, we set À 2-4 for GPT-2 and À = 2 for Pythia- IB. We present the selected fea- ture dimensions Table [ We also show the estimated DoF for each head in GPT-2 in Table We summarize the key findings from these results below.
Table DoF with X = 2- of each head. The row and column represent the layer and head, respectively. The maximum value of each layer is highlighted in bold.
Layer
2
5
7
10
I1
4.7
3.6
4.5
5.0
150.0
10.5
28.5
155.1
44.6
173.8
51.1
9.9
16.5
17.5
8.1
88.9
19.5
9.9
3.4
24.5
2.8
2.6
12.1
15.1
3.0
1.4
1,7
39,8
12.1
1.6
1.5
2.5
12.6
1.9
1.8
1.8
11.3
1.8
42.1
2.3
40
7.1
22.6
3.3
31.6
8.9
3.7
3.7
6.9
4.7
$2.4
21.4
7
3.3
6.9
68.9
14.4
43.5
17.2
107.4
6.0
4.3
1.8
32.2
3.0
11.9
33,0
24.9
6.5
19.6
9.3
5.5
3.2
10.2
3.1
5.7
16.7
22.4
3.9
10
18.5
13.1
15.9
22.7
29.9
14.1
20.2
5.1
3.9
7.9
6.6
26.1
19.8
18.2
3.3
8.7
6.7
43.2
2.3
17.1
13.5
12
5.1
21.8
22.5
10.6
30.0
39.8
27.2
31.3
5.0
20.0
2.8
2.6
The evolution of effective dimensionality across layers. The result indicates that the effective dimensionality varies significantly across layers. This emphasizes the difference of the complexity of the roles played by each layer. Taking a closer look at the results, we observe that the feature dimensions selected by our method are generally larger in the early and middle layers, and small in the latter layers. This shows that Transformers engage in complex token interactions from the shallow layers the middle layers, and relatively simple processing the later layers. This observation aligns with the insights obtained in previous studies on sully-commectee//comoluuii neural networks (Arora et al., 2018 Ravichandran al., 2019 Suzuki et al.) 2020). Furthermore, it is worth noting that, between the early layers and the middle layers, there are several layers with relatively small dimensions.
"The dataset can be accessed through the Hugging Face http:://uegggggfococo////aaaaate
legacy- datasyystteets/////o wttaaastssssiikiieiteee wikipedia
8

---

Table 3: The results of distillation for GPT-2. Fix, DoF and DoF Clip in the first column imply the strategies of dimension selection. For all downstream tasks, the evaluation metric is accuracy. The top three higoest- are highlighted in bold.
Strategy
Method
PiQA
logiQA
ARC-E
ARC-C
Winogrande
MMLU
WSC
Average
Original GPT-2
0,5985
0.3103
0.3325
0.3003
0.5122
0.2789
0.6538
0.4266
Diliang
0.5065
0.2550
0.2113
0.2244
0.4846
0.2639
0.4615
0,3409
Performer
0.5468
0.2934
0.3039
0.2747
0.4996
0.2517
0.5962
0.3952
Fix
direct
0.5832
0.3195
0.2921
0.2995
0.5335
0.2552
0.6154
0.4141
softmax
0.5718
0.2673
0.2479
0.3029
0.5020
0.2634
0.6154
0.3958
1,2
0.5822
0.3195
0.2483
0.2773
0.5107
0.2520
0.5962
0.3980
DoF
direct
0.5669
0.3011
0.3241
0.3012
0.5280
0.2712
0.6346
0.4182
softmax
0.5751
0.3026
0.3224
0.2995
0.5328
0.2564
0.6442
0.4190
1,²
0.5664
0.3088
0.2736
0.2824
0.4972
0.2608
0.5865
0.3965
DoF + Clip
direct
0.5892
0.3164
0,3136
0.2952
0.5075
0.2993
0.6346
0.4223
softmax
0.5860
0.3026
0.3401
0.2816
0.4996
0,2832
0.6346
0,4182
1,2
0.5822
0.3164
0.2942
0.3063
0.5091
0.2799
0.5673
0.4079
Comparison of DoF across heads in the same layer. We additionally observe that the degrees of freedom are largely different across heads as well. Interestingly, most of the heads have much smaller degrees of freedom than the maximum of each layer. This indicates that the heads performing complex processing are limited to few within single layer.
4.2 Dimension Selection with Layerwise Training Enhances the Performance
We now evaluate the effectiveness of our method in distillation from softmax attention for two pre- pre-trained models. In these experiments, we examine the following two claims: (i) selecting feature dimensions using our DoF-based method improve performance compared to using fixed dimensions, and (ii) our efficient layerwise training match the performance of full end-to-end training.
We consider three types of dimension selection strategies: Fix is the method that sets the feature dimensions to the same value across all layers, which is conventionally used in existing works. We the feature dimensions in Fix to be the same as the head size, as prior work (Chen al.) 2024). DoF is the method that selects the feature dimensions using Algorithm (i.c., proportional to the maximum DoFs among heads). We also consider DoF + Clip, in which we clamp the feature dimensions to be at most the head size. This strategy is expected to maintain the performance of Fix, while reducing the computational cost.
For training the feature maps during distillation, we consider three types of loss functions: the first is direct loss, which refer the cross-entropy loss for next-token prediction, which provides the best performance in the pre-training task but requires full end-to-end training. The remaining two are softmax loss and L loss, which are the layerwise losses introduced in Section 3.3 and more efficient and allow independent training per layer. We will show that the layerwise training achieves performance comparable the direct loss, while significantly reducing training cost.
As comparison baselines, we evaluate against Performer (Choromanski et al. 2021) and Diliang (Chen et al. 2024). Both methods employ PRF for linear attention but do not incorporate any data-driven dimension selection. Diliang additionally utilizes quasi Monte Carlo sampling of features. The learnalle parameters of DiJiang trained the same dataset used in distillation.
To assess downstream performance, we fine-tune distilled models (as well as the original Trans- formers) suite of seven tasks: PiQA (Bisk et al. 2020), logiQA (Liu et al.) 2020). ARC- Easy/Challenge (Clark et al. 20018). Winogrande (Sakaguchi et al. 2019), WSC (Levesque et al.) 201T), and MMLU (Hendrycks et al. 2021). Further details on the training settings and datasets are provided in Appendix
Results for GPT-2. We first show the experimental results for GPT-2 Table 3 We outline the noteworthy four observations below.
Enhanced performance through dimension selection: Distilled models using feature dimen- sions selected by Algorithm (DoF) achieve higher or near-equal performance compared to those using fixed dimensions (Fix). In particular, in the average performance across all tasks (see the rightmost column of Table 3), when using direct or softmax loss, the distilled models with DoF significantly outperforms the ones with Fix (c.g., 75 % mitigation of performance degradation

---

Table 4: Comparison of the inference time se- quence with 1,000 tokens for different feature di- mensions (for GPT-2).
Table 5: Comparison of the consuming time per each sample during training for different loss types (for GPT-2).
DoF Dol + Clip
Type of loss Speed (ms/sample)
dircct
softmax
1
Speed (sec/sequence)
15.589 15.379
14.662
109.4
90.4
60.6
Table 6:
The downstream accuracy for Pythia 1B. We use softmax loss for distillation.
Method
PiQA
logiQA
ARC-E
ARC-C
Winogrande
MMLU
WSC
Average
Original Pythia-IB
0.6213
0.2995
0.3295
0.3072
0.5146
0.2883
0.6346
0.4279
Fix
0.5691
0.2965
0.2849
0.2679
0.5051
0.2903
0.5962
0,4014
DoF
0.5860
0.3026
0.3110
0,3106
0.4949
0.2811
0.5962
0,4118
from the original model for softmax loss). This indicates that method effectively captures the varying complexity of each layer, leading to improved accuracy across downstream tasks.
Effectiveness of DoF Clip for efficiency: The DoF Clip strategy. which caps feature dimensions at the head size, achieves comparable or better performance than the fixed-dimension approach. Notably, it also yields faster inference (Table 4 showing that selectively reducing feature dimensions based on DoF makes the model efficient without compromising performance.
Stronger results compared to prior baselines: Our distilled models surpass both Performer and DiJiang, which use fixed-dimension linear attention and (quasi) Monte Carlo feature sampling. Additionally, the average performace our distilled model is comparable to (and sometimes win) the original GPT-2 model. This shows that data-adaptinn adaptive, layer-specific dimension selection combined with learned features yields a more accurate approximation of softmax attention.
Using the softmax loss, distilled models perform comparably those trained end-to-end with the standard pre-training loss (direct). As shown in T5] layerwise training is more efficient than direct training. Although the 1.2 loss yields slightly lower performance than both direct and softmax losses, it still outperforms prior baselines while offering notable training efficiency gains.
Efficient and effective layerwise feature training:
Results for Pythia- We also carried out experiments on distilling Pythia- (ithia-1B.ag In this experiment, the model was distilled using the softmax loss. The results presented in Table The distilled model with feature dimensions selected by our method performs better than the model with fixed dimension average, and become comparable to the original Pythia- Indeed, when comparing the average accuracy across all tasks, the performance degradation from the original Pythia- IB (Fix: -0.0265, DoF: -0.0161) is reduced by 39 %. This result indicates the efficacy of our dimension selection and layerwise training when distilling large models.
5 Conclusion
We proposed a principled method for selecting feature dimensions in linear attention, grounded in statistical theory on finiteddimensional kernel approximation, and demonstrated its effectiveness in distilling softmax attention. Our approach adaptively assigns dimensions to each layer based on the statistical degrees of freedom of attention kernel, and trains nonlinear features in a layerwise manner to reduce computation. Experiments GPT-2 and Pythia- 1B show that method improves performance without increasing inference cost, highlighting the importance of layer-specific design for efficient linear attention models.
Limitation and Future Work. While our method assigns feature dimensions adaptively per layer, we use a shared dimension across all heads within a layer to maintain implementation efficiency. Interestingly, our analysis revealed that the DoFs vary significantly even among heads in the same layer. Exploiting this head-level variability-for instance, through head pruning or sparse feature allocation-coe further enhance performance and efficiency. In addition, although our experiments focus on moderate-scale language models, applying our approach to larger models and models for other modalities remain important direction validate its scalability.
10

---

Acknowledgements
NN was partially supported by JST ACT-X (IPMJAX24CK) and JST BOOST (IPMJBS2418) RH was partially supported by JST CREST (IPMCR22115) TS was partially supported by JSPS KAKENHI (24K02905) and IST CREST (JPMICR2015), This research is supported by the National Research Foundation, Singapore and the Ministry of Digital Development and Information under the Al Visiting Professorship Programme (award number AIVP-2024-004). Any opinions, findings and conclusions recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore and the Ministry of Digital Development and Information.
References
S. Arora, R. Ge, B. Neyshabur, and Y. Zhang. Stronger generalization bounds for deep nets via compression approach. In International conference on machine learning, pages 254-263, PMLR, 2018.
Bach. On the equivalence between kernel quadrature rules and random feature expansions. Journal of machine learning research, 18(21)11-3,, 2017.
A. Bick, K. Li, E. P. Xing, Z. Kolter, and A. Gu. Transformers ssms: Distilling quadratic knowledge to subquadratic models. Advances in neural information processing systems, 38, 2024.
S. Biderman, H. Schoelkopf, Q. G. Anthony, H. Bradley, K. O' E. Hallahan, M. A. Khan, S. Purohit, U. S. Prashanth, E. Raff, et al. Pythia: suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397-2430.
PMLR, 2023.
Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty- Foerty- AAAI Conference on Artificial Intelligence, 2020.
H. Chen, Z. Liu, X. Wang, Y. Tian, and Y. Wang. Dijiang: Efficient large language models through compact kernelization. arXiv preprint arXi::2403.19928 2024.
K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, Gane, T. Sarlos, Hawkins, Q Davis, A. Mohiuddin, L. Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations, 2021.
Clark, Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. 1xXiv:1803.055577 2018.
Dao and Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In International Conference Machine Learning, pages 10041- 10071, PMLR, 2024.
L. Dong, S. Xu, and B. Xu. Speech-transformer:: a no-recurrence sequence-to-sequenee model for speech recognition. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), pages 5884-5888. IEEE, 2018.
A. Dosovitskiy, An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint ar arvi::2010.1 1929, 2020.
A. Gu and T. Dao. Mamba: Linear-ting time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024.
D. Han, Z. Wang, Z. Xia, Y. Han, Pu, C. Ge, J. Song, S. Song, B. Zheng, and G. Huang. Demystify mamba in vision: A linear attention perspective. Advances in neural information processing systems, 37, 2024.
D. Hendryeks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations, 2021.

---

J. Kasai, H. Peng, Y. Zhang, D. Yogatama, G. Ilharco, N. Pappas, Y. Mao, W. Chen, and N. A. Smith, Finetuning pretrained transformers into mns. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630-10643, 2021.
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156- 5165. PMLR, 2020.
H. J. Levesque, E. Davis, and L. Morgenstem. The Winograd schema challenge. In AAAI Spring Symposimm: Logical Formalizations of Commonsense Reasoning, volume 46, page 47, 2011.
Liu, I.. Cui, H. Liu, D. Huang, Y. Wang, and Y. Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08/22 2020.
S. Luo, S. Li, T. Cai, D. He. D. Peng, S. Zheng, G. Ke, L. Wang, and T.-Y. Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. Advances in Neural Information Processing Systems, 34:22795-22808
2021.
S. Massaroli, M. Poli, D. Fu, H. Kumbong, R. Parnichkun, D. Romero, A. Timalsina, Q. McIntyre,
B. Chen, A. Rudra, et al. Laughing hyena distillery: Extracting compact recurrences from convolutions. Advances in Neural Information Processing Systems, 36, 2024.
H. Peng, N. Pappas, D. Yogatama, R. Schwartz, N. A. Smith, and L. Kong. Random feature attention. 2021.
Z. Qin, W. Sun, Deng, D. Li, Wei, B. Lv, J. Yan, L.. Kong, and Zhong. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022.
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
A. Rahimi and Recht. Random features for large-scale kernel machines. Advances neural information processing systems, 20, 2007.
T. R. Ralambomihanta, S. Mohammadzadeh, M. S. N. Islam, W. Jabbour, and L. Liang. Scavenging hyena: Distilling transformers into long convolution models. arXiv preprint arXiv:2401.17574 2024.
K. Ravichandran, A. Jain, and A. Rakhlin. Using effective dimension to analyze feature transforma- tions in deep neural networks. In ICML. 2019 Workshop on Identifying and Understanding Deep Learning Phenomena, 2019.
K. Sakaguchi, L. B. Ronan, B. Chandra, and C. Yejin. Winogrande: An adversarial winograd schema challenge at scale. 2019.
H. Sakamoto and Sato. Data-driven h2 model reduction for linear discrete-time systems. arXiv preprint 2.t.:2200010077744 2024.
J. Sieber, C. A. Alonso, A. Didier, M. N. Zeilinger, and A. Orvieto. Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks. Advances in Neural Information Processing Systems, 37: 141111334444555555
2024.
T. Suzuki, H. Abe, T. Murata, S. Horiuchi, K. Ito, T. Wachi, S. Hirai, M. Yukishima, and T. Nishimura. Spectral pruning: Compressing deep neural networks via spectral analysis and generalization error. International Joint Conferences on Artificial Intelligence, 2020.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need. Advances in neural information processing systems, 2017.
Wang, D. Paliotta, May, A. Rush, and Dao. The mamba in the llaaa: Distilling and accelerating hybrid models. Advances in Neural Information Processing Systems, 37:62432- 62457, 2024.
Wang, B. 7. Li, M. Khabsa, H. Fang, and Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04F7 04768, 2020a.
12

---

W. Wang, F. Wei, L. Dong, H. Bao, N. Yang, and M. Zhou. Minilm: Deep self aelf-attenioon distillation for task-agnostic compression of pre-trained transformers. Advances in neural information processing systems, 33:5776-5788, 2020b.
J. Yang, B. Martinez, A. Bulat, G. Tzimiropoulos, et al. Knowledge distillation via softmax regression representation learning. International Conference on Learning Representations (ICLR), 2021.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims in the abstract and introduction are clearly aligned with the theoretical and empirical contributions presented in the paper (see Section and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper includes a "Limitations and Future Work" paragraph in Section 5 which discusses several open issues.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: The theoretical development includes detailed assumptions and complete proofs, especially in Section Section and Appendix A The derivation builds upon known kernel approximation theory (e.g., Bach 2017), and all critical assumptions are stated clearly.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper provides detailed experimental setups including datasets used, model configurations, batch sizes (Table learning rates, number of epochs, and training strategies (Appendix This level of detail supports reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
15

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
The code is provided the supplemental material. The all data used in the experiments are publicly available datasets (Wikipedia, PiQA, logiQA, ARC-E, ARC-C, Winogrande, MMLU and WSC).
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https
www.eeeeeeoeeemmmeeeeeeeeeee commiissions aessonoolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
describes the training settings detail, including learning rates, number of epochs, batch sizes, and hardware specifications. The dataset and its usage also specified.
Justification: Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: While the experimental results are averaged across tasks and compared with strong baselines, we did not report statistical significance such as emro bars or confidence intervals due computational constraints.
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
The paper provides detailed information on computational resources in Ap- pendix C It specifies the computational resources used for training and inference, including the type of compute workers, memory, and time of execution.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: The research was conducted in accordance with the NeurIPS Code of Ethics. We used only publicly available datasets and open-source pre-trained models. No personally identifiable information, private data, human subjects were involved. We have also ensured that all evaluations preserve fairness and transparency.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
17

---

Answer: [NA]
Justification: Our paper contributes to the efficiency and scalability of Transformer-based models, which does not have direct societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Our work does not involve the release of a new dataset or novel pretrained language model. Instead, we focus on method to distill existing models using already public data and models. As such, the paper poses no significant risk of misuse that would necessitate additional safeguards.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We properly cite the papers or URLs for the datasets and models used in our experiments. The URLs and licenses for the datasets and models are provided in Table I0
18

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com camdsatasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not introduce new assets such as datasets or model code. uses publicly available datasets and pre-trained models, and focuses theoretical and empirical improvements to existing architectures.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cood/mmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
19

---

Answer: [NA]
Justification: This paper does not involve any human subjects, and therefore IRB approval is not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs not used tools for developing augmenting the core methodology. The work focuses on distilling attention mechanisms and does not involve using LLMs in novel or generative capacity within the research method itself.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
20

---

Appendix
A Proof of Theorem 3
We begin by recalling the definitions of some symbols used in the proof.
K : Rdd x Rd IR is the positive definite kernel given by ( y) Pegrro((n) =)o(;;
where is probability measure on measurable set Z, and d: Rd 2 IR is feature map.
E: L2(p) L2(p) is integral operator defined by (z)) (K, Lo
R Rd x Rd
R is the approximation of K, defined as (xx,
M zm), where &M are ii.d. samples drawn from distribution with density q w.r.t. measure T.
1
Em=L, olx; 2m 2m)
M
L2(p) L2(p) is the empirical approximation of E, defined (ff(z) = -, Lz(p)
Na. (A 0) the value defined as Ng,1 supze2 ((अ) (0(s)), (2+ L2(dp)
We restate Lemma
2 which can be found in the proof of Proposition in Bach (017))
Lemma S. Let A ( (1) 1/2(11//------ E) -11-1/2, For any X, > 0, holds
P Axllop > f S S2N (1+ 12 log2 Mt/N,,))
MP2
exp
(1 -t/3)
This lemma produces the following high probability bound. Lemma 6. For any X, 8 E (0. if
4N
64N 8t
M
log
2
it holds
|| | St, and equivalently, with probability at least 1- 6.
115 M) - -
Proof. We first note that
n, (1+ 13 logg2 (1 + MMIN
6
Mt
exp
N1 (1 +133)
2.
(2
6
Mt²/2
S
f2
log 2 (1 + M+/N,
exp
Ng.1 (1,x11+6 + t//3)
(6)
Let us consider
A log
CN
8t2
for some constants B, C > 0 which will be determined later. First, we consider the factor (b). Then, for any / (0,3), we have
B/2
CN 8t2
5t2
B/2 144/3
5t2
B/6"
(b) S exp
log
t/3
CN
CN
where
(4
if 8t /CN,,1 $ 1,
bt
22
if 662 /CN 1.
21

---

Next, we consider the factor (a). Let D E [3/C, oo) be a constant to be determined later. Then, if Nz,2/t 2 D, we have
MI
B
CN
B
CD
B
CD
B
CD
log
log
log
2
log
Ng.A
t
8t2
t
3
3
3
which implies
(a) S 9
log (1 + log
On the other hand, if N2.1/1 S D, then MI/N,, 4.a 1/D since M 2 1, which yields
(a) S + log" (1+ ))
Finally, we choose B,C 0. If we set B = bi, we obtain the upper bound of the probability as
2
(a) 61 28
CN₂. (a).
2
Moreover, have
6
(a) 5 max
9 +
log" 2
1+
log 3
CD
log (1+ }
6
S 9 log" (min {1 + 3
2
log 3
CD
,11 }1)
If we set C== 64, D = 1/2, then it holds D E [3/C, co), and the right-hand side of 3 is smaller than 32, which implies (a) 32. Thus, we set
4Ng. t2
64N 9.a
bi Na.
64N 5t2
M 2
log
2
log
8t2
t
we have ||Ax S f with probability at least - 8.
The result | < implies that
-tI = 10/--11-11/E
EEE 1/2 StI.
From the right-hand side inequality, we have
Σ- 1
Indeed, for any f E L2(d)), we have
(5, + AI) (2-- $))fr) La(dp)
= =II 1/ -2:1-11:: - A1) "00 L2(dp)
= -(E+ f. -1 x/) -)))))))' +XT)113)) L2(dp)
2 0.
since (E + 1/2 is self-adjoint. Similarly, we have -- Σ > -t + AI). This completes the proof.
Below, we provide the proofs of Theorem
by dividing it into items (i) and (ii).
Theorem 7. Let d E (0, 1), A and
If
4 g. 2
64N Aa 8t2
M
log
then, holds
K - K
< A
18
( + 2222)
12(p)
with probability 1 1 28, where of ,0 are constants depending K.
22

---

Proof. Using Lemma 6 we have
K K
Σ E
=tr [I= =)]]
(p(p)
HS
(E+A1)(2 - -
st. tr
(ir 22( (-(E=-))
{]] - + +Ar tr [r + E1
St. [r|]]]] د Σ
AAtr 12+ 2
with probability Lemma 6 implies that | | - st. + lTpp which yields
K R ?
St. t. tr [>]] \T op + Atr [_ =l)
(p20)
11 (1. (1.rr(_]|/=)) + (1 + 1) tr [Σ} + A Atr (x])
SxA (x [
+ ? tr | P1llo
Now, (r[]] is bounded. Indeed, if the Mercer decomposition of K is given by K(æ,v) =
Σ/=1 13 (1) then have
tr|E] Σ F1 = / (EMSs)) appl) / / K(x, x)d((x) <o, 3=1
since x + K (ar, x)
is integrable with respect to p. Moreover, from the Markov inequality, we have
20 [rrt]
tr
tr - =
tr
Therefore, it holds
K K 2 (p8()
S SA. t 4+ tr [Σ] + +2 tr | lp S A A - + +
with probability 1 1-26, where C
tr | and 012) tr | op' This completes the proof.
Theorem 8. Ler ô E (0, > 0, t > 0 and a E [0, 1/2]. If
M 4N₂
64N, 9.a
log Dt2
12
then, for any v: RRd - Tke with |[ + AT) (2(2)
< c0, it holds
If o()) (dp(a)
$ VV((Ca, + +1 ++|(E+ + oll
LP(P)
with probability 1 - 6, where cK is a constant depending on K.
Proof. The left-hand side can be bounded
|ʃ o(z)k(zz (x,))(((()) f ((=))(z
De Do 12(0)
= ||(--2)) 1200)
(2(P)
Since E is self-adjoint operator, we have
[1 ( 20)01(2000 12000
v. (..( - ( - 27)0 -2)0
2)
(5, E)(0)
-(x+ 2 v. (E) =)² A7) -)) -17)))
< <
1 11) 01)00122000
23

---

Now, let us bound
( A1) Let = Σ E and B = Σ AZ, Then,
Lemma implies that
|BB-1/2AB-1/211 ol 3 with probability 1- 6. Thus, we have
(I+ -
op
18-"A²B-
B1/2-a
B 1/2 B-1122AB A² 1/2 B B 1/2 A2 1/2) B1/2-a
op
B1/2-0 8 1 (1
||B|| B 1-
op
op
with probability Moreover, since 2(1 - a) 2 1, we have
2(1-a)
B
(1-a)
|| Al 10p (((-aa
- | + }
op
op
$ 22 22111-))11 0 || + 2 -a))
Therefore, we conclude that
|/ ((x)) ))pppp)pp)) /((=)((x;)ddp(o)
L2(dp)
$ ( - (E=))ool (P(e)
|E 11)--($
a -)E-E)((++ +1
|| A1) (2
1/ |B||
3(1 a)
| + 21)
op
$ 2/2-aa
ΣΣ| 2(1-a) + \2(1-a) +||(E + xT)
S 22 |2 1
+ A )
+ 21)
op
S 22 max | |E||op>
E 1/2] + +
+|[(E + XI)" 23
This completes the proof.
B The Experimental Results on Next-Token Prediction
this section, we report the performance next-token prediction for the distilled models Section The results are presented in Table 7
Here, we highlight the following two key observations:
Among the three types of loss, the layerwise losses sometimes underperform compared to the "direct" loss. This is natural because the cross entropy loss for next-token prediction is used in "direct",, while L loss (Softmax loss) just aims make the attention kernel (attention weights) and its approximation close. As we reported in Section f for the downstream tasks, the models distilled with layerwise loss have comparable performance the models distilled with "direct" loss.
We observe that the models distilled using (layerwise) softmax loss outperform DiJiang, in which the learnable parameters are directly trained using next-token prediction loss. This result emphasizes the validity of our approach in training the features in linear attention.
C More Information of the Experiments
C.1 Computational Resources and Hyperparameters
We provide the details of the experimental settings follows:
24

---

Table 7: The performance of next-token prediction for the models distilled from GPT-2.
direct
softmax
L
Baseline method
Loss
Fix
3.9657
5.4082
6.2453
Diliang
5.5965
DoF
5.2651
4.0170
6,6802
Performer
8.1586
DoF + Clip
5.4547
4.0355
6.2378
Original GPT-2
3.3558
Table Summary of learning rates.
Model GPT-2
PiQA
logiQA
ARC-E
ARC-C
Winogrande
MMLU
WSC
le-4
5e-5
le-4
le-5
5e-4
le-6
le-6
Pythia-IB
le-5
5e-5
5e-5
5c-5
5e-5
5c-6
5c-5
Table 9: Summary of batch sizes.
Model
Distillation
PiQA
logiQA
ARC-E
ARC-C Winogrande
MMLU
WSC
GPT-2
128
128
64
64
64
128
64
128
Pythia-IB
64
64
32
32
32
64
32
64
Experiments for GPT-2 are conducted on four devices of A100 40GB. Experiments for Pythia- conducted on four devices of A100 80GB.
For Wikipedia dataset, we randomly sample 10% segment, and use it as one dataset.
The context lengths for GPT-2 and Pythia- IB were set to 1024 and 2048, respectively.
Training for distillation (learning features in proposed method learning parameters in DiJiang) is conducted over 1 epoch. This consumes about 0.5 day for GPT-2 and day for Pythia-
For the downstream tasks except for MMLU. training is conducted over three epochs for the tasks except for MMLU, and the best accuracy among three epochs is reported. For MMLU, we train the model for one epoch, and the accuracy for the last checkpoint reported.
As for the learning rates of distillation,
-when using DiJiang, we set O 02.
-when training feature maps, we set 0.02 for 21...., 2M and 0.2 for a1....,aaM.
The learning rates downstream task are chosen maximize the accuracy when we fine-tune the original model with softmax attention, and the same learning rates are used for the distilled model. The choices of learning rates are 1e-6, 5e-6, 1e-5, 5e-5, 1e-4, 5e-4 for GPT-2, and 1e-7, 5e-7, 1e-6, 5e-6, 1e-5, 5e-5 for Pythia- IB. The selected learning rates are summarized in Table
The batch sizes summarized in Table 9 gradient accumulation.
When out-of-memory enror occurs, utilized
C.2 Datasets
All the datasets used in this paper are publicly available from HuggingFace Datasets library. The dataset URLs and licenses are summarized in To
Table 10: Summary of datasets.
Datasel
URL
License
Wikipedia
https://hugginggtte http:///wwa.ccccotttttttttteeeee wwt.aataaee//leepaaeaeeeeo http://stt..llppo.c.aaassa//////
PIQA
hEtps: httpaa///huggg/n/// ht//aaaaseta/pgge
AFL-3.00
logiQA
http::///uuginn//ccc 00/
Apache Lioense, Version 2.0 CC BY-SA 4,0 CC BY.SA
ARC-E
ARC-C
kttps: http:://hhugssggaaccc.. co/ wtstaaatt/sllnnnaiaaaa
Winogrande
tttpat//bugguuggaccc
Apache License, Version 2.0
MMLU
https:/
MIT
C.3 Additional Results
In this section, provide the experimental results beyond Section 4
25

---

Table 11:
The approximation error when using L loss for GPT-2. All values are expressed in units
of x x10-
2
3
4
6
7
9
10
11
12
Total
Fix DoF
8.70
10.60
3.74
2.39
3.48
1.77
1.80
1.19
1.22
1.32
8.26
45.68
7.83
8.92
4.15
2.50
3.48
1.76
1.76
1.24
1.33
1.29
1.32
8.52
44.10
Table 12: Training time (in hours and minutes) required for feature training.
Model GPT-2
Fix / DoF
direct
L2
Softmax
Fix
10h 44m
5h 57m
8h 53m
DoF
10h 45m
6h 3m
8h 59m
DoF + Clip
10h 35m
5h 43m
8h 51m
Pythia-15 IB
Fix
23h 30m
DoF
24h 18m
Approximation error of K. Table I1 shows the loss values when GPT-2 is distilled into models whose dimensionality is determined by the Fix and DoF methods using the loss. Because the
1,2 loss represents the error between K and R, the values in this table directly indicate how well linear attention approximates softmax attention. From the table, we approximation errors under Fix also achieve small errors under the DoF. At the same time, for the layers with large approximation errors under the Fix method, the models whose dimensionality is selected by DoF exhibit smaller approximation errors. As a result, the total error across all layers is smaller for DoF than for Fix.
that the layers with small
Computational cost of Algorithm[] The total time required to run Algorithm T is 105.05 seconds for GPT-2 and 476.53 seconds for Pythia- IB. These costs are negligible compared to the overall distillation process, which includes feature map training.
Computational cost of feature training. We have summarized the actual training time required for feature training in Table 12 For GPT-2, the maximum training time is approximately 0.5 days, and for Pythia- IB, it is up to I day. Consistent with the trends shown in Table the computational cost follows the order: 2 < softmax < direct.
C.4 Additional Discussion
Why distilled models sometimes outperform the original model. Since the proposed method performs distillation by approximating the softmax attention with linear attention, the distilled models generally expected to perform worse than the original models. However, as shown Table 3 and Table 6) the distilled models outperform the original models in some tasks. Here, we discuss possible explanations for this phenomenon, in case readers are interested in understanding it.
One possible interpretation is that, the original model may not have reached the lowest possible test loss due to the limited amount of data available during downstream fine-tuning. Consequently, the error introduced by the approximate attention kernel in the distilled model might have incidentally contributed to reduction in test loss. We believe that such case occurred in experiments.
Another possible explanation that the original model may have slightly overfit the fine-tuning data for the downstream task, thereby limiting its test performance. In contrast, the distilled model might have benefited from an implicit regularization effect due to its smaller representational capacity, which could have led to improved generalization performance.
26


Reviews:
[{"review_id": "Uaz6Moa8G6", "reviewer": "Reviewer_RVhi", "summary": "In this work, the authors study an interesting practical problem in sequence modeling: Distilling Softmax Attention into Linear Attention. The authors conduct a series of theoretical analyses on the selection of optimal feature dimensions and propose a layer-wise training method. Empirical verifications of the proposed algorithms are provided.", "questions": "1. Comparing fix to DoF, if we fix the feature dimension in every layer and head to be the max of all optimal dimensions, will the fix approach be better than the DoF approach? If not, could the authors explain more about it?\n2. The authors could add discussions on [1, 2], which also provide analyses on the approximation quality of PRF-based Linear Attention and Softmax Attention in different perspectives, e.g., Theorem 3,4 in [1] and Theorem 3 in [2].\n\n[1] Rethinking Attention with Performers. ICLR 2021.\n\n[2] Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding. NeurIPS 2021.", "limitations": "The authors carefully discuss limitations and future work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- The studied problem is interesting and significant, which helps to improve the utility of linear attention in practice.\n- The proposed algorithms are well-motivated and supported by a series of interesting theoretical analyses.\n- The paper is well-structured and easy to follow. The authors explain well after each theoretical statement to help readers catch up quickly.\n\n## Weaknesses\n- **Lack of empirical investigations on the approximation error when using the proposed DOF (and Clip)**. It would be better for the readers to understand the effectiveness of the proposed approach to further investigate the approximation error between Linear Attention and Softmax Attention in practice. Although the authors provide an analysis of it in Theorem 3, the readers would be curious about the empirical comparisons, especially its relations to the number of feature dimensions.\n- **Limited scale of the experiments**. We can see the effectiveness of the proposed selection methods on GPT-2/Pythia-1B, but both are small-scale experiments. It is understandable that the computational resources would be a restriction for larger-scale experiments, but it would be better if the effectiveness of the approach could be verified further.\n- **Limited choices of the kernel feature map**. In this work, the authors only investigate the PRF empirically. What is the effect of the proposed approach on other types of kernel feature maps? Will it be consistent adaptable or not? Further investigations could enhance the quality of this submission.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback. We address the specific concerns and questions below.\n\n> **Lack of empirical investigations on the approximation error when using the proposed DOF (and Clip)**. It would be better for the readers to understand the effectiveness of the proposed approach to further investigate the approximation error between Linear Attention and Softmax Attention in practice. Although the authors provide an analysis of it in Theorem 3, the readers would be curious about the empirical comparisons, especially its relations to the number of feature dimensions.\n\nThank you for the insightful suggestion. In general, increasing the number of feature dimensions reduces the approximation error between linear and softmax attention, while decreasing it increases the error. Although our dimension selection strategy is theoretically justified, it is challenging to directly evaluate the effectiveness of the selected dimensions purely from an approximation error standpoint.\n\nIn this work, following common practice in language model distillation, we chose to evaluate our method based on downstream task performance, which reflects the overall effectiveness in practical scenarios. The improved results under our DoF-based selection empirically validate that the approximation quality is improved when dimensions are allocated adaptively.\n\nAs a reference, we provide below an empirical analysis from our GPT-2 experiments, showing the selected feature dimensions per layer and the change in approximation error (DoF/Fix ratio) before and after applying DoF-based selection. Here, the Fix baseline used a uniform dimension of 64. In layers where DoF selected larger dimensions than Fix, the approximation error decreased. On the other hand, in layers where smaller dimensions were selected, the error increased. This behavior is consistent with theoretical expectations. More importantly, the overall downstream performance improved, confirming that DoF-based allocation offers a better tradeoff between capacity and approximation quality.\n\n|      | 1    | 2    | 3   | 4   | 5   | 6   | 7   | 8   | 9   | 10  | 11  | 12  |\n|------|------|------|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n| Dimension | **130** | **182** | 35  | 44  | 42  | **65** | **92** | 33  | 28  | 34  | 46  | 39  |\n| DoF/Fix | 0.900 | 0.838 | 1.111 | 1.047 | 1.000 | 0.995 | 0.981 | 1.023 | 1.114 | 1.051 | 1.004 | 1.031 |\n\n> **Limited scale of the experiments**. We can see the effectiveness of the proposed selection methods on GPT-2/Pythia-1B, but both are small-scale experiments. It is understandable that the computational resources would be a restriction for larger-scale experiments, but it would be better if the effectiveness of the approach could be verified further.\n\nWe thank the reviewer for raising this important point. While we acknowledge the benefits of evaluating on larger models, we believe the current experiments already provide strong evidence for the effectiveness and generality of our method.\n\nOur primary aim is to establish a theoretically grounded and practically effective framework for distilling softmax attention into linear attention with adaptive feature dimensions. To validate this, we deliberately chose two widely-used and open-source models that differ in scale. Importantly, we evaluated performance on a diverse set of downstream tasks, showing consistent improvements over strong baselines.\n\nWe fully agree that extending to larger models such as LLaMA is valuable, and view it as promising directions for future work.  Nonetheless, we believe our current empirical results, supported by rigorous theory, already demonstrate the core value and practicality of our method.\n\n> **Limited choices of the kernel feature map**. In this work, the authors only investigate the PRF empirically. What is the effect of the proposed approach on other types of kernel feature maps? Will it be consistent adaptable or not? Further investigations could enhance the quality of this submission.\n\nOur theoretical framework is designed to apply to feature maps $\\phi$ that satisfy the condition $K(x, y) = \\mathbb{E}_{z \\sim \\tau}[\\phi(x; z)\\phi(y; z)]$, i.e., where the attention kernel can be expressed as an expectation over random features. Consequently, our proposed approach is not applicable to methods such as those in Katharopoulos et al. (2020) and Qin et al. (2022), which do not involve in randomized feature expectations.\n\nAmong approximations of attention kernel that satisfy the above condition, PRF are the most well-established and promising, which motivated our focus on them in this work. Another common alternative is the use of random fourier features as in Peng et al. (2021). However, as pointed out in Choromanski et al. (2021), this can yield negative values, which may adversely affect stability when training. Since such implementation aspects beyond feature dimension can significantly influence performance, we chose not to adopt RFF in our study. This is briefly discussed in lines 127--128 of our paper.\n\nThat said, our method is general, and if a promising alternative feature map for approximating an attention kernel $K$ (satisfying the expectation condition above) is developed, our dimensionality selection approach can readily be extended to it. We believe that including this discussion in the camera-ready version will make the paper more convincing. We appreciate the helpful comment.\n\n\n> 1. Comparing fix to DoF, if we fix the feature dimension in every layer and head to be the max of all optimal dimensions, will the fix approach be better than the DoF approach? If not, could the authors explain more about it?\n\nThank you for the interesting question. We conducted additional experiments on GPT-2, where we fixed the feature dimension of all layers and heads to the global maximum value observed in Table 1 (i.e., 182). The results are as follows:\n\nUsing the softmax loss, the average downstream accuracy across all tasks was 0.4149, with individual task scores of: 0.5724 (PiQA), 0.2842 (logiQA), 0.3232 (ARC-E), 0.2875 (ARC-C), 0.5004 (Winogrande), 0.3007 (MMLU), and 0.6346 (WSC). These are very close to the performance achieved by our DoF-based method. One possible reason is that the approximation of the layer with the largest DoF became a bottleneck, and increasing the feature dimensions of other layers did not improve the overall approximation quality of the network.\n\nUsing the direct loss (i.e., next-token prediction loss), the average performance dropped to 0.4149, with task-wise scores: 0.5626, 0.2826, 0.2997, 0.2875, 0.5020, 0.2735, and 0.6442. This is lower than the performance achieved by the “DoF” in Table 3. We hypothesize that this drop may be due to the more complex optimization landscape of the direct loss, which could make it harder and unstable to learn the feature maps. For reference, when we fixed all feature dimensions to 128 and used the direct loss, the average accuracy was 0.4190, which is higher than the 182-fixed version, and close to the DoF result.\n\n> 2. The authors could add discussions on [1, 2], which also provide analyses on the approximation quality of PRF-based Linear Attention and Softmax Attention in different perspectives, e.g., Theorem 3,4 in [1] and Theorem 3 in [2].\n\nThank you for the valuable comment to connect our theory to prior results. We summarize the relationship to [1, 2] as follows. We will add these discussions to the final version of the paper, and include [2] as an additional citation, which was previously omitted.\n\nTheorem 3 and 4 in [1] analyze the approximation quality of PRF-based linear attention in terms of the **sup-norm error** between the true attention kernel and its approximation. They show that to achieve an error less than $\\epsilon$, the required number of features $M$ depends on the head dimension $d$, the norm bound $R$ of queries and keys, and $\\epsilon$, with a particular dependence of $M = O(d \\log d)$ on $d$. In contrast, our Theorem 3 focuses on bounding the **$L^2(\\rho \\otimes \\rho)$-norm**, where $\\rho$ represents the data distribution over queries and keys. This allows us to capture the structure of the input distribution in each attention layer. As a result, the approximation error is governed not by the dimension $d$, but by the degrees of freedom (DoF), which reflect the intrinsic dimensionality of the data.\n\nTheorem 3 in [2] highlights an exponential dependence on the norm bound $R$ of queries and keys. In our analysis, such dependence is implicitly captured by constants such as $C_K$ and $||K||^2_{L^2(\\rho \\otimes \\rho)}$, which may also grow exponentially with $R$. Our work does not aim to reduce the exponential dependence on $R$; instead, we focus on how to select the feature dimension per layer to minimize the approximation error under a fixed computational budget.\n\nIt is also worth noting that [2] proposes mitigating the exponential dependence on $R$ by normalizing queries and keys and incorporating relative positional encodings. While this is a promising direction, our setting assumes distillation from pre-trained models, where such modifications to the architecture or inputs are not directly applicable. That said, developing distillation techniques that incorporate such normalization to reduce sensitivity to $R$ remains an important and interesting direction for future work. We will mention this in the revised paper.\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period."}, {"review_id": "37OHdV3gOn", "reviewer": "Reviewer_uHta", "summary": "This paper addresses the challenge of selecting the feature dimension when distilling pretrained softmax attention into more efficient linear attention models. The authors propose a method to determine the feature dimension for each layer based on the estimated degrees of freedom of the attention matrix. In addition, they advocate for layer-wise distillation rather than end-to-end training, arguing that it leads to better alignment and stability during knowledge transfer. The proposed approach is evaluated using GPT-2 and Pythia-1B models on the Wikipedia dataset, demonstrating promising results in terms of both efficiency and performance.", "questions": "1. The performance of baseline models (e.g., original, DiJiang, and Performer) in Tables 3 and 6 appears significantly lower than what has been publicly reported in previous works. Could the authors clarify why this is the case? Were different training setups, datasets, or hyperparameters used?\n2. When comparing the proposed method with DiJiang and Performer in Table 3, efficiency metrics such as FLOPs or latency should be reported as well. Since the motivation is to distill softmax attention into linear attention while preserving efficiency, this comparison is incomplete without those figures.\n3. In Table 3, it is unclear whether Fix should be considered a baseline or part of the proposed method. The central claim of the paper is that the DoF-based method can achieve better performance under the same feature dimension budget. However, in cases where Fix and DoF share the same budget, their performance is often quite similar. This undermines the claim that the DoF-based allocation provides a clear advantage. A more careful comparison and explanation are needed to clarify whether the observed gains are due to the DoF strategy itself or simply due to increased parameter capacity.\n4. What does “DoF + clip” mean in Table 3? The phrase “clamp the feature dimensions to be at most the head size” is unclear—does it imply that feature dimensions can normally exceed the head size? If so, how is this handled in implementation, and why would such a configuration be valid?\n5. Why are DiJiang, Performer, and “DoF + clip” missing from Table 6? Including these baselines would provide a more complete picture of the proposed method’s performance across settings.\n6. The motivation for layerwise training is not clearly supported. In Table 3, the Direct method often achieves the best performance, raising the question of whether layerwise distillation offers any real advantage. Does layerwise training lead to faster convergence, better stability, or other benefits not captured by final performance? If not, its inclusion appears secondary and may not contribute meaningfully to the overall method.\n\nMinor Suggestion:\n- It is unclear why all attention heads are required to have the same feature dimension. Allowing different dimensions per head could potentially lead to better parameter efficiency and adaptivity. Have the authors considered or experimented with head-wise dimension selection?", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper is clearly written and easy to follow with well-structured explanations.\n- The idea of selecting the optimal feature dimension for each layer based on degrees of freedom is novel and supported by empirical evidence. The observation that different layers and attention heads exhibit significantly varying degrees of freedom suggests substantial redundancy in standard ViT-style architectures, making the analysis insightful and impactful.\n\nWeaknesses:\n- The theoretical derivation leading to the first and only formal result is overly long and contributes little to the main findings. Much of the content is a restatement or adaptation of existing results from prior work, rather than offering new theoretical insights directly tied to the proposed method.\n- The proposal to use layer-wise training instead of end-to-end distillation lacks clear motivation and theoretical justification. It appears secondary to the main contribution and is not sufficiently developed.\n- The experimental results, while promising, do not strongly support the claimed advantages of the proposed method. The gains are modest, and more thorough ablations or broader comparisons would be needed to convincingly demonstrate the effectiveness of the approach.", "was_revised": true, "final_justification": "This paper presents a good motivation. However, the proposed method appears overly complicated and lacks clarity, making it difficult to follow. Additionally, the results are not clearly presented or sufficiently convincing to support the claims. I believe the paper needs further polishing and clarification before it can be considered for acceptance.", "rebuttal": "We thank the reviewer for the helpful feedback. We address the specific concerns and questions below.\n\n> The theoretical derivation leading to the first and only formal result is overly long...\n\nWe appreciate your important comment. We agree that parts of the theoretical section can be made more concise. However, we believe the current exposition plays an important role for two reasons.\n\nFirst, it introduces the key notation and concepts, such as the kernel approximation $\\hat{K}$, the feature maps $\\phi$, $\\Phi$, and the operators $\\Sigma$, $\\hat{\\Sigma}$, and $N_{q,\\lambda}$, which are fundamental for our analysis and algorithm.\n\nSecond, the derivation builds on prior work, and we include these results to **motivate and justify the use of degrees of freedom (DoF) as a core quantity for feature dimension selection**. This connection, though known in kernel theory, is not established in the context of linear attention, and we believe making it explicit clarifies the theoretical motivation behind our method.\n\nThat said, we will revise and streamline parts of the exposition (e.g., lines 140–144 and around Proposition 1) to improve clarity in the final version. We appreciate your suggestion.\n\n> The proposal to use layer-wise training ...\n\n> 6. The motivation for layerwise training is not clearly supported. \n\n**The main motivation for introducing layer-wise training is to reduce the computational cost of distillation**, which is discussed in lines 52–55, 64–65, and 220–223. End-to-end training requires full backpropagation through all layers simultaneously, which leads to significantly higher training cost, especially for larger models. In contrast, layerwise training allows us to optimize each attention layer independently, avoiding high cost for computing gradient. Our experiments demonstrate that layerwise training is **not only more efficient** (Table 5) **but also yields performance comparable to end-to-end training** (Table 3, 6).\n\n> The experimental results, while promising, do not strongly support ...\n\nThank you for your feedback. We would like to clarify that the experimental results do provide concrete evidence supporting the effectiveness of our proposed method.\n\nOur experimental results show that using our DoF-based feature dimension allocation leads to better downstream task performance compared to the “Fix”. Importantly, **Fix and DoF are matched in total feature dimension, and therefore incur the same computational cost**. This point may not have been sufficiently described, and **we will clarify it in Section 4.2 of the final version.** As further evidence, Table 4 confirms that the inference time for Fix and DoF is nearly identical. Thus, the observed performance improvement cannot be attributed to an increase in parameter capacity or computation budget, but rather to a more effective allocation of feature dimensions across layers.\n\nCrucially, in the context of distillation from a stronger teacher model, the key goal in distillation is minimizing performance drop from the teacher model. Therefore, when evaluating the benefits of our method over baseline “Fix”, the key metric is not the raw accuracy gain, but rather **how small the performance degradation from the original softmax-based model is**. From this perspective, the improvements are far from modest. For example, as shown in Table 6 (Pythia-1B), the Fix baseline exhibits a 0.0265 drop in average accuracy compared to the original model, whereas our DoF-based method reduces this gap to 0.0161, which implies **a 39% reduction in degradation**. We view this as a significant improvement given that the total inference cost remains unchanged.\n\nWe agree that further ablations could strengthen the case, but we believe the current results already demonstrate the practical value of our approach.\n\n> 1. The performance of baseline models in Tables 3 and 6 ...\n\nWe believe that the differences in baseline performance are mainly due to differences in experimental setup. In order to ensure a fair and controlled comparison, we fine-tuned all models under a unified setup. These settings may differ from those used in prior works, which often employ task-specific tuning. While such tuning can improve absolute performance, we did not apply it, as our goal is to demonstrate consistent relative improvements under equal resource constraints. Full details of our experimental setup are provided in Appendix C, and we also include our code in the supplementary materials to support reproducibility. The baseline implementations are taken from their official repositories.\n\nAlthough it is true that the downstream results for baseline models are sometimes lower than those reported in prior work our reproduced baselines are not consistently weaker: for example, when comparing results on DiJiang distilled by GPT-2 (117M) of our paper and results on DiJiang-160M by Chen et al. (2024), in some of the tasks, the former shows higher accuracy (e.g., logiQA, ARC-C, WSC). \n\n> 2. When comparing the proposed method with DiJiang and Performer in Table 3, ...\n\nThank you for the insightful suggestion. Our goal is to improve downstream performance of linear attention under a fixed computational budget, rather than to further accelerate it. Accordingly, Table 3 compares accuracy at equal total feature dimensions, and Table 4 confirms that our method (DoF) maintains inference efficiency comparable to the fixed baseline (Fix).\n\nWe agree that including efficiency metrics for methods like DiJiang would be valuable. However, DiJiang incorporates an acceleration technique using discrete cosine transforms which make it faster than standard linear attention. To emphasize the effect of adaptive dimension selection, we intentionally focused on a simple linear attention architecture without such modifications when comparing inference-time speed.\n\nIntegrating our dimension selection method into more efficient architectures like DiJiang is a promising direction for future work.  We also emphasize that our DoF-based dimension selection is not in competition with specific linear attention mechanisms like DiJiang or Performer. Rather, it is complementary, and can be applied to various linear attention models, including DiJiang, as a general strategy for improving performance under a fixed budget.\n\nFinally, we note that Performer and Fix share the same architecture and compute cost; the difference lies in whether kernel features are learned or fixed.\n\n> 3. In Table 3, it is unclear whether Fix ...\n\nThank you for this important question. We clarify that “Fix” serves as a baseline to evaluate the effect of DoF-based dimension allocation. Both Fix and DoF use trained feature maps under the same total feature dimension. In Fix, the feature dimension is uniformly set to the head size (as in prior work), while DoF allocates it per layer based on Table 1. Since the total dimensions of the both models are the same (see also line 246--248), the computational complexity of two models are the same. Thus, any performance difference reflects the impact of the allocation strategy itself, not differences in model capacity. \n\nFurthermore, we would like to emphasize again that DoF reduces the accuracy drop from the original model by 39% compared to Fix, which shows a substantial gain under matched budgets.\n\nWe acknowledge that it was not sufficiently described that the feature dimension of “Fix” is set to be the same as the head size. We will clarify this point in the final version. Thank you again for the helpful question.\n\n> 4. What does “DoF + clip” ...\n\nDoF + clip is introduced in line 280--282. It refers to a variant of our method in which the feature dimensions computed via Algorithm 1 are clamped to be at most the head size. That is, for each layer, we set the feature dimension to $\\min(\\text{DoF-based value}, \\text{head size})$, where DoF-based value corresponds to the values shown in Table 1.\nIt is indeed possible in our method that the computed DoF-based feature dimension exceeds the head size (e.g., GPT-2 Layer 2 in Table 1 has a value of 182, while the head size is 64). This does not cause any implementation issue, as the feature dimension in linear attention can operate with arbitrary feature dimensions, regardless of the head size.\n\n> 5. Why are DiJiang, Performer, and “DoF + clip” missing from Table 6? \n\nThank you for the suggestion. We excluded DiJiang, Performer, and DoF + Clip from the Pythia-1B results (Table 6) to focus on the key comparison between Fix and DoF, which directly evaluates our core contribution (i.e., feature dimension selection via Algorithm 1), under a larger model.\n\nThese baselines were thoroughly evaluated in the GPT-2 setting (Table 3), where we showed that DoF outperforms both DiJiang and Performer. These results showed that using DoF for layer-wise dimension allocation and feature training yields improved downstream accuracy under the same feature dimensions. We chose not to repeat these comparisons in Table 6. We also note that DoF + Clip is an auxiliary variant for reducing inference cost and is not central to our main claim.\n\n> Minor Suggestion: It is unclear why all attention heads are ...\n\nWe set the feature dimension to be the same across heads within a layer, as all heads are handled using a shared tensor (see line 205). This avoids overhead from irregular memory layouts or dynamic shapes. \n\nWhile head-wise allocation may offer further gains in adaptivity, it requires nontrivial architectural changes and scheduling complexity. We view this as an important direction for future work, but our current focus is on improving efficiency and accuracy within standard Transformer infrastructures.\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period. We would greatly appreciate it if you could consider increasing the score once all concerns have been resolved."}, {"review_id": "DL52aEHnHf", "reviewer": "Reviewer_i3tm", "summary": "The paper proposes a new distillation method to distill softmax attention models into linear attention models. The novelty lies in the dimension selection of the features used for distillation, which is computed via the effective degrees of freedom each layer (and head) needs.The paper also proposes two methods to learn the features from data and shows extensive empirical evidence to backup their claims.", "questions": "1. How computationally involved is the computation of the effective degrees of freedom (i.e., Algorithm 1)? You only report inference speed.\n1. Are you setting the feature dimension of the heads to the largest value computed? This is not entirely clear in the paper, you only mention that you fix the dimension.\n1. If yes (to the previous question), have you considered using the mean dimension across all heads in a layer (or any other suitable way)? This will reduce the performance of the head with the highest feature dimension but this might be mitigated by the other heads.\n1. Do you have a hypotheses why the distilled models can outperform the original model? Is this mainly due to the learned features?", "limitations": "The limitations are adequately addressed in a separate paragraph in the conlcusion.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n- The proposed idea is novel and offers a new efficient way to distill softmax attention models into more efficient linear ones.\n- The practical impact of the proposed method is potentially high and is bolstered by two new feature learning methods.\n- The paper is well-written and easy to follow.\n\n**Weaknesses:**\n- The computational advantages are only sparingly shown in the empirical section. The computation times/performance of Algorithm 1 is not show nor discussed. Including this would give further information into how expensive the distillation process is. While there are some results for the achieved inference time after distillation, it would be beneficial to also show the inference time of the original softmax model to compare directly. Finally, more information on the computational complexity of the feature learning methods would be appreciated (apart from Table 5).\n- The paper could benefit from a better setting in the current literature, i.e., the paper has relatively few citations. E.g. in Line 32, the references Wang et al. (2024) and Han et al. (2024) are not good references for the connection of linear attention to SSMs, better ones would be [1] and [2]. Additionally, mentioning more distillation methods would be nice, e.g., [3], [4].\n\n\n**References:**\n\n[1] Sieber et al., \"Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks\", NeurIPS 2024\n\n[2] Dao & Gu, \"Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality\", ICML 2024\n\n[3] Yang et al., \"Knowledge distillation via softmax regression representation learning\", ICLR 2021\n\n[4] Wang et al., \"Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers\", NeurIPS 2020\n\n**Minor comments/Typos:**\n- Please add the appendix to the main text for the final submission to improve the flow.\n- Caption of Table 1: The underlined values are the three lowest values, not the dimensions lower than the top three.\n- Add to captions of Table 4 and 5 that these values are for GPT-2.\n- Line 11: ... smaller errors ...\n- Line 34/35: Despite -> Even though; remove \"indeed\"", "was_revised": true, "final_justification": "The authors adequately addressed my concerns in the rebuttal, therefore I will maintain my positive score. I think the motivation of the paper is good and in its revised form also suitable for presentation at NeurIPS.", "rebuttal": "We thank the reviewer for the helpful feedback. We address the specific concerns and questions below.\n\n\n> The computational advantages are only sparingly shown in the empirical section. The computation times/performance of Algorithm 1 is not show nor discussed. Including this would give further information into how expensive the distillation process is.\n\n\n> 1. How computationally involved is the computation of the effective degrees of freedom (i.e., Algorithm 1)? You only report inference speed.\n\n\nThank you for your thoughtful suggestion. We have measured the runtime of Algorithm 1. The total time required is 105.05 seconds for GPT-2 and 476.53 seconds for Pythia-1B. These costs are negligible compared to the overall distillation process, which includes feature map training. We will include this information in the final version of the paper.\n\n> While there are some results for the achieved inference time after distillation, it would be beneficial to also show the inference time of the original softmax model to compare directly.\n\nWe thank the reviewer for the valuable suggestion. We agree that comparing the inference time of the distilled linear attention model with the original softmax-based model could be informative. However, we did not include such a comparison in the paper for the following reasons:\n\n1. **Implementation-dependence**: The inference time is highly implementation-specific. Modern Transformer libraries apply a wide range of kernel-level and architectural optimizations to softmax attention (e.g., FlashAttention), whereas our linear attention implementation does not yet incorporate such optimizations. As a result, a naive comparison of wall-clock time would reflect engineering disparities rather than algorithmic efficiency.\n2. **Well-established fact**: The linear-time nature of linear attention compared to the quadratic complexity of softmax attention has been well documented in prior work. Our focus is not on re-establishing this fact, but rather on *how to optimally allocate feature dimensions* within a given linear attention budget to improve performance.\n3. **DoF maintains the efficiency of Fix**: A central goal of our work is to enhance the performance of linear attention *without increasing its computational cost*. As shown in Table 4, our method adaptively allocates feature dimensions across layers while keeping the overall inference cost effectively unchanged. This confirms that the proposed DoF-based strategy improves accuracy while preserving the efficiency advantage of linear attention.\n\n> Finally, more information on the computational complexity of the feature learning methods would be appreciated (apart from Table 5).\n\nWe have summarized the actual training time (in seconds) required for feature learning in the table below. We will include this table in the final version of the paper. For GPT-2, the maximum training time is approximately 0.5 days, and for Pythia-1B, it is up to 1 day. Consistent with the trends shown in Table 5, the computational cost follows the order: L2 < softmax < direct.\n\n| Model | Fix / DoF | direct | L2 | Softmax |\n| --- | --- | --- | --- | --- |\n| GPT-2 | Fix | 38659.7177 | 21420.2063 | 31991.2506 |\n|  | DoF | 38689.3816 | 21753.6426 | 32322.82 |\n|  | DoF + Clip | 38093.8039 | 20563.6596 | 31847.8818 |\n| Pythia-1B | Fix | - | - | 84601.079 |\n|  | DoF | - | - | 87506.1713 |\n\n> The paper could benefit from a better setting in the current literature, i.e., the paper has relatively few citations. E.g. in Line 32, the references Wang et al. (2024) and Han et al. (2024) are not good references for the connection of linear attention to SSMs, better ones would be [1] and [2]. Additionally, mentioning more distillation methods would be nice, e.g., [3], [4].\n\nThank you for the important comment regarding related work. We agree that the current discussion of related work can be improved. We will revise Section 1 to include the more relevant references (Sieber et al., 2024; Dao & Gu, 2024) about the connection of linear attention to SSMs, as well as additional distillation methods (Yang et al., 2021; Wang et al., 2020). \n\n> **Minor comments/Typos:** …\n\nThank you for pointing out these typos and for your helpful suggestions to improve the readability of the paper. We will incorporate all of these changes (including corrections to the captions, wording suggestions, and appendix integration) in the final camera-ready version.\n\n> Are you setting the feature dimension of the heads to the largest value computed? This is not entirely clear in the paper, you only mention that you fix the dimension.\n\nThank you for pointing this out. Yes, we set the feature dimension for each layer to the maximum value among its heads. This is indeed mentioned in Algorithm 1 and also in line 207 of the paper, where we state: “Hence, we define the degrees of freedom …”. However, we understand that this point may not have been sufficiently emphasized. In the revised version, we will restate this point more clearly in experimental section (Section 4).\n\n> 3. If yes (to the previous question), have you considered using the mean dimension across all heads in a layer (or any other suitable way)? This will reduce the performance of the head with the highest feature dimension but this might be mitigated by the other heads.\n\nWe are grateful for this suggestion. We have not explored using the mean feature dimension across all heads in a layer. Since our distillation procedure is conducted independently for each head, i.e., the loss functions defined in Section 3.3 are applied separately to each head, we believe it is unlikely that other heads can compensate for the reduced performance of a head whose feature dimension is insufficient. Therefore, lowering the feature dimension of a complex head would likely degrade its ability to approximate the attention kernel, which may in turn negatively affect the overall model performance.\n\n\nThat said, we agree that this is an interesting direction, and exploring more flexible or shared allocation strategies across heads could be valuable for future work.\n\n> 4. Do you have a hypotheses why the distilled models can outperform the original model? Is this mainly due to the learned features?\n\nOne explanation is that, in practice, the original model may not have reached the lowest possible test loss due to the limited amount of data available during downstream fine-tuning. As a result, the deviation caused by the approximated attention kernel in the distilled model may have incidentally reduced the test loss. We believe such a case occurred in our experiments.\n\nAnother possible (though unverified) hypothesis is that the original model may have slightly overfit the fine-tuning training data, which limited its test performance. In contrast, the distilled models may have benefited from an implicit regularization effect since they have smaller representational capacity, thereby improving generalization.\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period."}, {"review_id": "0k78p83ewU", "reviewer": "Reviewer_nGDS", "summary": "The submission proposes a data-adaptive and layer-specific way to distill softmax attention in pre‑trained Transformers into linear attention via Kernel Approximation.\n\nThe core idea is to automatically choose the random feature dimension used in linear attention based on the statistical degrees of freedom (DoF) of the attention kernel. The authors derive a bound showing that the kernel approximation error depends on $N_{q, \\lambda}$, a function of the sampling density, and show that sampling from a density proportional to  $\\text{tr}(\\Sigma(\\Sigma+\\lambda I)^{-1})$ (the DoF) yields an optimal sample complexity. \n\nThey then propose Algorithm 1 to estimate per‑layer DoF using a Gram matrix computed on sampled queries/keys from a pre‑trained model and allocate feature dimensions across layers under a budget constraint. To avoid expensive end‑to‑end distillation, the paper introduces layer‑wise training of random features via either an $L_2$ loss that matches attention kernels or a softmax (cross‑entropy) loss that matches attention distributions. Experiments distill GPT‑2 and Pythia‑1B models and evaluate on seven downstream tasks. Models using DoF‑based dimension selection and layer‑wise training match or exceed baselines with fixed dimensions (e.g., Performer and DiJiang) and sometimes approach the original model’s accuracy. The paper concludes with an analysis of effective dimensionality across layers and discusses limitations such as shared dimensions across heads.", "questions": "– To really unleash the benefits of linear attention, consider evaluating linear attention approaches on long-context tasks beyond standard benchmarks.", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**:\n\n- Clear motivation and problem formulation. The paper articulates the challenge of choosing a fixed random feature dimension for linear attention and motivates why different layers may require different dimensions. This is supported by theoretical analysis relating sample complexity to DoF.\n\n- Principled algorithm with theoretical guarantees. The authors derive an error bound (Theorem 3) for approximating the softmax kernel by random features and show that sampling from the optimal density reduces the required number of features. Algorithm 1 provides a practical estimation of DoF and allocates per‑layer feature dimensions under a fixed cost.\n\n**Weaknesses**:\n\n- Limited evaluation scope. The experiments cover GPT‑2 (117M parameters) and Pythia‑1B (1B parameters), but there is no evaluation on larger models (e.g., GPT‑2 XL or LLaMA) or non‑language modalities. It’s unclear how the method scales to models with many more layers or to vision Transformers. The authors mention this limitation but do not provide preliminary evidence. It would be great if the authors could try their proposed approach on slightly larger models (such as LLaMA-3, Qwen, deepseek) and evaluate the models on a slightly broader tasks such as instruction following.\n\n- Compared to the baseline (Fix), in Table 2 and Table 3, the proposed method is not significantly better than the simple baseline.", "was_revised": true, "final_justification": "I would like to thank the authors for their rebuttal, and I will keep my original score.", "rebuttal": "We thank the reviewer for the insightful feedback. We address the specific concerns and questions below.\n\n\n> Limited evaluation scope. The experiments cover GPT‑2 (117M parameters) and Pythia‑1B (1B parameters), but there is no evaluation on larger models (e.g., GPT‑2 XL or LLaMA) or non‑language modalities. It’s unclear how the method scales to models with many more layers or to vision Transformers. The authors mention this limitation but do not provide preliminary evidence. It would be great if the authors could try their proposed approach on slightly larger models (such as LLaMA-3, Qwen, deepseek) and evaluate the models on a slightly broader tasks such as instruction following.\n\nWe appreciate the reviewer’s insightful comment. While we acknowledge the benefits of evaluating on other modalities, larger models (e.g., LLaMA, Qwen, DeepSeek) and broader tasks (e.g., instruction following), we believe that the current experiments already provide strong evidence for the effectiveness and generality of our method.\n\nOur primary aim is to establish a theoretically grounded and practically effective framework for distilling softmax attention into linear attention with adaptive feature dimensions. To validate this, we chose two widely-used and open-source models that differ in scale. The results across both models consistently show that our dimension selection method improves performance. Importantly, we evaluated performance on a diverse set of downstream tasks, showing consistent improvements over strong baselines.\n\nWe fully agree that extending our evaluation to larger models and broader tasks is a natural and valuable next step. Nonetheless, we believe our current empirical results, supported by rigorous theory, already demonstrate the core value and practicality of our method.\n\n\n> Compared to the baseline (Fix), in Table 2 and Table 3, the proposed method is not significantly better than the simple baseline.\n\nWe thank the reviewer for raising this important point. We believe the current experimental results offer sufficient evidence for the effectiveness of our method over the Fix baseline, for the following reasons:\n\n\nFirst, our evaluation covers **multiple diverse downstream tasks**, and we report both task-wise scores and their average. This provides a comprehensive assessment of model quality. The proposed method achieves better average performance overall, which demonstrates its effectiveness under fixed resource constraints. \n\n\nSecond, and most importantly, we believe the significance of our results lies not just in absolute accuracy improvements, but in **how much they help close the performance gap between linear attention and softmax attention**. In the context of distillation from a stronger teacher model, it is natural for a distilled model to underperform the original teacher model (while some exceptions exist in the experiments). Thus, the central goal is not to surpass the teacher, but to preserve as much of its performance as possible under tighter resource constraints. From this perspective, for instance, in Table 6 (Pythia-1B), the Fix baseline shows an average accuracy gap of 0.0265 compared to the original model. In contrast, our proposed DoF-based method reduces this gap to 0.0161, which implies **a 39% reduction in performance degradation**. We consider this a meaningful result, demonstrating that our method helps recover the performance typically lost when replacing softmax attention with linear attention.\n\n\n> To really unleash the benefits of linear attention, consider evaluating linear attention approaches on long-context tasks beyond standard benchmarks.\n\nThank you for your thoughtful suggestion. We agree that evaluating linear attention models on long-context tasks is an important direction. However, our work focuses on distilling pre-trained Transformers, which are typically trained on inputs with a fixed context length. As a result, adapting the distilled model to operate on longer contexts than the original model is inherently challenging, and this difficulty is not specific to our proposed feature dimension selection method. Indeed, prior work such as Chen et al. (2024), which also distills Transformers into linear attention models, conducts evaluation only on standard language benchmarks.\n\nMore importantly, the goal of our study is not to extend the usable context length of pre-trained Transformers through distillation. Rather, our aim is to construct more computationally efficient models that operate within the same context length, while minimizing performance degradation. We believe this goal is well supported by our current experimental results.\n\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period. We would greatly appreciate it if you could consider increasing the score once all concerns have been resolved."}]

Author rebuttal:

Meta-review: This paper proposes a data/layer-specific method to distill softmax attention into linear attention via Kernel Approximation, where the key is to select the optimal feature dimensions based on degree of freedom via a series of theoretical analyses. Empirical verifications of the proposed algorithms are provided.

The reviewers in general agree that the work proposes an interesting and solid method to distill softmax attention models into more efficient linear ones. In the discussions, the authors further addressed some concerns from reviewers, eg, more detailed experimental results. Please incorporate these improvements into the final version if accepted.
Decision: Accept (poster)