# Post-Norm can Resharpen Attention

arXiv: https://arxiv.org/abs/2510.08341 · HTML: https://mlsanity.com/paper/2510.08341/

## Key facts

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

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Universal Adversarial Triggers for Attacking and Analyzing NLP](https://mlsanity.com/paper/1908.07125/) (2019)
- [CogView: Mastering Text-to-Image Generation via Transformers](https://mlsanity.com/paper/2105.13290/) (2021)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [Emergent Linear Representations in World Models of Self-Supervised Sequence Models](https://mlsanity.com/paper/2309.00941/) (2023)
- [Massive Activations in Large Language Models](https://mlsanity.com/paper/2402.17762/) (2024)
- [Theoretical Limitations of Self-Attention in Neural Sequence Models](https://mlsanity.com/paper/1906.06755/) (2019)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)
- [How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model](https://mlsanity.com/paper/2305.00586/) (2023)
- [What Algorithms can Transformers Learn? A Study in Length Generalization](https://mlsanity.com/paper/2310.16028/) (2023)
- [Thinking Like Transformers](https://mlsanity.com/paper/2106.06981/) (2021)
- [Calibrated Language Models Must Hallucinate](https://mlsanity.com/paper/2311.14648/) (2023)
- [Adam Can Converge Without Any Modification On Update Rules](https://mlsanity.com/paper/2208.09632/) (2022)
- [Overcoming a Theoretical Limitation of Self-Attention](https://mlsanity.com/paper/2202.12172/) (2022)
- [Self-Attention Networks Can Process Bounded Hierarchical Languages](https://mlsanity.com/paper/2105.11115/) (2021)
- [Why Do We Need Weight Decay in Modern Deep Learning?](https://mlsanity.com/paper/2310.04415/) (2023)
- [RNNs can generate bounded hierarchical languages with optimal memory](https://mlsanity.com/paper/2010.07515/) (2020)
- [A Formal Framework for Understanding Length Generalization in Transformers](https://mlsanity.com/paper/2410.02140/) (2024)
- [Counting Like Transformers: Compiling Temporal Counting Logic Into Softmax Transformers](https://mlsanity.com/paper/2404.04393/) (2024)
- [Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs](https://mlsanity.com/paper/2502.15938/) (2025)
- [Peri-LN: Revisiting Normalization Layer in the Transformer Architecture](https://mlsanity.com/paper/2502.02732/) (2025)
- [Weight decay induces low-rank attention layers](https://mlsanity.com/paper/2410.23819/) (2024)

## 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.