Adaptive Head Budgeting for Efficient Multi-Head Attention Bilal FAYE1 , Abdoulaye MBAYE2 , Hanane AZZAG3 , Mustapha Lebbah4
arXiv:2604.22583v1 [cs.LG] 24 Apr 2026
e-mail: [email protected], [email protected], [email protected], [email protected]
Abstract— Transformers have become the dominant architecture across a wide range of domains, largely due to the effectiveness of multi-head attention in capturing diverse representation subspaces. However, standard multi-head attention activates all heads uniformly for every input, regardless of task requirements or input complexity. In many scenarios, particularly for coarse-grained tasks such as text classification, the relevant information is often global and does not require the full diversity of attention heads. As a consequence, using a fixed number of heads can introduce unnecessary computational cost or lead to suboptimal performance when the allocation does not match the input. To address this limitation, we introduce BudgetFormer, a Transformer architecture equipped with an adaptive multi-head attention mechanism that dynamically allocates computational resources. Our approach learns, for each input, both a head budget corresponding to the number of attention heads required, and a relevance distribution that selects the most informative heads. We also propose a training strategy based on an exploration and exploitation trade-off, allowing the model to discover effective head configurations before converging to efficient usage patterns. Experiments on text classification tasks of varying complexity show that our method reduces inference cost in terms of FLOPs and memory, while also achieving performance that can surpass standard full multi-head attention. These results highlight the potential of adaptive head allocation as a principled approach to improving both efficiency and effectiveness in Transformer models.
I. I NTRODUCTION Transformers have become the dominant architecture in natural language processing and beyond, driven by the effectiveness of self-attention mechanisms in modeling long-range dependencies [1, 2]. In particular, multi-head attention enables the model to capture diverse representation subspaces by projecting inputs into multiple parallel attention heads. This design has been central to the success of large-scale models across tasks such as language understanding, generation, and classification. However, the computational cost of self-attention scales quadratically with the sequence length, making it a major bottleneck in practice [1]. This limitation becomes especially pronounced in autoregressive generation, where tokens are processed sequentially and inference latency accumulates over time. To mitigate this issue, techniques such as key-value caching are commonly used to reuse past computations and reduce redundant operations during decoding [3, 4]. Despite these optimizations, the cost of attention remains significant, particularly in large models and long-context settings. A broad range of methods has been proposed to improve the efficiency of Transformers. Model compression
techniques such as knowledge distillation reduce model size while preserving performance [5]. Sparse and approximate attention mechanisms aim to alleviate the quadratic complexity by restricting attention patterns or using low-rank approximations [6]–[8]. Other approaches include token pruning, which removes less informative tokens during inference, and early exiting strategies that adaptively reduce the depth of computation [9, 10]. While effective in certain settings, these methods often require architectural modifications, introduce approximation errors, or rely on heuristics that may not generalize well across tasks. In this work, we focus on a complementary and largely underexplored dimension of efficiency: the adaptive use of attention heads. Standard multi-head attention activates all heads uniformly for every input, regardless of its complexity or the nature of the task. This can be suboptimal, especially in coarse-grained tasks such as text classification, where the relevant information is often global and does not require the full diversity of attention heads. Using a fixed number of heads may therefore lead to unnecessary computation or inefficient allocation of model capacity. To address this limitation, we introduce BudgetFormer, a Transformer architecture equipped with an adaptive multi-head attention mechanism that dynamically allocates computational resources at the head level. For each input, the model estimates a head budget corresponding to the number of attention heads required, and selects the most informative heads based on learned relevance scores. In addition, we propose a training strategy based on an exploration and exploitation trade-off, allowing the model to effectively discover and refine head usage patterns. Our contributions can be summarized as follows:
We propose an adaptive multi-head attention mechanism that learns to allocate a variable number of attention heads per input based on its complexity. • We introduce a training strategy that balances exploration and exploitation to learn efficient and robust head selection policies. • We demonstrate that our approach reduces inference cost in terms of FLOPs and memory usage, leading to more frugal models with lower computational and environmental footprint. • We validate our method on text classification tasks of varying complexity, showing that BudgetFormer can outperform standard full multi-head attention while using fewer computational resources. •
II. R ELATED W ORK Improving the efficiency of Transformer models has become a major research direction due to the high computational and memory cost of self-attention. Existing approaches can be broadly categorized into model compression, sparse and approximate attention, token-level adaptivity, and conditional computation mechanisms. While these methods have shown promising results, they often operate at the level of tokens, layers, or full attention maps, leaving the adaptive allocation of attention heads relatively underexplored. A. Model Compression and Pruning Model compression techniques aim to reduce the size and computational cost of Transformers while preserving performance. Knowledge distillation methods train smaller student models to mimic larger teachers [5]. Structured pruning approaches remove redundant components such as weights, neurons, or attention heads based on importance criteria [11]. Head pruning in particular has revealed that many attention heads are redundant and can be removed with minimal performance degradation. For instance, pruning strategies based on search or saliency metrics can eliminate a significant fraction of heads without loss in accuracy [12]. More recent works extend this idea by combining head pruning with block or token pruning, highlighting the redundancy present in both attention maps and head structures [13]. However, these approaches are typically static, requiring pruning decisions to be made offline and applied uniformly across all inputs. This limits their ability to adapt computation dynamically based on input complexity. B. Sparse and Approximate Attention Another line of work focuses on reducing the quadratic complexity of self-attention by introducing sparsity or approximation. Methods such as Longformer, BigBird, and Performer replace dense attention with structured or kernel-based approximations [6]–[8]. These approaches achieve sub-quadratic complexity while maintaining strong empirical performance. Despite their efficiency gains, sparse attention methods often rely on predefined patterns or approximations that may restrict the expressiveness of the model. In particular, fixed sparsity structures can limit the ability to capture global dependencies when needed [6]. C. Token Pruning and Adaptive Sequence Processing Token-level methods aim to reduce computation by dynamically selecting or pruning tokens during inference. Techniques such as PoWER-BERT and subsequent works remove less informative tokens based on learned importance scores. More recent approaches introduce progressive or dynamic token pruning strategies that adaptively refine the token set across layers [14]. Recent work has further explored adaptive token retention
and pruning in both NLP and vision settings, demonstrating significant reductions in FLOPs while maintaining accuracy [15]. Additionally, dynamic pruning methods have been proposed to jointly prune tokens, heads, and attention blocks during inference, highlighting the redundancy present across multiple dimensions of the Transformer architecture [13]. However, token pruning methods may struggle in tasks where all tokens contribute to the final prediction, such as fine-grained reasoning or dense prediction tasks. Moreover, pruning decisions are often irreversible within a forward pass, which can lead to information loss. D. Conditional Computation and Early Exiting Conditional computation approaches aim to adapt the amount of computation to the difficulty of each input. Early exiting methods allow models to produce predictions at intermediate layers, reducing average inference depth [16]. Similarly, adaptive-depth Transformers dynamically select the number of layers to apply per input, achieving substantial reductions in computation [10]. Mixture-of-Experts models extend this idea by routing inputs to a subset of expert networks, enabling scalable conditional computation [17]. More recent works also explore skipping layers or dynamically adjusting network depth based on input complexity [18]. While effective, these approaches primarily operate at the level of layers or feed-forward modules. They do not explicitly address the allocation of attention heads within each layer, which remains fixed in standard architectures. E. Discussion Across these lines of work, a common theme is the presence of significant redundancy in Transformer computations, whether at the level of tokens, layers, or attention structures. In particular, recent analyses show that only a subset of attention heads contributes meaningfully to global information processing, while many heads focus on local or redundant patterns. Despite this observation, existing methods either remove heads statically or treat all heads uniformly during inference. This suggests a gap in current approaches: the lack of fine-grained, input-dependent allocation of attention heads. In contrast, our approach focuses on dynamically allocating attention heads on a per-input basis. Rather than pruning or approximating attention globally, we learn to estimate a head budget and select the most relevant heads for each input. This enables a more flexible and fine-grained form of conditional computation that is particularly well suited for coarse-grained tasks such as classification, where the required level of attention diversity may vary significantly across inputs. III. BACKGROUND A. Multi-Head Self-Attention Let X ∈ RB×N ×D denote a sequence of N input tokens, where B is the batch size and D the model dimension. In
Transformer encoders, self-attention operates by projecting X into queries, keys, and values through linear mappings: Q = XWQ ,
K = XWK ,
V = XWV
with WQ , WK , WV ∈ R . These representations are partitioned into H heads, each of dimension dh = D/H, allowing the model to attend to information from multiple representation subspaces. For each head h, attention is computed as: Qh Kh⊤ √ Vh (2) Attnh (X) = Softmax dh The outputs of all heads are concatenated and projected back to the model dimension: (3)
This mechanism allows each head to capture distinct interaction patterns across the sequence, which has been identified as a key factor behind the empirical success of Transformer models [1]. B. Computational Complexity The computational cost of multi-head attention arises from both projection operations and pairwise interactions between tokens. Given an input of length N , the computation of attention scores involves forming the matrix product Qh Kh⊤ ∈ RN ×N for each head. This induces a quadratic dependency on the sequence length. More precisely, for a single layer, the dominant cost can be expressed as: Cattn (X) ≈ O(BN D2 ) + O(BHN 2 dh )
Mattn (X) ≈ O(BHN 2 )
(1)
D×D
MHA(X) = Concat(Attn1 (X), . . . , AttnH (X))WO
Consequently, the overall memory footprint of a Transformer encoder layer is dominated by the quadratic term:
(4)
where the first term corresponds to the linear projections and the second to the attention computation and aggregation. Using D = Hdh , the quadratic term becomes O(BN 2 D), which dominates for long sequences. An important observation is that this cost scales linearly with the number of heads H. All heads are computed independently, and their contributions are aggregated uniformly, regardless of their individual relevance to the input. As a result, increasing H improves representational capacity but also directly increases computational cost. C. Memory Requirements Beyond computation, memory consumption is another critical factor in Transformer models. In encoder architectures, all tokens are processed simultaneously, and intermediate attention representations must be stored during the forward pass. The attention score tensors for each head have size RN ×N , leading to a total storage cost proportional to BHN 2 . In addition, attention probabilities, intermediate projections, and output representations contribute linearly in BN D. During training, these activations must be retained for backpropagation, effectively doubling memory usage.
(5)
This scaling makes attention particularly expensive in settings with long sequences or large numbers of heads. D. Inference Efficiency in Encoder Models In encoder-based tasks such as text classification, inference typically processes the entire sequence in a single forward pass. While this avoids the sequential overhead of autoregressive decoding, the full attention computation remains necessary for all tokens and all heads. In this setting, the computational cost per layer remains proportional to HN 2 , and the model evaluates all attention heads regardless of the input structure. However, not all inputs require the same level of representational diversity. For instance, in coarse-grained classification tasks, the decision often relies on global semantic cues that can be captured by a subset of attention heads. This suggests that the uniform use of all heads may lead to over-computation, where some heads contribute marginally to the final representation while still incurring full computational and memory cost. E. Motivation The above analysis highlights two structural inefficiencies in standard multi-head attention. First, the cost of attention grows linearly with the number of heads, making head multiplicity a direct driver of computational and memory overhead. Second, the architecture assumes that all heads are equally useful for every input, which is unlikely to hold in practice, especially in tasks where the required level of abstraction varies across examples. These observations motivate the design of adaptive mechanisms that can modulate the number of active heads depending on the input. Instead of treating all heads uniformly, it becomes natural to consider a formulation in which only a subset of heads is selected or weighted more strongly, reducing unnecessary computation while preserving task-relevant information. In the next section, we build on this perspective and introduce an adaptive attention mechanism that learns to allocate a head budget and select informative heads dynamically. IV. M ETHOD : B UDGET F ORMER In this section, we introduce BudgetFormer, a Transformer encoder equipped with adaptive head-level computation. Unlike standard multi-head attention, which activates all heads uniformly for every input, our approach learns to dynamically allocate a computational budget over attention heads. This allows the model to adapt its level of computation to the complexity of each input, while maintaining full model capacity during training.
For active heads (mi = 1), attention is computed as:
A. Adaptive Head Budgeted Attention Given an input representation X ∈ R a global summary via mean pooling:
B×N ×D
, we compute
N
h=
1 X Xi ∈ RD . N i=1
(6)
The budget network predicts a continuous control variable: s = σ(fθ (h)),
(7)
D
where fθ : R → R is a feed-forward network and σ is the sigmoid function. The scalar s ∈ (0, 1) represents the fraction of attention heads to activate. We then compute head scores using a learnable projection gϕ : RD → RH : t , (8) z = gϕ (h) + ϵ · σmax 1 − T where ϵ ∼ N (0, 1), t is the training step, T is the total number of steps, and σmax controls the initial exploration noise magnitude. Head selection probabilities are obtained via a temperature-scaled softmax: exp(zi /τ (t))
p i = PH
j=1 exp(zj /τ (t))
,
(9)
where the temperature follows: t , τ (t) = τmin + (τmax − τmin ) exp −γ T
(10)
with τmax > τmin > 0 controlling the transition from exploration (uniform distribution) to exploitation (peaked distribution). The importance weight assigned to each head is defined as: wi = s · H · pi , where H is the total number of heads. Each head produces an output: Qi K ⊤ Attni (X) = Softmax √ i Vi . dh
(11)
The final output is obtained by concatenation followed by a linear projection: Y = Concat(Ã1 , . . . , ÃH )WO .
(14)
Optionally, we enforce sparsity by selecting only the top-k heads during inference: k = max(1, ⌊s · H⌋),
The final output is: Y = Concat(Ã1 , . . . , ÃH )WO .
(15)
Let Sk ⊂ {1, . . . , H} denote the indices of the k largest values of pi . We define a binary mask: ( 1 if i ∈ Sk mi = (16) 0 otherwise
(19)
Only the subset Sk is evaluated, with: |Sk | = k = ⌊s · H⌋,
(20)
which reduces the effective number of computed heads from H to k. The combination of the noise scale σmax and the temperature schedule τ (t) defines a gradual transition from exploration to exploitation, while the budget variable s controls the global computational allocation per input. B. Training Objective Optimizing only the task loss is insufficient in our setting, as it does not constrain how computational resources are allocated across attention heads. In particular, the model may converge to degenerate solutions where all heads are uniformly used or where the budget collapses to extreme values. We therefore define the following objective: L = Ltask + Lbudget + Lentropy ,
(21)
where Ltask is task-dependent (e.g., classification or regression), Lbudget controls the global allocation of heads, and Lentropy regulates head specialization. The budget s is constrained within a target interval [smin , smax ] using a quadratic hinge formulation. We first define the violation as: (22)
The budget loss is then given by: Lbudget = α(s) · v(s)2 ,
(13)
(17)
For inactive heads (mi = 0), the corresponding outputs are not computed: Ãi = 0. (18)
v(s) = max(0, smin − s) + max(0, s − smax ). (12)
We scale each head by its importance weight: Ãi = wi · Attni (X).
Ãi = wi · Attni (X).
(23)
where the scaling factor adapts to the magnitude of the violation: α(s) = min(αmax , αbase + v(s)).
(24)
This formulation allows the model to freely explore any value of s within the interval without penalty, while progressively increasing the constraint when s deviates from the desired range. The adaptive scaling prevents unstable behavior and avoids collapse toward trivial budgets. To control the distribution over heads, we introduce an entropy regularization term: Lentropy =
H X i=1
pi log pi ,
(25)
where pi are the head selection probabilities. Its influence is modulated over training through a time-dependent coefficient: 2t β(t) = βmax −1 . (26) T At early stages (t ≈ 0), β(t) < 0, which favors high-entropy distributions and encourages exploration across heads. Around mid-training, β(t) ≈ 0, reducing its effect. At later stages (t > T /2), β(t) > 0, which promotes low-entropy distributions and leads to sparse and specialized head usage. The entropy term is thus defined as: Lentropy = β(t)
H X
pi log pi .
(27)
i=1
The combination of the violation-based budget constraint and the entropy schedule enables a stable transition from exploration to exploitation, while explicitly controlling the computational footprint of the model.
while BudgetFormer reduces this to:
leading to:
Minference (X) ≈ O(BkN 2 ),
(35)
Minference ≈ s. MMHA
(36)
This reduction directly translates into lower memory usage and improved scalability for long sequences. Finally, since energy consumption is approximately proportional to the number of floating-point operations, BudgetFormer also reduces the inference-time carbon footprint: CO2 ∝ Cinference ∝ s. (37) Overall, BudgetFormer preserves the full expressivity of multi-head attention during training while enabling a controllable and input-adaptive reduction in computation, memory, and energy usage at inference time. V. E XPERIMENTS
C. Complexity Analysis
A. Experimental Setup
We analyze the computational and memory complexity of BudgetFormer and compare it to standard multi-head attention. In standard attention, the cost of a single layer is dominated by: CMHA (X) ≈ O(BN D2 ) + O(BHN 2 dh ), (28)
We evaluate BudgetFormer on text classification tasks by comparing it to a standard Transformer encoder using full multi-head attention. All models are trained and evaluated on five widely used benchmark datasets, covering a range of domains and classification granularities. The datasets used in our experiments are summarized in Table I. They include topic classification (DBpedia, AG News), sentiment analysis (IMDB, Yelp Review Full), and natural language inference (SNLI). For datasets without an official validation split, we use the test set as validation. As a baseline, we use a Transformer encoder composed of L = 4 layers, each with H = 8 attention heads and model dimension D = 768. BudgetFormer follows the same architecture, replacing the standard attention layer with the proposed adaptive head budgeted attention. The budget predictor fθ is implemented as a two-layer feed-forward network with a ReLU activation, mapping RD → R. The head scoring function gϕ is a single linear projection mapping RD → RH . Both models are trained for 10 epochs using the AdamW optimizer with a learning rate of 2 × 10−5 and a batch size of 16. The number of training steps T depends on the dataset size and is used consistently in the scheduling functions defined in Section IV. For BudgetFormer, the budget is constrained within [smin , smax ] = [0.1, 0.9], allowing the model to explore a wide range of computational allocations without bias toward extreme values. The training hyperparameters are set as follows: αbase = 0.001, αmax = 0.05, βmax = 0.05, σmax = 0.5, τmax = 2.0, τmin = 0.1, and γ = 5.0. In terms of model size, the baseline Transformer requires 197.58 MB of memory, while BudgetFormer requires 206.70 MB. This increase is due to the additional parameters introduced by fθ and gϕ , which remain lightweight compared to the attention layers and introduce negligible computational
where all H heads are computed for every input. The second term dominates and scales linearly with H. In BudgetFormer, additional computations arise from the budget and gating networks: Cbudget ≈ O(BD2 ) + O(BDH),
(29)
which are independent of the sequence length N and negligible compared to the attention cost. During training, all heads are evaluated: Ctrain ≈ CMHA + Cbudget ,
(30)
ensuring stable gradients and full exploration of the head space. The overhead induced by fθ and gϕ remains marginal relative to the quadratic attention term. During inference, only the top-k heads are computed, with: k = ⌊s · H⌋.
(31)
The attention cost becomes: Cinference ≈ O(BN D2 ) + O(BkN 2 dh ).
(32)
This yields a proportional reduction: k Cinference ≈ = s. CMHA H
(33)
Hence, the computational cost scales linearly with the predicted budget s, allowing input-dependent efficiency. A similar reduction applies to memory. In standard attention: MMHA (X) ≈ O(BHN 2 ),
(34)
Dataset DBpedia [19] AG News [19] IMDB [20] SNLI [21] Yelp Full [19]
Description Ontology classification News topic classification Sentiment analysis Natural language inference Review rating prediction
Train 560,000 120,000 25,000 549,367 650,000
Val 70,000 7,600 25,000 9,842 50,000
Test 70,000 7,600 25,000 9,824 50,000
Classes 14 4 2 3 5
TABLE I: Datasets used for evaluation. Dataset DBpedia AG News IMDB SNLI Yelp
Accuracy Transformer BudgetFormer 0.9830 0.9859 0.9099 0.9022 0.8354 0.8356 0.7835 0.8106 0.5810 0.6190
FLOPs (Test) Transformer BudgetFormer 4.40 · 1013 4.26 · 1013 4.78 · 1012 4.65 · 1012 1.61 · 1014 1.45 · 1014 12 6.18 · 10 6.05 · 1012 3.23 · 1014 2.58 · 1014
Carbon (gCO2) Transformer BudgetFormer 0.1468 0.1420 0.0159 0.0155 0.5376 0.4840 0.0206 0.0202 1.0751 0.8591
smean 0.085 0.212 0.601 0.364 0.198
TABLE II: Comparison between standard Transformer and BudgetFormer on test sets. BudgetFormer uses adaptive head selection at inference (top-k). FLOPs and carbon correspond to full evaluation over the test set.
overhead. All experiments are conducted on a single NVIDIA A100 GPU with 80GB of memory. B. Main Results We report the main results on five text classification benchmarks in Table II. We compare BudgetFormer against a standard Transformer encoder with identical architecture and training setup. For BudgetFormer, inference is performed using top-k head selection, where k = ⌊s · H⌋, enabling actual computational savings. BudgetFormer consistently achieves competitive or improved performance compared to the standard Transformer, while reducing inference cost. On DBpedia and SNLI, we observe clear accuracy gains, with improvements of +0.29 and +2.71 points respectively. On Yelp, the gain is even more pronounced (+3.8 points), suggesting that adaptive head selection is particularly beneficial for more complex or noisy datasets. On AG News, performance remains close to the baseline with a slight drop (-0.7), while still reducing computational cost. This indicates that for simpler datasets, aggressive head reduction may slightly affect performance, but remains controlled. From an efficiency perspective, BudgetFormer systematically reduces FLOPs and carbon emissions at inference. The reduction is directly correlated with the learned budget smean . For instance, on DBpedia, the model uses only 8.5% of heads on average, leading to lower computational cost with improved accuracy. On IMDB, where longer sequences require richer representations, the model allocates a higher budget (s ≈ 0.6), preserving performance while still reducing cost by approximately 10%. Importantly, these gains are achieved without modifying the training pipeline. During training, all heads remain active, ensuring stable optimization and full gradient flow. The additional overhead introduced by the budget and gating networks is negligible compared to the overall model size (approximately +9 MB in parameters), and does not significantly impact training cost. Overall, these results demonstrate that BudgetFormer
effectively adapts computational resources to input complexity, achieving a favorable trade-off between accuracy and efficiency. The model learns when fewer heads are sufficient and when more capacity is required, leading to both improved generalization and reduced inference cost. C. Efficiency Analysis We analyze the behavior of BudgetFormer along two complementary dimensions: (i) the evolution of the learned budget during training, and (ii) its adaptation to input complexity at inference. Training dynamics. We first study the evolution of the average budget smean on both training and validation sets, jointly with the validation accuracy. Figure 1 reports these curves for representative datasets. At early stages of training, the budget smean is relatively high, reflecting an exploration phase where multiple attention heads are actively used. As training progresses, smean gradually decreases, indicating a transition toward a more selective and efficient allocation of heads. Importantly, we observe a strong alignment between training and validation curves, with no noticeable gap. This suggests that the learned budget generalizes well and does not overfit to the training data. At the same time, the validation accuracy steadily improves and remains stable as smean decreases. This indicates that reducing the number of active heads does not harm performance. On the contrary, the model learns to discard redundant heads while preserving or improving predictive accuracy, highlighting an effective transition from exploration to exploitation. Adaptation to input complexity. We then evaluate how the predicted budget varies with input difficulty. Based on the results from Section II, we focus on two representative datasets: SNLI, and Yelp, which exhibit different levels of task complexity. For each dataset, we construct three categories of inputs: Simple, Medium, and Hard. Figure 2 presents the distribution
(a) SNLI
(b) Yelp
Fig. 1: Training dynamics showing the evolution of smean (train and validation) and validation accuracy over epochs.
(a) SNLI
(b) Yelp
Fig. 2: Distribution of the predicted budget s across input complexity levels (Simple, Medium, Hard) for SNLI, and Yelp.
of s across these categories. Across all datasets, we observe a consistent increase of s with input complexity. Simple inputs require only a small fraction of attention heads, while more complex inputs trigger higher budgets. This trend is particularly clear on SNLI, where logically challenging examples require more heads, and on Yelp, where nuanced sentiment leads to higher computational demand. These results demonstrate that BudgetFormer effectively adapts its computational effort to the input. The model allocates more resources when necessary while remaining efficient on simpler examples, leading to a form of conditional computation at the head level. D. Ablation Study We conduct two complementary ablations to isolate the roles of the budget predictor fθ and the head selection network gϕ . Ablation 1: Fixed budget (no fθ ). We remove the learned budget and fix s ∈ {0.1, 0.25, 0.5, 0.75, 1.0} while keeping gϕ trainable. Results are reported in Table III. We observe a strong degradation when s increases on several datasets
Dataset DBpedia AG News IMDB SNLI Yelp
s = 0.1 0.9846 0.8905 0.7296 0.6704 0.5818
s = 0.25 0.9478 0.8974 0.8036 0.7598 0.5813
s = 0.5 0.6202 0.8675 0.8246 0.5838 0.2737
s = 0.75 0.3245 0.8139 0.8289 0.3541 0.1013
s = 1.0 0.1670 0.7554 0.8162 0.3190 0.0743
TABLE III: Accuracy with fixed budget s (no learned fθ ).
(DBpedia, SNLI, Yelp). This shows that allocating more heads does not necessarily improve performance. Without adaptive control, larger budgets introduce noise through gϕ , leading to inefficient head utilization. In contrast, BudgetFormer (Table II) learns small but optimal budgets (e.g., smean ≈ 0.085 on DBpedia), achieving higher accuracy with fewer active heads. Ablation 2: Random head selection (no learned gϕ ). We fix s to the learned value from BudgetFormer and replace gϕ with random head selection. Results are shown in Table IV. The performance collapses across all datasets when head selection is random, even with the correct budget. This demonstrates that gϕ is essential to identify relevant heads. The budget s alone is insufficient: performance depends on
Dataset DBpedia AG News IMDB SNLI Yelp
Random gϕ 0.7511 0.8180 0.7493 0.3370 0.3244
BudgetFormer 0.9859 0.9022 0.8356 0.8106 0.6190
Data fraction 10% 25% 50% 100%
Transformer 0.6151 0.6814 0.7243 0.7835
BudgetFormer 0.6158 0.7075 0.7738 0.8106
smean 0.445 0.317 0.228 0.364
TABLE VI: Scaling training data size on SNLI.
TABLE IV: Impact of removing learned head selection gϕ (random gating).
which heads are selected, not only how many. These ablations highlight two key properties: (i) the budget must be adaptive (learned via fθ ), as fixed allocations are suboptimal and can introduce noise, (ii) head selection must be structured (learned via gϕ ), as random selection severely degrades performance. Together, fθ and gϕ enable BudgetFormer to allocate computation both quantitatively (how many heads) and qualitatively (which heads), explaining the gains observed in Table II. E. Generalization Across Model and Data Scales We evaluate whether BudgetFormer maintains its advantages when scaling (i) the model capacity and (ii) the amount of training data. All experiments are conducted on SNLI for controlled comparison. Scaling model capacity. We vary both the number of layers and attention heads, and compare against a standard Transformer of identical architecture. Results are summarized in Table V. BudgetFormer consistently Model 4L-8H 4L-12H 6L-12H 12L-12H
Transformer 0.7835 0.7847 0.7922 0.7808
BudgetFormer 0.8106 0.8080 0.8065 0.8193
smean 0.364 0.310 0.260 0.105
TABLE V: Scaling model depth (L) and heads (H) on SNLI.
outperforms the Transformer across all configurations. A key observation is that smean decreases as model capacity increases. For instance, with 12 layers and 12 heads, the model achieves its best accuracy (0.8193) while using only ∼10% of the heads on average. This indicates that larger models contain redundant heads, and BudgetFormer effectively exploits this redundancy by selecting only the most relevant ones. In contrast, the standard Transformer does not benefit as much from scaling, suggesting inefficient use of additional capacity. Scaling data size. We now vary the fraction of the training set used (10%, 25%, 50%, 100%). Results are shown in Table VI. BudgetFormer shows stronger robustness in low-data regimes. At 10% of the data, both models perform similarly, but BudgetFormer quickly surpasses the Transformer as more data becomes available. Notably, smean
adapts to the data regime: it is higher when data is scarce (0.445 at 10%), indicating that the model uses more heads to compensate for uncertainty, and decreases as more data becomes available (down to 0.228 at 50%). This reflects an adaptive trade-off between exploration and efficient computation. These results highlight two important properties. First, BudgetFormer scales better with model capacity by avoiding redundant computation and focusing on a subset of useful heads. Second, it adapts its computational budget to the amount of available data, using more resources when necessary and becoming more selective as learning stabilizes. This dynamic behavior leads to consistently better accuracy while maintaining controlled computational cost, demonstrating strong generalization across both model and data scales. F. Qualitative Analysis We conduct a qualitative analysis on DBpedia to better understand the behavior of the scaling factor s and the head selection distribution q across transformer blocks and semantic classes. DBpedia is particularly interesting as it exhibits relatively low smean values across several classes. Variation of smean across blocks and classes. Figure 3 presents the average s values across blocks and classes. We observe two key trends. First, smean varies significantly across classes, indicating that the model adapts its scaling behavior depending on the semantic difficulty of the class. Classes that are intuitively easier tend to exhibit lower smean , while more abstract or ambiguous classes show higher values. Second, smean consistently decreases from early to deeper blocks (Block 0 to Block 3). This suggests that earlier layers require stronger modulation (higher s), while deeper layers operate with more refined and stable representations, requiring less scaling. Variability of s. To assess whether s is dynamically adapted or fixed, we analyze its standard deviation across samples. Figure 4 shows that the standard deviation of s varies across both blocks and classes. This indicates that s is not a static parameter but adapts dynamically depending on the input instance. Higher variability is typically observed in intermediate layers and for more complex classes, suggesting that the model adjusts its scaling more aggressively when the classification task is harder or less certain.
Fig. 3: Mean s across blocks and classes on DBpedia.
Fig. 4: Standard deviation of s across blocks and classes.
Entropy of head selection. We further analyze the entropy of the head selection distribution q (Figure 6). The entropy measures how evenly attention is distributed across heads: High entropy indicates that multiple heads contribute similarly. • Low entropy indicates that only a few heads dominate.
•
Our results show relatively low entropy across most classes, particularly in deeper layers. This suggests that the model tends to concentrate importance on a small subset of heads rather than distributing it uniformly. This behavior aligns with the low smean observed earlier: when only a few heads are relevant, strong scaling is not required across all heads. Attention visualization. To better understand this phenomenon, we visualize attention maps for a representative example: ”The Amazon rainforest is the largest tropical rainforest in the world.” Figure 5 shows attention patterns across blocks and heads,
sorted by head importance q. We observe that: A small number of heads (with high q) capture most of the meaningful structure in the sentence. • These heads focus on key tokens such as Amazon, rainforest, and largest. • Remaining heads exhibit diffuse or less informative attention patterns. •
This supports the hypothesis that, for DBpedia, the task can often be solved using a limited number of highly specialized heads. Consequently, the model does not require strong global scaling, leading to lower smean values. Overall, our qualitative analysis reveals that: (i) s adapts both across layers and semantic classes, (ii) its variability reflects input-dependent modulation, (iii) head importance is highly concentrated (low entropy), and (iv) only a few heads are necessary to capture the relevant information for classification in DBpedia.
(a) Block 0
(b) Block 1
(c) Block 2
(d) Block 3
Fig. 5: Attention maps across transformer blocks for a DBpedia example. Each row corresponds to one block, and heads (H) are sorted by importance (q).
Fig. 6: Entropy of head selection distribution q.
These findings explain why DBpedia exhibits lower smean compared to other datasets: the task often relies on sparse and specialized attention rather than distributed processing.
VII. C ONCLUSION AND F UTURE W ORK
VI. L IMITATIONS
In this work, we introduced a dynamic scaling mechanism for transformer attention, enabling adaptive modulation of attention heads through learned scaling factors and head importance distributions. Our analysis showed that the model naturally learns to concentrate information on a small subset of heads, leading to sparse and efficient representations. Through qualitative analysis on DBpedia, we highlighted several key properties of the proposed approach: (i) the scaling factor s adapts across layers and semantic classes, (ii) its variability reflects input-dependent behavior, and (iii) attention tends to be concentrated on a few dominant heads, especially in simpler tasks. These findings suggest that transformer models exhibit an inherent tendency toward sparsity, which can be effectively exploited through adaptive mechanisms. Future Work. Several directions can be explored to extend this work. First, applying the proposed method to more complex tasks such as question answering, multi-hop reasoning, or long-context understanding would provide a better assessment of its generalization capabilities. These tasks require fine-grained token interactions, which are not fully captured by the current global pooling design. Second, integrating the approach into modern large language model architectures could further validate its scalability and effectiveness in real-world settings. Third, extending the framework to other modalities, such as vision or multimodal learning, would allow us to investigate whether similar sparsity patterns emerge beyond text. Finally, an important direction is the development of hybrid efficiency strategies that combine attention head pruning with token-level pruning. Such approaches could leverage both head sparsity and token sparsity to significantly reduce computational cost while maintaining performance, paving the way toward more frugal and scalable transformer models.
While our approach demonstrates strong performance and insightful behavior on classification tasks such as DBpedia, it presents several limitations. A key limitation lies in the design of fθ and gθ , which both rely on global pooling mechanisms. By construction, this implies that all tokens contribute equally to the computation of the scaling factor s and the head importance distribution q. As a result, token-level heterogeneity is not explicitly modeled. This design is well-suited for coarse-grained tasks, where global semantic information is sufficient to make accurate predictions. However, for more complex tasks such as question answering, reasoning, or fine-grained sequence understanding, this global aggregation becomes a bottleneck. In such settings, the importance of tokens is highly non-uniform and often depends on subtle contextual interactions. Consequently, the current formulation may fail to capture fine-grained token-level dependencies, limiting its expressiveness in scenarios where localized or structured information is critical. Another limitation is that the current framework focuses solely on head-level modulation. While this provides useful insights into attention sparsity, it does not exploit potential redundancies at the token level, which could further improve efficiency. Addressing these limitations requires moving towards mechanisms that preserve token-level importance while remaining computationally efficient. Designing compact representations that encode token contributions without significantly increasing computational cost is an important direction for future work.
R EFERENCES [1] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” Advances in neural information processing systems, vol. 30, 2017. [2] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” in Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pp. 4171–4186, 2019. [3] Z. Dai, Z. Yang, Y. Yang, J. G. Carbonell, Q. Le, and R. Salakhutdinov, “Transformer-xl: Attentive language models beyond a fixed-length context,” in Proceedings of the 57th annual meeting of the association for computational linguistics, pp. 2978–2988, 2019. [4] R. Child, S. Gray, A. Radford, and I. Sutskever, “Generating long sequences with sparse transformers,” arXiv preprint arXiv:1904.10509, 2019. [5] V. Sanh, L. Debut, J. Chaumond, and T. Wolf, “Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter,” arXiv preprint arXiv:1910.01108, 2019. [6] I. Beltagy, M. E. Peters, and A. Cohan, “Longformer: The long-document transformer,” arXiv preprint arXiv:2004.05150, 2020. [7] M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al., “Big bird: Transformers for longer sequences,” Advances in neural information processing systems, vol. 33, pp. 17283–17297, 2020. [8] K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, D. B. Belanger, L. J. Colwell, and A. Weller, “Rethinking attention with performers,” in International Conference on Learning Representations, 2021. [9] L. Hou, Z. Huang, L. Shang, X. Jiang, X. Chen, and Q. Liu, “Dynabert: Dynamic bert with adaptive width and depth,” Advances in Neural Information Processing Systems, vol. 33, pp. 9782–9793, 2020. [10] Y. Liu, F. Meng, J. Zhou, Y. Chen, and J. Xu, “Faster depth-adaptive transformers,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 35, pp. 13424–13432, 2021. [11] P. Ganesh, Y. Chen, X. Lou, M. A. Khan, Y. Yang, H. Sajjad, P. Nakov, D. Chen, and M. Winslett, “Compressing large-scale transformer-based models: A case study on bert,” Transactions of the Association for Computational Linguistics, vol. 9, pp. 1061–1080, 2021. [12] A. Parnami, R. Singh, and T. Joshi, “Pruning attention heads of transformer models using a* search: A novel approach to compress big nlp architectures,” arXiv preprint arXiv:2110.15225, 2021. [13] G. Jaradat, M. Tolba, G. Alsuhli, H. Saleh, M. Al-Qutayri, T. Stouraitis, and B. Mohammad, “Hybrid dynamic pruning: A pathway to efficient transformer inference,” arXiv preprint arXiv:2407.12893, 2024. [14] S. Goyal, A. R. Choudhury, S. Raje, V. Chakaravarthy, Y. Sabharwal, and A. Verma, “Power-bert: Accelerating bert inference via progressive word-vector elimination,” in International conference on machine learning, pp. 3690–3699, PMLR, 2020. [15] R. Liao, C. Zhao, J. Li, W. Feng, Y. Lyu, B. Chen, and H. Yang, “Catp: Cross-attention token pruning for accuracy preserved multimodal model inference,” in 2025 IEEE Conference on Artificial Intelligence (CAI), pp. 1100–1104, IEEE, 2025. [16] W. Zhou, C. Xu, T. Ge, J. McAuley, K. Xu, and F. Wei, “Bert loses patience: Fast and robust inference with early exit,” Advances in Neural Information Processing Systems, vol. 33, pp. 18330–18341, 2020. [17] W. Fedus, B. Zoph, and N. Shazeer, “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,” Journal of Machine Learning Research, vol. 23, no. 120, pp. 1–39, 2022. [18] T. Lawson and L. Aitchison, “Learning to skip the middle layers of transformers,” arXiv preprint arXiv:2506.21103, 2025. [19] X. Zhang, J. Zhao, and Y. LeCun, “Character-level convolutional networks for text classification,” Advances in neural information processing systems, vol. 28, 2015. [20] A. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts, “Learning word vectors for sentiment analysis,” in Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pp. 142–150, 2011. [21] S. Bowman, G. Angeli, C. Potts, and C. D. Manning, “A large annotated corpus for learning natural language inference,” in Proceedings of the 2015 conference on empirical methods in natural language processing, pp. 632–642, 2015.