Reviewer 1
Methods & technical-correctness reviewer (strict)
Summary
The paper proposes DIFF Transformer, which replaces softmax attention with differential attention. In each head, queries and keys are split into two groups. Two softmax maps are computed, and the second, scaled by a learnable scalar λ, is subtracted from the first (Eq. 1). λ is re-parameterised through exponentiated dot products of learnable vectors plus a layer-dependent λ_init (Eq. 2). The outputs are normalised per head (GroupNorm/RMSNorm) and multiplied by a fixed (1−λ_init) factor (Eq. 3). The motivation is an analogy to differential amplifiers: subtracting two maps is meant to cancel shared 'attention noise' on irrelevant context. The authors make several claims: - Appendix G argues that gradient flow matches standard attention, so Transformer hyperparameters can be reused. - Under matched parameters and FLOPs, the method reaches equal LM loss with roughly 65% of the parameters or tokens (Fig. 3). - It beats a LLaMA-style baseline on several fronts at 3B scale: downstream tasks, 64K long-context NLL, multi-needle retrieval, many-shot ICL and its order robustness, GPT-4o-judged contextual hallucination, activation outliers / low-bit attention quantisation, and post-hoc math reasoning. - Ablations at 1.4B scale examine head count, GroupNorm, and λ initialisation.
Strengths
- The modification is minimal and drop-in. The macro layout, FFN, and parameter/FLOP budget match a LLaMA-style baseline: heads are halved, so h = d_model/2d. The method also reuses FlashAttention kernels, with a measured throughput cost of only 5–12% (Table 7). This makes the core comparison reasonably controlled and the idea easy to adopt.
- Controlled comparisons against a same-recipe Transformer are run at several scales: 830M–13.1B in Fig. 3a, and a 350B-token 3B model in Fig. 3b and Table 8. The gains are consistent in direction across sizes and token budgets, which is more convincing than a single-scale comparison.
- The evaluation is broad and probes context utilisation rather than only perplexity. The retrieval gains are large and grow with distractors: at N=6, R=2, accuracy is 0.85 vs 0.55 in Table 2. Order-permutation variance of ICL is much smaller (TREC margin 4.0 vs 19.0 in Fig. 7a). These behavioural differences are consistent with the paper's stated motivation.
- The activation-outlier analysis is a useful side finding. Top-1 attention logits drop from 318.0 to 38.8 (Table 5), and the 4-bit attention-logit quantisation results in Fig. 8 point to practical implications for low-bit attention kernels.
- The ablations do address one obvious confound. Adding GroupNorm to an 8-head, d=256 Transformer leaves loss unchanged (3.086 vs 3.087, Table 6). Halving the head count alone also does not change loss.
Weaknesses
- Recommendation: marginal accept (6). The empirical gains are consistent and the method is simple, but the paper's mechanistic and theoretical claims are considerably stronger than what is derived or controlled for. The central key reason is that the 'noise cancellation' story is asserted by analogy and never established. A differential amplifier cancels noise only if both channels carry the same common-mode component. Nothing in Eqs. 1–2 or the training objective forces softmax(Q2K2ᵀ) to model the irrelevant-context mass of softmax(Q1K1ᵀ). The difference is simply a signed attention matrix whose rows sum to 1−λ, with entries that can be negative. No analysis shows that the two maps share a common component, or that the second map specialises in 'noise'. Useful evidence would include: - correlation between A1 and λA2 on non-answer tokens; - the learned λ values per layer; - the fraction of negative entries.
- The second key reason is that the ablations do not isolate the subtraction itself, yet the paper concludes the gains come 'from the differential attention mechanism'. The most direct control is missing: the same DIFF architecture with λ fixed to 0. That model keeps per-head Q/K dimension d=128, V dimension 2d, head-wise RMSNorm, and the (1−λ_init) scale, but has a single softmax. The comparison Transformer row in Table 6 uses QK dimension 256, so it differs from DIFF in QK width as well as in the subtraction. Also missing is a model with two independent softmax maps that are added or concatenated rather than subtracted. Without these controls, the gain could come from structural changes (narrower QK with wider V, extra normalisation interacting with it) rather than from cancellation.
- The gradient-flow argument in Appendix G is a heuristic, not a proof, although Section 2.1 says 'Appendix G proves' it. Several steps are not justified. - The key step, A ≈Θ A1 ≈Θ A2 ≈Θ A1−λA2, is asserted without justification. When λ approaches 1, or A1 ≈ A2 (the regime in which 'cancellation' is supposed to occur), A1−λA2 can be arbitrarily smaller than A1. - The GroupNorm Jacobian is reduced to a scalar Θ(1), dropping its projection structure. - The analysis holds only 'at the early training stage'. - 'Equivalent in magnitude up to constant factors' does not imply equivalent AdamW dynamics, because AdamW is only per-coordinate scale-invariant. The empirical fact that removing GroupNorm causes 'training instability' (Table 6, loss 3.122) suggests the dynamics are not Transformer-like without the added normaliser.
- The λ parameterisation in Eq. 2 is under-specified and under-analysed. - λ = exp(λq1·λk1) − exp(λq2·λk2) + λ_init is unconstrained. It can become negative, turning the operation into a sum, or exceed 1, making row sums negative. The paper gives no evidence of the range λ actually takes. - 'Synchronize the learning dynamics' is never explained. - A learnable-vs-fixed λ ablation is not reported. - The fixed (1−λ_init) post-norm multiplier ties output scale to the initial λ rather than the learned one. The ablation shows insensitivity only to λ_init, not to λ's parameterisation or learnability.
- The attention-score analysis (Table 3, Fig. 1) is close to circular and under-defined. Once one map is subtracted from another, low mass on context tokens is expected by construction. The paper also does not say how 'normalized attention scores' are computed when DIFF weights are signed, nor which layers, heads, or tokens are aggregated. It is further unclear whether the BOS sink is counted as noise: Fig. 1 shows BOS separately. Lower 'attention noise' is therefore not independent evidence for the denoising mechanism. Retrieval accuracy is the meaningful evidence, and it does not distinguish between mechanisms.
- The headline '~65% of parameters/tokens' is a fitted-curve interpolation, and its statistical support is weak. - Fig. 3a uses only 10B training tokens (40K steps × 0.25M) for every size up to 13.1B, so the larger models are heavily under-trained. The fitted savings may not transfer to compute-optimal regimes. - Each curve is fit to 5 single-seed points. - The 6.8B Transformer point visibly sits above its fit. - No confidence intervals on the ratio are given. The controlled 350B-token comparison (Table 8) shows a modest 0.8-point zero-shot average gain, and BoolQ is worse (60.1 vs 62.9).
- The method is not compared with other mechanisms that target the same phenomenon. Attention sinks and outliers are known to arise from heads needing to 'do nothing'. Bondarenko et al. (2024), which the paper cites, propose clipped softmax and gated attention specifically for this. Sink/register-style fixes (Xiao et al., 2023) are a related cheap alternative. Since the outlier and 'noise' claims overlap with these works, at least one such baseline at 1.4B scale would sharpen the causal claim that subtraction specifically is what helps.
- Several downstream claims rest on small, single-run evaluations without uncertainty estimates. - Hallucination uses 100 GPT-4o-judged samples per dataset on non-instruction-tuned base models. - ICL robustness uses 10 seeds over a single demonstration set. - The Section 1 math claim (in the paper body, Appendix C) selects the best checkpoint per model from a single run. Table 1 compares against external models with different data, tokenizers, and recipes, and partially reports numbers from a technical report, so it cannot support architectural conclusions. Confidence intervals or multiple seeds would be needed to treat the 0.07–0.19 hallucination gaps as robust.
Questions
- Could you train the λ≡0 control: identical DIFF architecture, QK dim 128, V dim 256, head-wise RMSNorm, (1−λ_init) scale, single softmax. Ideally also run a variant with '+λ' (sum) instead of '−λ' at 1.4B scale. If the λ≡0 model recovers most of the 0.025 loss gap in Table 6, the mechanistic claim would need to be revised. If it does not, my confidence in the subtraction being causal rises substantially.
- What values does λ take after training, per layer? Does it stay in (0,1)? What fraction of entries of A1−λA2 are negative, and how strongly are A1 and A2 correlated on non-answer tokens? This would directly test the common-mode-noise interpretation.
- In Appendix G, under what conditions does A1−λA2 have the same order of magnitude as A1 or A? In particular, what happens when A1 ≈ A2 and λ ≈ 1? Can the 'proof' be stated as a proposition with explicit assumptions, and does it hold beyond initialisation? If not, I would suggest rewording 'proves' to 'heuristically argues'.
- How are the normalised attention scores in Table 3 and Fig. 1 computed for signed DIFF weights: absolute values, clipping, or renormalisation? Which layers and heads are averaged, and is the BOS token included in 'noise'?
- For the quantisation study, is absmax quantisation applied separately to each of the two logit matrices in DIFF, and per-head or per-tensor? The logit outliers may simply reflect DIFF's smaller per-map QK dimension (128 vs 256 in comparable Transformer variants). Could you report the same statistics for the 8-head, d=256 Transformer and for the λ≡0 control?
- Can you provide seed variance for the Table 6 ablations and bootstrap confidence intervals for the scaling-law savings ratios? Several ablation differences are ≤0.004 in loss, so it is unclear which are distinguishable from noise.
Flag For Ethics Review
No ethics review needed.