Transformer attention is typically implemented using softmax normalization, which enforces attention weights with unit sum normalization. While effective in many settings, this constraint can limit flexibility in controlling attention magnitudes and may contribute to overly concentrated or unstable attention patterns during training. Prior work has explored modifications such as attention sinks or gating mechanisms, but these approaches provide only limited or indirect control over attention reweighting. We propose Affine-Scaled Attention, a simple extension to standard attention that introduces input-dependent scaling and a corresponding bias term applied to softmax-normalized attention weights. This design relaxes the strict normalization constraint while maintaining aggregation of value representations, allowing the model to adjust both the relative distribution and the scale of attention in a controlled manner. We empirically evaluate Affine-Scaled Attention in large-scale language model pretraining across multiple model sizes. Experimental results show consistent improvements in training stability, optimization behavior, and downstream task performance compared to standard softmax attention and attention sink baselines. These findings suggest that modest reweighting of attention outputs provides a practical and effective way to improve attention behavior in Transformer models.
Figures & tables
Figure 1: Comparison of baseline softmax, attention sink, and Affine-Scaled Attention. Attention skew toward the first token (red) is progressively reduced. Here, i denotes the index over key tokens, and the attention weight distributions are illustrated for a single query.
Figure 2: Training dynamics of attention logits. QKT logits mean over training steps for the 3B baseline and attention sink models, averaged across layers, heads, and token positions.
Figure 3: Per-query attention weight sum distributions for (a) Attention sink and (b) Affine-Scaled Attention. Dashed lines indicate the mean and median.
Figure 4: Per-layer first-token attention weights for the 1B model.
Figure 5: Layer-wise and head-wise heatmaps for the 3B model showing effective attention reweighting across methods.
Figure 6: Layer-wise mean attention entropy for the 3B model across baseline, attention sink, Affine-Scaled Attention methods.
Figure 7: KD training loss curves over training steps for (a) 0.5B, (b) 1B, and (c) 3B models across baseline, attention sink, and Affine-Scaled Attention methods, showing that Affine-Scaled Attention consistently achieves lower loss throughout training.
Accuracy (%)
Model
Method
Δ Params (M)
Loss
C4 PPL
Wino
PIQA
Hella
ARC-c
ARC-e
BoolQ
Avg.
0.5B
Baseline
0.0
2.82
71.7
52.6
68.6
45.1
27.6
51.5
58.6
50.6
Sink
0.0
2.83
81.5
54.9
68.2
44.3
26.7
51.9
58.7
50.8
Affine-Scaled
0.4
2.82
70.6
53.2
68.8
45.1
28.2
52.5
60.2
51.3
1B
Baseline
0.0
2.73
61.9
55.6
72.4
53.1
35.2
62.5
61.5
56.7
Sink
0.0
2.73
68.3
56.6
72.7
53.2
34.0
62.0
61.7
56.7
Table 1: Comparison of baseline, attention sink, and Affine-Scaled Attention methods for 0.5B/1B/3B models. We report the parameter overhead ( Δ Params), training Loss, C4 perplexity, and downstream accuracies on common sense reasoning tasks, along with the overall average. Best accuracy within each model size is highlighted in bold.
Figure 8: Gradient norm dynamics during training for (a) 0.5B and (b) 1B models across baseline, attention sink, and Affine-scaled Attention.
Accuracy (%)
Model
Method
Δ Params (M)
Loss
C4 PPL
Wino
PIQA
Hella
ARC-c
ARC-e
BoolQ
Avg
1B
Baseline
0.0
2.73
61.9
55.6
72.4
53.1
35.2
62.5
61.5
56.7
Gated
67.1
2.71
56.6
57.7
71.9
53.9
33.9
62.5
62.8
57.1
Affine-Scaled + Gated
68.2
2.71
56.2
58.1
72.0
54.5
35.3
64.3
64.0
58.0
3B
Baseline
0.0
2.65
44.9
59.0
75.1
60.8
38.4
68.1
63.6
60.8
Gated
264.2
2.65
44.6
61.5
74.4
61.4
40.4
68.8
63.1
61.6
Table 2: Comparison of baseline, Gated Attention, and Affine-Scaled + Gated Attention methods for 1B and 3B models. We report parameter overhead ( Δ Params), training loss, C4 perplexity, and downstream accuracies on common sense reasoning tasks, along with the average. Bold indicates the best result within each model size.
Act. Function
Sigmoid
Linear clipping
CSR Avg.
57.6
58.0
Table 3: Average common sense reasoning performance of the 1B model comparing sigmoid and linear clipping nonlinear functions.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Attention heatmaps for the 8th layer of the 1B model, across (a) Baseline, (b) Attention sink, and (c) Affine-Scaled Attention. The baseline concentrates attention on the first token. Attention sink reduces the weight on the first token, but the maps still show substantial concentration on the first two positions (the sink token and the first token). In contrast, Affine-Scaled Attention mitigates this positional over-concentration via input-adaptive modulation.
Figure 10: KD loss curves in the late training phase (steps 7k–9k) for 0.5B / 1B / 3B models. We compare baseline, attention sink, and Affine-Scaled Attention.
Model
Method
Mean
Std
Variance
CV
0.5B
Baseline
0.563
0.507
0.257
0.90
Sink
0.551
0.509
0.259
0.92
Affine-Scaled
0.561
0.451
0.203
0.80
1B
Baseline
0.907
1.983
3.934
2.19
Sink
0.902
1.988
3.951
2.20
Affine-Scaled
0.842
1.702
2.897
2.02
Appendix
Table 4: Gradient norm statistics during early training (steps 0–500). Variance and coefficient of variation (CV) are computed from raw gradient norms without smoothing.
Model
Method
#Spikes ( k=6 )
#Spikes ( k=9 )
Qwen-0.5B
Baseline
40
30
Sink
44
31
Affine-Scaled
42
22
LLaMA-1B
Baseline
53
28
Sink
67
35
Affine-Scaled
54
27
Appendix
Table 5: Extreme gradient norm spike counts during early training (steps 0–500). A spike is defined as a training step where the gradient norm exceeds a robust run-specific threshold, τ=median(g)+k⋅MAD(g) , with k∈{6,9} . MAD denotes the median absolute deviation, a robust measure of variability that is less sensitive to outliers than standard deviation. Lower spike counts indicate fewer extreme gradient events and smoother optimization dynamics.
Figure 11: Comparison of a sigmoid and a custom clipping-linear function.