Title: Scale-invariant attention
Venue: NeurIPS 2025
Accepted: yes
Keywords: long context

LLM

Transformer

scale invariance
Abstract: One persistent challenge in LLM research is the development of attention mechanisms that are able to generalise from training on shorter contexts to inference on longer contexts. We propose two conditions that we expect all effective long-context attention mechanisms to have: scale-invariant total attention, and scale-invariant attention sparsity. Under a Gaussian assumption, we show that a simple position-dependent transformation of the attention logits is sufficient for these conditions to hold. Experimentally we find that the resulting scale-invariant attention scheme gives considerable benefits in terms of validation loss when zero-shot generalising from training on short contexts to validation on longer contexts, and is effective at long-context retrieval.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenge of long-context generalization in LLMs by proposing scale-invariant attention, which satisfies two key properties: scale-invariant total attention (the expected sum of unnormalized attention weights over any token range \\([t, t\\Delta)\\) is asymptotically constant) and scale-invariant attention sparsity (attention is denser in local ranges and sparser in global ranges, with entropy growing sub-logarithmically). Under a Gaussian assumption, the authors derive a position-dependent transformation of attention logits (\\(L_t = a_t \\bar{L}_t + m_t\\)) to satisfy these properties. Experimental results show that the proposed \"scale-invariant p-RoPE\" outperforms baselines (e.g., RoPE, LogN, ALiBi) in validation loss across various training/validation context lengths, enables strong zero-shot generalization from short to long contexts, and performs well in the needle-in-a-haystack retrieval task. The main contributions include introducing the two scale-invariant properties, deriving the logit transformation, and validating its effectiveness empirically\n\nSTRENGTHS AND WEAKNESSES:\n## Strength\n1. The paper rigorously defines scale-invariant total attention and sparsity, provides mathematical proofs (Lemmas 1, 2, and Theorem 1) for the conditions under which these properties hold, and derives a concrete logit transformation, ensuring theoretical soundness.\n2. The proposed method involves a simple position-dependent transformation of attention logits, with only one hyperparameter (\\(\\tau\\)) after simplification, making it easy to implement and integrate with existing frameworks like p-RoPE.\n3. The paper effectively distinguishes its approach from related methods (e.g., LogN’s position-independent scaling vs. its position-dependent transformation), highlighting its unique ability to balance local and global attention.\n\n## Weakness\n1. The paper only validates the proposed method on models with 162M and 304M parameters, like GPT2, while state-of-the-art LLMs typically have billions of parameters. The authors acknowledge this limitation, noting that resource constraints prevented testing on larger models, which raises questions about whether the conclusions generalize to commercially deployed large-scale models.\n2. Although the paper compares with methods like RoPE, LogN, and ALiBi, it omits comparisons with several recent long-context attention mechanisms (e.g., Infini-Attention, Ring Attention) that also aim to address length generalization. This gap makes it harder to situate the proposed method within the broader landscape of current long-context LLM research .\n3. The theoretical derivation of the scale-invariant attention transformation hinges on the assumption that base attention logits (\\(\\bar{L}_t\\)) follow an IID Gaussian distribution. While the paper includes a QQ plot for Llama-3B to partially verify this, it lacks in-depth analysis of how the method performs when logits deviate significantly from Gaussianity (e.g., skewed distributions or heavy tails), which could limit its robustness in real-world scenarios .\n4. The issue that this paper aims to address is significant, while the presentation and experimental desmonstration is poor for me. How the proposed method motivated by studying at entropy? how it affect the design? and why this factor not evaluated at experiments?\n\nQUESTIONS:\n(1) The paper only evaluates models with 162M and 304M parameters, but contemporary LLMs often have billions of parameters. How does scale-invariant attention perform on larger models (e.g., 10B+ parameters)? Larger models may exhibit different attention behaviors due to increased capacity; verifying scalability is crucial for assessing real-world applicability. For example, do the hyperparameters (e.g., \\(\\tau=10\\)) remain optimal, or do they need adjustment for larger architectures?\n\n(2) The logit transformation is derived under the assumption that base logits (\\(\\bar{L}_t\\)) are IID Gaussian. However, real attention logits in trained models may deviate from Gaussianity. The paper mentions a QQ plot for Llama-3B, but how robust is the proposed method when logits are non-Gaussian (e.g., skewed or with heavier tails)? Would the transformation still satisfy scale-invariant properties, or is additional adaptation needed?\n\n(3) While the paper compares with several baselines, it does not include some recent long-context methods, such as Infini-Attention or Ring Attention. How does scale-invariant attention compare to these methods in terms of performance, computational efficiency, and generalization? Including such comparisons would better situate the work within the broader landscape of long-context LLMs.\n\n(4) The paper mentions that scale-invariant RoPE underperforms scale-invariant p-RoPE due to RoPE’s low-frequency components interfering with the transformation. However, it does not elaborate on why low-frequency components cause this interference. Could this be mitigated by modifying RoPE’s frequency spectrum (e.g., adjusting the base \\(\\theta\\))? A deeper analysis of this interaction would clarify the choice of p-RoPE as the base.\n\n(5) The needle-in-a-haystack task uses fine-tuning, which is unusual (most studies use prompting). The paper states that prompting alone was insufficient, but it does not explain why (e.g., model size, pretraining data). Would larger models or more pretraining enable the proposed method to perform this task with prompting? This information is important for assessing its practical utility in retrieval scenarios without fine-tuning.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new attention mechanism—scale-invariant attention—aimed at improving large language models’ ability to generalize from training on short contexts to inference on long contexts. The authors formalize two desirable properties for attention over long sequences: scale-invariant total attention, which ensures that attention is evenly distributed across different positional ranges, and scale-invariant attention sparsity, which ensures that attention becomes increasingly sparse over more distant contexts. To achieve these properties, the authors propose a simple transformation of the attention logits using position-dependent multiplicative and additive terms. They provide theoretical proofs under Gaussian assumptions that this transformation satisfies both scale-invariant total attention and (weak) scale-invariant attention sparsity.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n1. The paper provides a solid mathematical foundation for its proposed attention mechanism, with clear definitions and provable guarantees for the desired properties.\n2. The proposed logit transformation is lightweight and easy to implement, requiring only position-dependent scaling and biasing of logits.\n3. The method shows strong zero-shot performance without requiring long-context pretraining, which has implications for cost-efficient LLM deployment.\n\n**Weaknesses**\n1. Experiments are conducted on relatively small-scale models (162M and 304M parameters), which may limit confidence in applicability to larger foundation models.\n2. The authors do not analyze whether their logit transformation impacts convergence speed or training stability, which may be useful for practitioners.\n3. The three lines drawn in Figure 2 have low differentiation, so it is recommended that the author use different colors to mark them.\n4. The performance in Table 2 is slightly lower at 4k. It is recommended that the author analyze the reasons.\n\nQUESTIONS:\n1. Can the authors compare their method to memory-based approaches or retrieval-augmented LLMs, particularly on long-context retrieval tasks?\n2. Does the position-dependent logit scaling affect optimization dynamics? Empirical data on convergence speed or stability would be helpful.\n3. Can the authors provide more detailed empirical or theoretical analysis on why p-RoPE pairs better with scale-invariant attention than RoPE? Frequency spectrum visualizations could help.\n4. While the study focuses on GPT-style models, could this method extend to encoder-decoder or encoder-only models like BERT or T5?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes \"scale-invariant attention\" to improve the ability of large language models (LLMs) to generalize from short training contexts to longer inference contexts. This method aims for attention to function stably across various scales, from local to global, and is specifically designed to satisfy two desirable properties: \"scale-invariant total attention\" and \"scale-invariant attention sparsity\". Based on a Gaussian assumption, the authors theoretically derive a method to achieve these properties by applying a simple, position-dependent affine transformation to the attention logits. In experiments, the proposed method, when combined with p-RoPE, was demonstrated to outperform prominent existing methods (such as ALiBi and LogN) in terms of validation loss. It showed particularly excellent performance in zero-shot generalization from a short context (4k) to a long context (64k). Furthermore, it maintained high performance in a long-context information retrieval task.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The effectiveness of the proposed method is demonstrated through a multi-faceted and convincing set of experiments, using multiple model sizes, various training/validation context lengths, and a range of strong baseline methods for comparison. The result showing superior performance in a 16x zero-shot generalization task is particularly impressive and highlights the strength of the method.\n- The proposed transformation on logits is a simple affine transformation using the logarithm of the position index, which can be easily integrated into existing Transformer architectures. This makes the research highly practical.\n\n**Major Weaknesses**\n- The theoretical derivation hinges on the assumption that attention scores follow an IID (independent and identically distributed) Gaussian distribution. The paper's justification for this assumption is insufficient. While Appendix J presents a QQ plot for a Llama-3B model, this is just a single case. It is particularly questionable whether logits remain IID with respect to position after applying positional encodings like RoPE. The main text lacks a discussion on how valid this assumption is for real-world models and what the implications would be for the method's performance if the assumption is violated.\n\n**Minor Weaknesses**\n- While Appendix H explores the lengthscale hyperparameter τ, the analysis of why τ=10 is optimal is somewhat superficial. The intuition that it represents the \"size of 'chunks' we attend over\" is helpful, but a more in-depth analysis of the relationship between the optimal τ, model architecture, and data characteristics would have further strengthened the paper.\n\nQUESTIONS:\n- Regarding the IID Gaussian assumption that is central to your theory: while the QQ plot in Appendix J is interesting, do you believe the IID assumption holds even after applying position-dependent encodings like RoPE? How would you predict the effectiveness or performance of your proposed transformation to be affected if this assumption is only partially met?\n- The LogN trick also aims to control the attention distribution for long contexts. Your method differs by using a position-dependent transformation and shows superior results in your experiments. Do you envision any scenarios or task types where the simpler, position-independent LogN trick might be preferable? Or does your position-dependent approach have fundamental advantages that would make it superior in almost all long-context situations?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to transform attention scores to achieve scale invariance. It introduces two notions of scale invariance, namely equal sum and sparsity of attention weights across token ranges. The paper derives a close-form transformation function under some assumptions that enables scale-invariant total attention, while empirically demonstrates that the same transformation also achieves scale-invariant attention sparsity. The experiments demonstrates that the training small-scale models with the proposed attention scheme achieves lower validation loss compared to the baselines, especially when the context is longer than during training. Further, the model performs better than the baselines on long-context retrieval following fine-tuning.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths:\n- The transformation applied on attention scores is mathematically proven to satisfy the definition of scale invariance under certain assumptions.\n- The empirical results suggest the potential effectiveness of the method.\n\n## Weaknesses:\n- **Motivation:** The choice of scales (powers of 10) appears arbitrary and seems mainly for mathematical convenience. Additionally, it is unclear why a language model should satisfy the scale invariance properties. The paper asserts that scale invariance facilitates generalization on long contexts, but this claim is not substantiated by the toy experiments presented in the paper.\n\n- **Experiments:** The paper lacks solid experiments. Validation loss is the only metric, yet there is no evidence showing its correlation with the accuracy of any downstream tasks. The NIH experiment is also not convincing, as the model is fine-tuned to overfit this very mechanical task. Ideally, one would perform scaling experiments and run a suite of LLM benchmarks to show the method works and scales well. While there may be constraints on computational resources, legitimate experiments are essential to support the claims for acceptance in a peer-reviewed conference. Nonetheless, this work could be a valuable tech report and may gain traction if the method turns out to be effective with community effort.\n\n- **Visualization:** It helps to plot the distribution of attention weights to help visualize scale invariance.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nThe paper has addressed the limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}