Gradient Boosting within a Single Attention Layer

Saleh Sargolzaei

cs.LGApr 3, 2026

Gradient Boosting within a Single Attention Layer

Authors: Saleh Sargolzaei

Organizations: University of Windsor

Abstract

Transformer attention computes a single softmax-weighted average over values -- a one-pass estimate that cannot correct its own errors. We introduce gradient-boosted attention, which applies the principle of gradient boosting within a single attention layer: a second attention pass, with its own learned projections, attends to the prediction error of the first and applies a gated correction. Under a squared reconstruction objective, the construction maps onto Friedman's gradient boosting machine, with each attention pass as a base learner and the per-dimension gate as the shrinkage parameter. We show that a single Hopfield-style update erases all query information orthogonal to the stored-pattern subspace, and that further iteration under local contraction can collapse distinct queries in the same region to the same fixed point. We also show that separate projections for the correction pass can recover residual information inaccessible to the shared-projection approach of Tukey's twicing. On 10M-token subsets of WikiText-103 and OpenWebText, gradient-boosted attention improves test perplexity by $6.0\%$ and $5.6\%$ over standard attention, outperforming both Twicing Attention and a parameter-matched wider baseline on both benchmarks, with two rounds capturing most of the benefit. We further show, both theoretically and empirically, that the mechanism requires the additive residual structure of Pre-LN transformers: under Post-LN, the same architecture degrades perplexity by $9.6\%$ .

Figures & tables

Figure 1: (a) Standard attention computes a single softmax-weighted average. (b) Gradient-boosted attention (

M{=}2

) adds a second pass where the prediction residual

\mathbf{r}=\mathbf{x}-\hat{\mathbf{y}}_{0}

is projected through separate

\mathbf{W}_{Q}^{(1)},\mathbf{W}_{K}^{(1)},\mathbf{W}_{V}^{(1)}

to produce queries, keys, and values. A learned gate

\mathbf{g}

controls the per-dimension correction magnitude.

Component	MART ( Friedman, 2001 )	Ours
Initial estimate	$f_{0}(x)$ (e.g., mean of targets)	$\hat{\mathbf{y}}_{0}=\mathrm{Attn}_{0}(\mathbf{x})$
Pseudo-residual	$r_{m}=y-F_{m-1}(x)$	$\mathbf{r}_{m}=\mathbf{x}-F_{m-1}$
Base learner	$f_{m}$ fit to $r_{m}$	$\mathbf{c}_{m}=\mathrm{Attn}_{m}(\mathbf{r}_{m})$
Shrinkage	Scalar $\eta_{m}\in(0,1]$	Gate $\mathbf{g}_{m}=\sigma(\mathbf{W}_{g}^{(m)}[\cdot])\in[0,1]^{d}$
Update	$F_{m}=F_{m-1}+\eta_{m}f_{m}$	$F_{m}=F_{m-1}+\mathbf{g}_{m}\odot\mathbf{c}_{m}$

Table 1: Correspondence between gradient boosting (MART) and gradient-boosted attention.

	$d{=}16,K{=}4$		$d{=}32,K{=}8$		$d{=}64,K{=}16$
	$\sigma{=}0.5$	$\sigma{=}0.8$	$\sigma{=}0.5$	$\sigma{=}0.8$	$\sigma{=}0.5$	$\sigma{=}0.8$
Random chance	25.0	25.0	12.5	12.5	6.3	6.3
Bayes optimal	80.7	61.9	68.8	45.3	58.1	32.6
One-step (trained)	71.1	55.0	53.1	36.4	46.5	26.1
Converged (trained, backprop)	25.2	25.2	12.1	12.1	6.7	6.7
Converged (trained, DEQ)	25.2	25.2	12.1	12.1	6.7	6.7

Table 2: Retrieval accuracy (%) on the synthetic pattern retrieval task. The Bayes-optimal predictor (softmax attention with

\beta=1/\sigma^{2}

, identity projections) provides an analytical ceiling. Both converged variants—whether trained by backpropagation through 20 iterations or by DEQ implicit differentiation—collapse to near chance, while one-step attention approaches the Bayes-optimal ceiling.

		Test Perplexity
