Pith. sign in

REVIEW 4 major objections 5 minor 13 references

Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency

T0 review · 4 major / 5 minor · reviewed 2026-08-06 · deepseek-v4-flash

Pith's one-line read A linear-attention layer needs only as many random features as its attention kernel has degrees of freedom.

desk verdict A clean DoF-based dimension allocation idea with a real gap: the theory covers q-lambda-sampled random features, while the actual method trains features with gradient descent, so the optimality claim is not established for what is implemented. read the letter →

arxiv 2507.03340 v1 pith:RJ5532KH submitted 2025-07-04 cs.LG stat.ML

classification cs.LGstat.ML
keywords linearattentionsoftmaxdistillationdegreesoffreedomkernelapproximationrandomfeaturesfeaturedimensionselectionlayerwisetrainingtransformerefficiency
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper argues that the feature dimension of each linear-attention layer should be set by the statistical degrees of freedom of that layer's attention kernel, $N^*_\lambda=\operatorname{tr}(\Sigma(\Sigma+\lambda I)^{-1})$, where $\Sigma$ is the kernel integral operator over the layer's query/key distribution and $\lambda$ is the allowed error. This would settle a question that previous distillation methods left to guesswork: instead of fixing one feature count for all layers, spend features where the kernel is genuinely complex. If the claim is right, distilled linear-attention models get better accuracy for the same compute, and the per-layer counts themselves become a diagnostic of how attention complexity evolves through the network. The paper proves an error bound for kernel approximation and attention-weighted sums in terms of this quantity, estimates it from a Gram matrix of sampled queries and keys, and trains nonlinear features layer by layer to approximate the kernel.

What carries the argument

The load-bearing object is the degrees of freedom $N^*_\lambda=\operatorname{tr}(\Sigma(\Sigma+\lambda I)^{-1})$ of the attention kernel, where $\Sigma$ is the integral operator of $K(x,y)=\exp(x^\top y/\sqrt{d})$ over the distribution of queries and keys. This trace counts the effective number of kernel directions whose eigenvalue exceeds the tolerance $\lambda$, and it controls how many features are needed in the approximation. The optimal sampling density $q_\lambda(z)\propto \langle \phi(\cdot;z),(\Sigma+\lambda I)^{-1}\phi(\cdot;z)\rangle_{L^2(\rho)}$ is the theoretical target that makes $M$ proportional to $N^*_\lambda$, while Algorithm 1 approximates the trace by $\operatorname{tr}(\hat\Sigma(\hat\Sigma+\lambda I)^{-1})$ using the empirical Gram matrix of sampled queries and keys. The layerwise-trained PRF features are the practical stand-in for exact sampling from $q_\lambda$.

What would settle it

Run distillation twice with the same total budget on one pretrained model: once with Algorithm 1's DoF allocation and once with the per-layer feature counts permuted in reverse order of estimated DoF. If the reverse allocation performs as well as the DoF allocation on downstream tasks, the claim that the degrees-of-freedom profile is what drives the gain is falsified. A complementary calculation on a synthetic kernel with known eigenvalues can check whether the empirical approximation error at the DoF-selected $M$ matches the Theorem 3 bound; a much larger error would indicate that the practical estimator and trained features do not realize the theoretical guarantee.

Watch

Extended reading notes

Core claim

The central claim is that the number $M$ of random features needed in a linear-attention layer is proportional, up to a logarithmic factor, to the degrees of freedom $N^*_\lambda = \operatorname{tr}(\Sigma(\Sigma+\lambda I)^{-1})$ of the attention kernel's integral operator. Sampling features from the density $q_\lambda(z)\propto \langle \phi(\cdot;z),(\Sigma+\lambda I)^{-1}\phi(\cdot;z)\rangle_{L^2(\rho)}$ makes this count optimal, and Theorem 3 shows that with $M \gtrsim N^*_\lambda\log N^*_\lambda$ both the kernel and attention-weighted sums are approximated with error of order $\lambda$ with high probability. Because the degrees of freedom depend on the input distribution, the required dimension differs across layers and heads. The paper turns this into Algorithm 1, which estimates the degrees of freedom from a finite Gram matrix, takes the maximum over heads in each layer, and rescales the resulting dimensions to a fixed inference budget; it then learns PRF features layerwise with an L2 or softmax kernel-matching loss. Empirically, on GPT-2 and Pythia-1B, the DoF-based dimension choice matches or beats fixed-dimension distillation while keeping inference cost unchanged.

