Post-Norm can Resharpen Attention
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866918313992388608 |
|---|---|
| author | Zsámboki, Pál Levi, Benjamin Smith, David Ansel Josef Kagalwala, Mitansh Kell, Arlington Liechty, Samuel Wang, Cong |
| author_facet | Zsámboki, Pál Levi, Benjamin Smith, David Ansel Josef Kagalwala, Mitansh Kell, Arlington Liechty, Samuel Wang, Cong |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_08341 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Post-Norm can Resharpen Attention Zsámboki, Pál Levi, Benjamin Smith, David Ansel Josef Kagalwala, Mitansh Kell, Arlington Liechty, Samuel Wang, Cong Machine Learning Artificial Intelligence 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. |
| title | Post-Norm can Resharpen Attention |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2510.08341 |