Model	Params	WikiText-103	OpenWebText
Standard ( $d{=}256$ )	7.4M	$72.2\pm 0.3$	$114.9\pm 0.5$
Twicing ( $d{=}256$ )	7.4M	$69.6\pm 0.1$	$110.7\pm 0.3$
Standard ( $d{=}288$ , param-fair)	8.8M	$69.0\pm 0.1$	$110.2\pm 0.5$
Gradient-boosted ( $M{=}2$ )	8.7M	$\mathbf{67.9\pm 0.1}$	$\mathbf{108.5\pm 0.9}$

Table 3: Test perplexity (lower is better) on two benchmarks. All models use Pre-LN. Perplexity values are not comparable across datasets (different tokenizers). Gradient-boosted attention outperforms all baselines on both benchmarks, including a wider model with matched parameter count.

Figure 2: Left: WikiText-103 test perplexity (zoomed axis). Gradient-boosted attention outperforms all baselines including Twicing and a parameter-matched wider model. Right: Retrieval accuracy on the synthetic pattern retrieval task as a function of boosting rounds. The dotted line marks the Bayes-optimal ceiling (58.1%). Four rounds nearly match it (58.1%); the jump from 1 to 2 rounds captures most of the improvement.

Model	Params	Pre-LN	Post-LN
Standard ( $d{=}256$ )	7.4M	$72.2$	$77.2$
Standard ( $d{=}288$ , param-fair)	8.8M	$69.0$	$71.8$
Gradient-boosted ( $M{=}2$ )	8.7M	$\mathbf{67.9}$	$84.6$
$\Delta$ Boosted vs. Standard		$-6.0\%$	$+9.6\%$
$\Delta$ Param-fair vs. Standard		$-4.4\%$	$-7.0\%$

Table 4: Effect of normalization placement on WikiText-103 test perplexity. Under Pre-LN, gradient-boosted attention improves over standard by

6.0\%

. Under Post-LN, it degrades by

9.6\%

. The parameter-fair baseline improves under both placements, ruling out a generic optimization difficulty with extra parameters under Post-LN.

Rounds ( $M$ )	1	2	3	4	5
Accuracy (%)	46.0	55.5	57.6	58.1	57.0
$\Delta$ from $M{=}1$	—	+9.5	+11.6	+12.1	+11.0

Table 5: Effect of boosting rounds on pattern retrieval accuracy (%). Synthetic pattern retrieval task,

d=64

K=16

\sigma=0.5

. Bayes-optimal ceiling: 58.1%.

Gate type	None	Scalar	MLP	Baseline ( $M{=}1$ )
Accuracy (%)	55.4	56.5	55.6	48.2
$\Delta$ from baseline	+7.2	+8.3	+7.4	—

Table 6: Effect of gate type on pattern retrieval accuracy (%). All gated variants use

M{=}2

rounds. Same task as Table 5 .

Figure 3: Learned gate values per dimension for each transformer layer, averaged over 50 test sequences. The dashed line marks

g=0.5

. Gate magnitudes and variation differ across layers, with layer 1 applying the strongest and most selective correction.

Figure 4:

\ell_{2}

distance in

\mathbb{R}^{d_{h}}

(

d_{h}{=}64

) from the boosted output to the convex hull of round-0 value vectors

\mathrm{conv}(\mathbf{v}_{1}^{(0)},\ldots,\mathbf{v}_{t}^{(0)})

, measured per attention head per position. Zero indicates the output lies within the hull; all 2,400 measured outputs are strictly positive, confirming escape at every layer. Annotated values are per-layer means. The ranking mirrors gate magnitudes (Figure 3 ).

Figure 5: Left: Distribution of attention entropy across all layers and heads for standard attention, boosted round 0, and boosted round 1. Round 0 is more diffuse than standard; round 1 is more focused. Right: Mean entropy per layer. The boosted model learns a division of labor: round 0 casts a wider net than standard attention, while round 1 sharpens, especially in layers 1–2.