Load-bearing premise

The method depends on believing that the layer-by-layer complexity numbers estimated from a sample of real queries and keys, together with the features learned one layer at a time, behave as well as the exact quantities in the theory; if those estimates are off, the promised efficiency gain disappears.

Editorial extensions

If this is right

  • A fixed total feature budget is spent more efficiently: layers with small effective kernel dimensionality receive fewer features and layers with large dimensionality receive more, reducing approximation error without increasing total compute.
  • Distilled models using DoF-selected dimensions match or outperform fixed-dimension baselines on downstream tasks, and capping dimensions at the head size yields faster inference while retaining accuracy.
  • The selected dimension profile shows that attention complexity is layer-dependent and head-dependent: early and middle layers need more features, later layers fewer, and only a few heads in a layer carry high degrees of freedom.
  • Layerwise training with a softmax or L2 kernel-matching loss approaches end-to-end distillation quality at lower training cost, which helps when distilling long-sequence models.
  • The approximation guarantee for attention-weighted sums applies to a broad class of value functions, so the error control covers actual attention outputs and not only individual kernel entries.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • Extension: because the DoF estimate uses only unlabelled queries and keys, the same allocation rule should transfer to larger models and other modalities without retraining; a direct test is whether the selected dimension profile stays stable as the sample size $J$ and context length grow.
  • Extension: the large head-to-head spread in estimated DoF suggests a per-head allocation or head-pruning variant that could push the same compute budget further; the paper notes this direction but does not implement it.
  • Extension: the DoF profile could be validated as a complexity diagnostic by correlating it with independent measures such as attention entropy or effective rank; a mismatch would indicate the Gram-matrix estimate, not the kernel itself, is driving the allocation.
  • Extension: if the finite-sample estimate is biased on short contexts, the selected dimensions should grow with sequence length; measuring $\hat N_\lambda$ across lengths would reveal whether the fixed-tolerance rule needs a length-dependent correction.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 5 minor

Summary. The paper proposes a method to choose per-layer feature dimensions in linear attention distillation by estimating the statistical degrees of freedom N*_λ = tr(Σ(Σ+λI)^{-1}) of the attention kernel. It presents a theoretical bound (Theorem 3), derived from Bach (2017), showing that the number of random features M needed for a target approximation error scales with N*_λ when features are importance-sampled from the optimal density q_λ. It then gives Algorithm 1 to estimate N*_λ from a finite Gram matrix, allocates layerwise dimensions proportional to the estimate under a fixed average compute budget, and trains nonlinear PRF features layerwise with L2 or softmax losses. Experiments on GPT-2 and Pythia-1B compare DoF-based allocation with fixed-dimension baselines and report average downstream accuracy gains.

Significance. If the central claim were fully established, the paper would be a useful step toward data-adaptive and layer-specific linear attention: the DoF measure is interpretable, the layerwise training is considerably cheaper than end-to-end distillation, and the empirical results on two pretrained models are encouraging. The paper also explicitly ships the proof of its main theorem and hyperparameter details, which aids reproducibility. However, the significance is conditional because the theoretical guarantee applies to random features sampled from q_λ, while the evaluated method trains features by gradient descent; no argument or experiment connects the two. The DoF estimate also lacks a concentration analysis, and the budget-constrained optimality claim goes beyond the per-layer error bound. These gaps make the theoretical framing stronger than what the experiments can currently support.

