Mitigating Attention Localization in Small Scale: Self-Attention Refinement via One-step Belief Propagation (2025)

Abstract: 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 multi-hop 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.

## introduction transformer architecture @xcite forms the backbone of modern large language models (llms), demonstrating remarkable performance across diverse natural language processing tasks. at the core of these models lies the selfattention mechanism, which enables dynamic modeling of contextual relationships between tokens. despite its empirical success, several studies have identified limitations associated with attention localization, a phenomenon in which attention distributions collapse onto a limited subset of tokens. such localized attention tends to have low entropy @xcite , low-rank sparsity @xcite and sparsed attention (fig. 6 in the appendix), which negatively affect the representational power, training stability, and downstream model performance @xcite . this problem is particularly pronounced in small-scale transformer variants. their limited depth and width inherently restrict the ability to propagate informa1 the code is released at nakyunglee20/saobp tion across layers, making them more susceptible to localized attention than larger models. in standard self-attention mechanisms, each query token primarily attends to a few highly salient keys @xcite , thus modeling predominantly one-hop foot_0 dependencies. to address this limitation, we hypothesize that explicitly incorporating global context-defined here as multi-hop information flow-could mitigate the localization issue, particularly in compact architectures. to test this hypothesis, we propose self-attention onestep belief propagation (saobp), a method that integrates global interactions by leveraging belief propagation (bp) principles combined with a repulsive potts prior (fig. 1 ). our empirical evaluations show that saobp effectively alleviates attention collapse and consistently enhances performance across downstream tasks. to systematically investigate how saobp mitigates localization through multi-hop information flow and interpret its impact on model performance, we introduce global token dependency (gtd). gtd quantifies the relative attention mass contributed by intermediate, multi-hop transitions (≥ 2 hops) within the stochastic attention graph. our empirical analysis illustrated in fig. 3 and fig. 4 demonstrates that gtd serves as a principled diagnostic tool, capable of detecting specific layers and heads where attention collapses into overly localized patterns. ## preliminaries self-attention and attention entropy selfattention is mathematically defined by eq. ( 1 ), where q, k, and v represent learnable weight matrices, and b, h, and l denote the batch size, number of attention heads, and sequence length: the internal dynamics of self-attention can be analyzed using the concept of attention entropy @xcite , defined formally in eq. ( 2 ). attention entropy quantifies the uncertainty in attention distributions, offering insight into how information is structured within models. we leverage this metric to investigate attention entropy collapse. constructing belief propagation a belief propagation algorithm can be constructed with the following four components @xcite . let i and j be variable nodes connected by a factor node f ij . 1. messages: messages represent probabilistic information passed between variable and factor nodes during the inference process. formally, messages are denoted by mf @xmath0, m @xmath1, where m is the message space. ## experimental settings since our method modifies the self-attention computation mechanism, it is necessary to pretrain models from scratch. we conduct experiments using three bert-style transformer architectures that vary in model size: bert-mini, bert-small, and bert-medium ( ## results and analysis we hypothesize that explicitly modeling multi-hop token dependencies within the self-attention mechanism can alleviate attention localization and improve model performance. to empirically validate this hypothesis, we address the following three research questions. the main analysis is illustrated in fig. 4 . does gtd effectively capture the locality of attention? before analyzing our proposed algorithm, we first evaluate whether gtd reliably captures indirect, global token interactions and serves as a meaningful indicator correlated with model performance. given that few prior studies explicitly quantify intermediate relationships in self-attention, we compare gtd against established graph-theoretic metrics widely utilized in network analysis @xcite to assess its reliability. clustering coefficient (cc) quantifies the mean normalized triangle density in the graph, indicating the extent of local token clustering. higher cc values implies that tokens are organized into clusters with dense interconnections whereas lower cc corresponds to sparser connectivity @xcite . betweenness centrality (bc) measures the average unnormalized centrality of nodes based on their presence on shortest paths, highlighting whether attention is concentrated through a small number of intermediate tokens. lower bc values suggest more evenly distributed information flow @xcite . fig. 2 shows that gtd exhibits a positive correlation with cc (r ≈ 0.70) and a moderate negative correlation with bc (r ≈ -0.43). this finding aligns with our intended definition of gtd, as higher gtd values indicate globally interconnected attention graphs: multi-hop transitions lead to shorter token-to-token distances (higher cc) and reduce over-reliance on specific bridging tokens (lower bc). however, as shown in to examine gtd's practical utility in interpreting model performance, we analyze its correlations with the model performance across various tasks and checkpoints, as illustrated in fig. 3 . we consistently observe significant correlations (|r| > 0.5, @xmath3.02), showing that while the extent to which models utilize global contextual information may vary by task, gtd remains a meaningful and interpretable indicator of overall model effectiveness. does saobp mitigate entropy collapse? fig. 4 presents comprehensive experiments examining the effectiveness of saobp algorithms in regulating attention entropy and gtd values. the bp-high variant consistently maintains or enhances both indirect entropy and overall mean entropy, particularly in deeper layers across model sizes. these results confirm that the repulsive potts factor employed in bp-high effectively counters the entropy collapse commonly observed in transformer layers. in terms of gtd, bp-high achieves the highest or second-highest values across layers, demonstrating that message passing mechanism effectively incorporates multi-hop contextual information. overall, we observed that heads with gtd values around 0.6˘0.8 tended to achieve better performance on downstream tasks, whereas heads falling below ∼0.5 or above ∼0.85 were often associated with degraded performance or unstable training. among other saobp variants, bp-elemmul exhibits moderate multi-hop dependency, closely following baseline gtd values in earlier layers but aligning progressively with bp-high in deeper layers, leading to stable entropy levels. this behavior aligns with performance outcomes shown in collectively, these findings substantiate our claim that bp-high effectively mitigates entropy collapse and modulates gtd within a beneficial range, ensuring stable and expressive attention dynamics across various downstream tasks. our results also highlight design trade-offs, motivating future exploration of adaptive factor functions. does saobp improve downstream performance? as summarized in ## conclusion we propose saobp, a novel self-attention refinement framework that leverages structured message passing to explicitly capture intermediate token relationships frequently overlooked by standard atten-tion mechanisms. to rigorously analyze the impact of our proposed approach, we introduce the gtd, a metric designed to quantify the extent of multi-hop information flow and interpretably assess attention patterns in terms of entropy dynamics. through empirical analyses, we demonstrate that saobp effectively alleviates entropy collapse, promotes globally coherent attention distributions, and enhances performance across a diverse range of nlp tasks. notably, performance improvements from saobp are particularly significant in small-scale models, where architectural constraints inherently limit the emergence of global context. our findings suggest that explicit multi-hop regularization through bp can effectively compensate for reduced depth, enhancing expressivity and generalization in compact transformers. these insights underscore the broader potential of integrating graphbased inference methodologies with attention architectures to develop lightweight, interpretable, and high-performing small-scale language models.