# Mitigating Attention Localization in Small Scale: Self-Attention Refinement via One-step Belief Propagation

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 33 references that are in the map.
- First posted on arXiv 9 Sep 2025.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [RACE: Large-scale ReAding Comprehension Dataset From Examinations](https://mlsanity.com/paper/1704.04683/) (2017)
- [MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases](https://mlsanity.com/paper/2402.14905/) (2024)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Discriminative Embeddings of Latent Variable Models for Structured Data](https://mlsanity.com/paper/1603.05629/) (2016)
- [A Comprehensive Survey of Small Language Models in the Era of Large Language Models: Techniques, Enhancements, Applications, Collaboration with LLMs, and Trustworthiness](https://mlsanity.com/paper/2411.03350/) (2024)
- [Structured Attention Networks](https://mlsanity.com/paper/1702.00887/) (2017)
- [SparseBERT: Rethinking the Importance Analysis in Self-attention](https://mlsanity.com/paper/2102.12871/) (2021)
- [MobiLlama: Towards Accurate and Lightweight Fully Transparent GPT](https://mlsanity.com/paper/2402.16840/) (2024)
- [Belief Propagation Neural Networks](https://mlsanity.com/paper/2007.00295/) (2020)
- [Attention Entropy is a Key Factor: An Analysis of Parallel Context Encoding with Full-attention-based Pre-trained Language Models](https://mlsanity.com/paper/2412.16545/) (2024)
- [Self-attention Networks Localize When QK-eigenspectrum Concentrates](https://mlsanity.com/paper/2402.02098/) (2024)
- [Entropy-Guided Attention for Private LLMs](https://mlsanity.com/paper/2501.03489/) (2025)
- [Taming Transformer Without Using Learning Rate Warmup](https://mlsanity.com/paper/2505.21910/) (2025)

## Built on by (most cited first, 3 of 3)

- [NerVE: Nonlinear Eigenspectrum Dynamics in LLM Feed-Forward Networks](https://mlsanity.com/paper/2603.06922/) (2026)
- [Stabilizing Transformer Training Through Consensus](https://mlsanity.com/paper/2601.22614/) (2026)
- [Neural expectation-maximization and meta-models](https://mlsanity.com/paper/2610.12741/) (2026)

## Abstract (arXiv)

Transformer-based self-attention mechanism serves as the core of modern language models, yet it often suffers from localization, where attentions collapse onto a limited subset of tokens and fail to capture long-range dependencies. To address this issue, we propose Self-Attention One-step Belief Propagation (SAOBP), a refinement framework that injects multi-hop relationships through a belief propagation process. To interpret and quantify these interactions, we introduce Global Token Dependency (GTD) that captures the relative contribution of multihop connections within the attention graph. Empirical results indicate that SAOBP helps prevent entropy collapse in deeper layers and adaptively maintains GTD at task-appropriate levels, thereby supporting improvements in model performance. Importantly, we observe competitive gains in small-scale models, highlighting its potential for improving inference quality in resource-constrained scenarios.