major comments (4)
  1. [3.1, 3.3] Theorem 3 and the preceding discussion guarantee the M ∝ N*_λ scaling only for i.i.d. features z_1,...,z_M drawn from the optimal density q_λ with weights 1/q_λ. Section 3.3 instead trains features z and weights α by minimizing ℓ(z) (L2 or softmax loss), and the paper gives no argument that these trained features behave like q_λ samples or that their approximation error decays with M according to the DoF curve. Since both the Fix and DoF experiments use trained features, the empirical comparison is consistent with the theory but does not validate it; a direct comparison against random features sampled from an estimated q_λ (or a theoretical analysis of the trained features) is needed to support the central claim.
  2. [3.2, Algorithm 1] The estimate \tilde N_λ = tr(\tilde Σ(\tilde Σ + λI)^{-1}) is computed from J randomly sampled queries/keys, but no concentration bound links \tilde N_λ to N*_λ, and J is not reported in the experiments. The selected dimensions in Table 1 therefore have an unquantified dependence on J and on the arbitrarily chosen λ (2^{-4} for GPT-2, 2^{-8} for Pythia-1B). The paper should either provide a finite-sample guarantee for the DoF estimate or report a sensitivity analysis over J and λ to show the allocation is robust.
  3. [3.1, 3.2] The statement that the proposed dimension selection 'achieves smaller error under a fixed computational budget' is not derived from Theorem 3. The theorem bounds the approximation error of a single kernel for a given λ and shows M should scale with N*_λ; it does not solve a budget-constrained allocation across layers or control the final attention-output error. The step in Section 3.2 that sets t^{-1} to match the average cost is a heuristic, and the paper should either prove an optimal-allocation statement or soften the claim to 'matches the per-layer error target at equal total cost.'
  4. [Theorem 3, Eq. (1)] Item (ii) of Theorem 3 bounds the unnormalized integral ∫ v(x)K(x,·)dρ(x), whereas the attention output in Eq. (1) is the normalized ratio of such integrals. The footnote about W_K invertibility does not address the error in the denominator, so the theorem as stated does not directly guarantee approximation of attention outputs. The authors should either extend the bound to the normalized quantity or state explicitly that the theory covers the kernel and the numerator only.
minor comments (5)
  1. [3.1, Theorem 3] Theorem 3's sample-complexity condition uses M ≥ (4N_{q,λ}/t) log(64N_{q,λ}/(δt)) in the statement, while Lemma 5 and the proof of Theorem 6 use M ≥ (4N_{q,λ}/t^2) log(64N_{q,λ}/(δt^2)); please harmonize the statement with the proof.
  2. [3.1] The quantity N_{1,λ} is used ('the degrees of freedom are always smaller than N_{1,λ}') but N_{1,λ} is not defined; presumably it is N_{q,λ} for q≡1, i.e., sampling from τ.
  3. [Algorithm 1] 'Fed the sequences in X' should be 'Feed the sequences in X'.
  4. [Table 7] Table 7 reports next-token prediction results where DoF with direct loss (5.2651) is worse than Fix with direct loss (3.9657), while the text in Section 4.2 emphasizes the opposite for downstream tasks; this discrepancy should be discussed.
  5. [Appendix C] The values of J (Gram matrix sample size) and the procedure for choosing λ are not reported; please include them in Appendix C.

Circularity Check

0 steps flagged · score 0.0 of 10
full rationale

The derivation chain is not circular. The paper defines N*_lambda = tr(Sigma (Sigma + lambda I)^{-1}) as the degrees of freedom, relies on Bach (2017)'s Proposition 1 and Lemma 2 to show that, for i.i.d. draws from the optimal density q_lambda, the number of features M needed to reach an approximation error lambda is proportional to N*_lambda up to logarithmic factors, and then Theorem 3 converts this into an L2 error bound for the kernel and for kernel-based integration. Algorithm 1 estimates N*_lambda from a finite Gram matrix of queries/keys and sets M_s proportional to the estimate while matching a fixed average cost. The empirical claim that DoF-based allocation outperforms fixed allocation is tested on held-out tasks (PiQA, logiQA, ARC, etc.) and is not obtained by fitting the method to those tasks. No equation defining the selected dimension is also the quantity being predicted, and no fitted parameter is renamed as a prediction. The paper does cite work by one of its authors (Suzuki et al., 2020), but only as motivational background for layerwise complexity variation, not as load-bearing support for the main claim. The most serious concern is the one the paper itself admits in Section 3.3: the trained features are not shown to behave like i.i.d. draws from q_lambda, so Theorem 3 does not strictly justify the empirical feature-training setup. That is a gap in theoretical support, not a circular argument: the trained-feature experiments do not assume the conclusion they are testing, and the dimension-selection formula is derived before, and independently of, the downstream accuracy evaluation.

Assumptions & free parameters 2 free parameters · 3 assumptions · 0 invented entities

The central claim rests on two unproven practical assumptions: the finite-sample DoF estimate is accurate, and trained features behave like optimal importance samples. No new entities are introduced; the degrees-of-freedom quantity is from Bach (2017) and the PRF features from Choromanski et al. (2021).

free parameters (2)
  • lambda (regularization parameter) = 2^-4 for GPT-2, 2^-8 for Pythia-1B
    The DoF N*_λ and thus all selected feature dimensions depend on the choice of λ. The paper sets λ by hand without a sensitivity analysis or a principled selection criterion.
  • Gram matrix sample size J = not reported
    Algorithm 1 requires sampling J queries and keys to form the Gram matrix. J controls the accuracy of the DoF estimate, but the paper does not specify its value or analyze its effect.