Figure 6: Two tokens where gradient-boosted attention corrects a prediction error. Blue bars show standard attention, gray bars show boosted round 0, and red bars show boosted round 1. Standard and round 0 spread weight broadly; round 1 concentrates on the most relevant context. Each title shows the target, the standard model’s prediction, and the boosted model’s prediction with cross-entropy loss. Layer 1, head-averaged.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Training tokens	10M (subset of WikiText-103)
Vocabulary	BPE, 16,384 types
Sequence length	256
Layers	4
Model dimension $d$	256 (288 for param-fair)
Attention heads	4

Appendix

Table 7: Hyperparameters for language modeling experiments.

Figure 7: Validation perplexity on WikiText-103 across training epochs (mean of 2 seeds). The dashed line marks the parameter-fair baseline’s final perplexity. Gradient-boosted attention reaches this level 3 epochs earlier (20% fewer optimizer steps), despite starting from a comparable first epoch.

Figure 8: Retrieval accuracy for standard vs. gradient-boosted attention (

M{=}2

, MLP gate) across configurations, with the Bayes-optimal ceiling shown for reference. Numbers above bars indicate the accuracy improvement in percentage points. Boosted attention closes most of the gap to the Bayes-optimal ceiling, especially at moderate noise levels (

\sigma\leq 0.5

K,V source	Test PPL	Val PPL
Residual $\mathbf{r}_{m}$ (ours)	$\mathbf{67.9}$	$\mathbf{68.1}$
Original input $\mathbf{x}$	$69.3$	$69.5$

Appendix

Table 8: Effect of key/value source on WikiText-103 test perplexity. Deriving all projections from the residual outperforms the cross-attention variant by 1.4 points.

Attention Mechanisms

Attention mechanisms let neural networks weight information from other tokens or regions according to relevance, enabling context-dependent processing.

See all 512 papers

Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training
Sep 28, 2026Junlin Chen, Daize Dong, Huanwei Di +9Softmax Attention Efficient Attention
Rutgers University · Carnegie Mellon University · Oracle +2
Beyond Linearity in Attention Projections: The Case for Nonlinear Queries
Mar 11, 2026Marko KarbevskiTransformer Attention Decoder-Only Language Models
In Simplicity Technologies
RoPE attention is an exact forward-pass gradient step with softmax intact
Sep 6, 2026Julie Huang, Maggie Chlon, Leon ChlonSoftmax Attention Rotary Positional Embeddings
Hassana Labs · Department of Information Engineering, Oxford University

Sep 28, 2026J/K move · Enter open · S save

Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training

Junlin Chen, Daize Dong, Huanwei Di, Haolong Jia+8

Rutgers University · Carnegie Mellon University · Oracle+2

Softmax Attention Efficient Attention Attention Mechanisms

Open paper AskPDF arXiv

BF16 is now standard in large-scale pretraining, including in fused attention kernels such as FlashAttention, and these kernels are widely trusted. When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN. Recomputing the attention backward of just two layers in FP32 removes almost all of the excess gradient. Part of the cause is known: a fused multiply-add in the forward softmax, so far treated as an extreme-input NaN case and never fixed in FlashAttention-3. Repairing it stops the blow-up, but the query gradient is still wrong by more than its own size, and training still drives attention logits to thousands of times their size under accurate gradients. The remaining error comes from a broken conservation law. The softmax score gradient sums to zero along every row, which makes the query gradient blind to where the keys sit as a group; rounding it to BF16 leaves a small nonzero sum that leaks the mean key into the gradient, and the leak grows exactly as late training makes keys large and attention sharp. We introduce GProj (gauge projection), which restores the zero sum after the cast with two rank-one corrections per row. It cuts the remaining median query/key gradient errors from 219%/13% to 0.34%/0.37%, on par with FP32 attention, for 4.7% more time per training step. In matched from-scratch runs it trains to the same loss as FP32 attention, while FlashAttention-3 and key smoothing both destabilize.

Abstract

Attention Mechanisms

Explore similar work

Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training

Beyond Linearity in Attention Projections: The Case for Nonlinear Queries

RoPE attention is an exact forward-pass gradient step with softmax intact