mlsanityData to 7 Oct 2026

Post-Norm can Resharpen Attention

arXiv 2510.08341 ↗First posted 9 Oct 2025cs.LG cs.AIcited by 0

Not on the map. Why?

Abstract (arXiv)

Length Generalization is the essential capacity of autonomous agents to perform tasks in longer contexts than those encountered during training. To systematically study this feat, we test how well models can approximate the next token distributions in algorithmic tasks. This is to take into account the realistic possibility of multiple next tokens being legal. We present a prototypical benchmark for this line of study: in the Set Complement Task, the model needs to output a uniform distribution over tokens not in the input. We prove a theorem that states simple transformers can length generalize on this task, however, with performance degradation due to attention dispersion. A mechanistic reading of how dispersion takes effect lets us discover a remedy: Post-Norm can Resharpen Attention. We present experimental evidence to support this idea. We also show that Exponential Moving Averages can help the issue of noisy gradients that arises when many next tokens are legal. We validate the general applicability of our proposed methods on a suite of formal language experiments. Our source code will be available upon publication.

Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 34 of its 45 references that are in the map. First posted on arXiv 9 Oct 2025.

Builds on

34 of 45 references.

References · 34

  1. 2017Attention Is All You Needcited by 59,932
  2. 2022Training language models to follow instructions with human feedbackcited by 15,712
  3. 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
  4. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
  5. 2021Prefix-Tuning: Optimizing Continuous Prompts for Generationcited by 3,878
  6. 2019Root Mean Square Layer Normalizationcited by 1,425
  7. 2021Swin Transformer V2: Scaling Up Capacity and Resolutioncited by 1,215
  8. 2022Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 smallcited by 1,130
  9. 2020On Layer Normalization in the Transformer Architecturecited by 906
  10. 2019Universal Adversarial Triggers for Attacking and Analyzing NLPcited by 709
  11. 2021CogView: Mastering Text-to-Image Generation via Transformerscited by 691
  12. 2022Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Taskcited by 417
22 more
  1. 2023Emergent Linear Representations in World Models of Self-Supervised Sequence Modelscited by 313
  2. 2019Theoretical Limitations of Self-Attention in Neural Sequence Modelscited by 286
  3. 2024Massive Activations in Large Language Modelscited by 277
  4. 2024Retrieval Head Mechanistically Explains Long-Context Factualitycited by 219
  5. 2023How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language modelcited by 199
  6. 2023What Algorithms can Transformers Learn? A Study in Length Generalizationcited by 198
  7. 2021Thinking Like Transformerscited by 182
  8. 2023Calibrated Language Models Must Hallucinatecited by 114
  9. 2022Adam Can Converge Without Any Modification On Update Rulescited by 112
  10. 2022Overcoming a Theoretical Limitation of Self-Attentioncited by 110
  11. 2021Self-Attention Networks Can Process Bounded Hierarchical Languagescited by 105
  12. 2023Why Do We Need Weight Decay in Modern Deep Learning?cited by 70
  13. 2020RNNs can generate bounded hierarchical languages with optimal memorycited by 59
  14. 2024A Formal Framework for Understanding Length Generalization in Transformerscited by 44
  15. 2024Counting Like Transformers: Compiling Temporal Counting Logic Into Softmax Transformerscited by 39
  16. 2025Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMscited by 34
  17. 2025Peri-LN: Revisiting Normalization Layer in the Transformer Architecturecited by 32
  18. 2024Weight decay induces low-rank attention layerscited by 27
  19. 2024Language Model Probabilities are Not Calibrated in Numeric Contextscited by 18
  20. 2022Language models are better than humans at next-token predictioncited by 14
  21. 2025EMA Without the Lag: Bias-Corrected Iterate Averaging Schemescited by 2
  22. 2025Distribution Prompting: Understanding the Expressivity of Language Models Through the Next-Token Distributions They Can Producecited by 1

Built on by

Not cited yet.