assumptions (3)
  • domain assumption The softmax attention kernel K(x,y)=exp(x^T y / sqrt(d)) admits the PRF representation K(x,y)=E_{z~N(0,I)}[φ(x;z)φ(y;z)] with square-integrable φ.
    Section 2 cites Choromanski et al. (2021) for the PRF feature map; the theoretical framework of Section 3 requires this representation.
  • domain assumption The empirical Gram matrix eΣ formed from J samples of queries and keys is an accurate approximation of the integral operator Σ on L2(ρ).
    Section 3.2 defines eN_λ = tr(eΣ(eΣ+λI)^{-1}) as a proxy for N*_λ without providing concentration bounds or specifying J.
  • domain assumption The layerwise-trained features z_m and weights α_m approximate the optimal sampling density qλ well enough for the theoretical guarantees to be indicative.
    Section 3.3 replaces importance sampling from qλ with gradient-based training and does not prove that the trained features satisfy the assumptions of Theorem 3.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency." pith.science (2026). https://pith.science/paper/RJ5532KH

@misc{pith2026250703340,
  author       = {Pith},
  title        = {Pith review of: Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/RJ5532KH}},
  note         = {Machine review of arXiv:2507.03340}
}
read the original abstract

Linear attention has attracted interest as a computationally efficient approximation to softmax attention, especially for long sequences. Recent studies have 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 error 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.

Figures

Figures reproduced from arXiv: 2507.03340 by the authors.

Figure 1
Figure 1. An overview of our method. We treat exp x ⊤y/√ d  , which appears in softmax attention, as a kernel K(x, y), and perform distillation by approximating it with the inner product of nonlinear features. The required feature dimension to achieve a 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 specifi… view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

13 extracted references · 5 canonical work pages

  1. [4]

    Dosovitskiy

    A. Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale.arXiv preprint arXiv:2010.11929,

  2. [5]

    Kasai, H

    J. Kasai, H. Peng, Y. Zhang, D. Yogatama, G. Ilharco, N. Pappas, Y. Mao, W. Chen, and N. A. Smith. Finetuning pretrained transformers into rnns. InProceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10630–10643,

  3. [7]

    direct” loss. This is natural because the cross entropy loss for next-token prediction is used in “direct

    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 L2 loss (Softmax loss) just aims to make the attention kernel (attention weights) and its appr...

  4. [11]

    S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity.arXiv preprint arXiv:2006.04768,

  5. [12]

    • K : Rd × Rd → R is the positive definite kernel given byK(x, y) = Ez∼τ [ϕ(x; z)ϕ(y; z)], where τ is a probability measure on a measurable setZ, and ϕ : Rd × Z →R is a feature map

    12 —— Appendix —— A Proof of Theorem 3 We begin by recalling the definitions of some symbols used in the proof. • K : Rd × Rd → R is the positive definite kernel given byK(x, y) = Ez∼τ [ϕ(x; z)ϕ(y; z)], where τ is a probability measure on a measurable setZ, and ϕ : Rd × Z →R is a feature map. • Σ : L2(ρ) → L2(ρ) is an integral operator defined by(Σf )(x) ...

  6. [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,

  7. [2011]

    J. Liu, L. 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.08124,

  8. [2017]

    J. Wang, D. Paliotta, A. May, A. M. Rush, and T. Dao. The mamba in the llama: Distilling and accelerating hybrid models. arXiv preprint arXiv:2408.15237,

Show all 13 references
  1. [2018]

    L. Dong, S. Xu, and B. Xu. Speech-transformer: a no-recurrence sequence-to-sequence model for speech recognition. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), pages 5884–5888. IEEE,

  2. [2019]

    Sakamoto and K

    H. Sakamoto and K. Sato. Data-driven h2 model reduction for linear discrete-time systems.arXiv preprint arXiv:2401.05774,

  3. [2020]

    H. Chen, Z. Liu, X. Wang, Y. Tian, and Y. Wang. Dijiang: Efficient large language models through compact kernelization. arXiv preprint arXiv:2403.19928,

  4. [2021]

    Clark, I

    P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge.arXiv:1803.05457v1,

  5. [2024]

    Ravichandran, A

    K. Ravichandran, A. Jain, and A. Rakhlin. Using effective dimension to analyze feature transformations in deep neural networks. In ICML 2019 Workshop on Identifying and Understanding Deep Learning Phenomena,

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.