PaperlayerView exact revision on arXiv

Source-linked AI summary

Softpick: No Attention Sink, No Massive Activations with Rectified Softmax

Zayd M. K. Zuhri, Erland Hilman Fuadi, Alham Fikri Aji

arXiv:2504.20966v4cs.LG

TL;DR

Softmax’s sum-to-one normalization is associated with attention sinks and massive activations, which pose challenges for quantization and low-precision training. The paper proposes softpick, a rectified replacement for softmax in transformer attention, and reports zero sink rate, sparse attention maps, and advantages for quantized models. It also reports that softpick does not yet scale reliably under the current training recipe.

  • Problem

    Softmax’s sum-to-one attention normalization is associated with attention sinks and massive activations that challenge quantization and low-precision training.

    Source: §8 · p. 8

  • Method

    Softpick is a rectified, non-sum-to-one drop-in replacement for softmax in transformer attention.

    Source: §3.1 · p. 2, §9 · p. 9

  • Results

    Softpick achieves a 0% sink rate and sparse attention maps, and quantized softpick models down to 2-bit precision outperform softmax.

    Source: §1 · p. 2, §9 · p. 9

  • Takeaways & Limitations

    The results indicate improved robustness under low-bit post-training quantization, especially at very low precision.

    Source: §9 · p. 9

    Under the current training recipe, softpick is competitive with softmax at 340M but underperforms at 1.8B and does not improve long-context retrieval in the reported experiments.

    Source: Limitations · p. 9

Abstract

from arXiv · show

We introduce softpick, a rectified, not sum-to-one, drop-in replacement for softmax in transformer attention mechanisms that eliminates attention sink and massive activations. Our experiments with 340M and 1.8B parameter models demonstrate that softpick achieves 0\% sink rate consistently. The softpick transformers produce hidden states with significantly lower kurtosis and creates sparse attention maps. Quantized models using softpick outperform softmax on standard benchmarks, with a particularly pronounced advantage at lower bit precisions. Our analysis and discussion shows how softpick has the potential to open new possibilities for quantization, low-precision training, sparsity optimization, pruning, and interpretability. Our code: https://github.com/zaydzuhri/softpick-attention

1 Introduction

The introduction questions whether attention scores must be nonzero probabilities summing to one, given softmax’s association with attention sinks and extreme hidden-state activations. It proposes softpick as a drop-in alternative and reports benchmark and model-behavior findings, while noting a performance gap at 1.8B.

1 Introduction

  • Softmax’s sum-to-one constraint is linked to attention sinks, while massive hidden-state activations are problematic for quantization and low-bit training.The introduction describes sinks as significant attention to a token that may be semantically irrelevant, often BOS.Source: §1 · pp. 1–2
  • Softpick reshapes softmax to allow rectified outputs without strict sum-to-one behavior and is proposed as a drop-in replacement for attention.Source: §1 · p. 2
  • Quantized softpick models down to 2-bit precision outperform softmax.Source: §1 · p. 2
  • Softpick produces more legible attention maps, a 0% sink rate, and no massive hidden-state activations.Source: §1 · p. 2
  • At 1.8B scale, the authors investigate softpick’s subpar performance and offer two possible hypotheses.Source: §1 · p. 2

2 Background

The background connects attention sinks and massive activations to softmax’s normalization, then reviews alternatives that relax normalization or alter attention architecture. It notes that existing approaches can add parameters or have reliability and depth-related drawbacks.

2 Background

3 Method

Softpick replaces softmax in attention with rectified outputs and a denominator that does not impose strict sum-to-one normalization. Its design retains gradient flow for negative inputs and supports an online implementation compatible with FlashAttention.

3.1 Softpick Function

  • Softpick uses rectified outputs and a numerically safe denominator, with a maximum-value shift and ε to avoid division by zero.The function is described as a drop-in replacement for softmax in attention.Source: §3.1 · p. 2

3.2 Design Rationale

  • The design adds null attention by rectifying shifted exponential outputs, while absolute values in the denominator preserve gradient flow for negative inputs.The numerator-denominator asymmetry also removes the strict requirement that attention scores sum to one.Source: §3.2 · p. 3

3.3 Derivative

  • The paper gives a Jacobian for numerically safe softpick and notes that its derivative is straightforward to implement when inputs are recomputed, as in FlashAttention.The denominator includes the absolute shifted exponentials and ε.Source: §3.3 · p. 3

3.4 FlashAttention

  • Softpick is compatible with FlashAttention because ReLU and absolute value preserve the multiplicative property for positive-valued multipliers, enabling an online version.The paper provides FlashAttention-2 forward- and backward-pass algorithms using softpick.Source: §3.4 · pp. 3–4

4 Experiments

The experiments train matched softmax and softpick Llamastyle transformers from scratch at 340M and 1.8B parameters, alongside short runs for sink-behavior analysis.

Experimental setup

  • Four Llamastyle transformers were trained from scratch: softmax and softpick variants at both 340M and 1.8B parameters.Each 340M model used 52 billion tokens and each 1.8B model used 104 billion tokens from the 100B FineWeb-Edu subset.Source: §4 · p. 4
  • Short comparison runs used the 340M configuration to analyze softpick against other sink-mitigation or sink-free attention variants.Unless otherwise specified, these runs used 2×A100 or 4×MI210 GPUs.Source: §4 · p. 4
  • The comparison runs were limited to 10k steps because baseline softmax sinks emerge within that window, and they were not benchmarked.The stated purpose was to diagnose sink behavior and activation outliers while keeping ablation compute manageable.Source: §4 · p. 4

5 Results

Softpick’s quantized models compare favorably at 340M, with the advantage narrowing at lower precisions at 1.8B. Full-precision benchmark results and training loss are less favorable at the larger scale.

Training Loss and Gradients

  • The final training-loss gap is 0.004 for the 340M model and 0.12 for the 1.8B model; gradient clipping at maximum norm 1.0 prevents observed instability from large gradients.The 340M model was trained for 100,000 steps (52B tokens).Source: §5 · p. 4

Benchmarks

  • On ARC-e, Lambada, PIQA, SciQ, and Wikitext validation, softpick is equal to or slightly better than softmax at 340M, but worse at 1.8B.The reported measures are benchmark accuracy, Lambada perplexity, and Wikitext validation perplexity.Source: §5 · p. 5

Quantization

  • At 340M, quantized softpick models consistently outperform softmax, and the performance gap closes at lower precisions at 1.8B.Figure 3 covers 2-, 3-, 4-, and 8-bit quantization.Source: §5 · p. 5, Figure 3 · p. 5

6 Analysis

Softpick produces sparse attention maps and a 0% sink rate at both model sizes, while reducing hidden-state activation extremes. Against other sink-free variants, softpick and scaled softpick most strongly suppress sinks and substantially reduce activation extremes.

6.1 Attention Maps

  • Softpick attention scores are functionally sparse, concentrated in specific spots and regions, with surrounding exact zeros.Some softpick heads still attend to a wider range of tokens.Source: §6.1 · p. 5
  • Softpick attention maps contain 99.34% exact zeros in the 340M model and 95.63% in the 1.8B model.The measure averages sparsity across 10 text samples and counts exact zeros only in the causal lower triangle; softmax zeros arise from floating-point underflow.Source: §6.1 · p. 6

6.4 Comparison with Other Sink-Free Methods

  • Softpick and scaled softpick most strongly suppress sinks while substantially reducing activation extremes relative to softmax; other variants help less consistently.Table 3 compares sink rate, activation statistics, and exact-zero sparsity across sink-free attention variants.Source: Table 3 · p. 6, §6.4 · p. 7

6.2 Attention Sink

  • Softmax heads with heavy BOS sinks can become trigger-token-specific or shut off under softpick when not needed.The paper describes one head shutting off on one input and attending to specific tokens on others.Source: §6.2 · p. 6
  • Softpick has a 0% sink rate at both thresholds and both model sizes, based on 1,000 SlimPajama validation samples.The thresholds were 0.3 and 0.2; sink rate measures heads whose average first-column attention exceeds the threshold.Source: §6.2 · p. 6

6.3 Massive Activations

  • For the 340M model, softpick reduces hidden-state kurtosis from 33,510 to 340 and lowers activation extrema by an order of magnitude.The last-layer softpick activations match softmax, while earlier layers have smaller activations.Source: §6.3 · p. 7

7 Scalability

Softpick does not improve passkey retrieval over softmax in these experiments, and retrieval collapses beyond the models’ trained context range. Dead heads decline during training, but remain more prevalent in the larger model.

7.1 Long Context and Underscoring

  • 0.0 passkey accuracy at length 4,986 for all methods; softpick is comparable to, but does not outperform, softmax across sequence lengths.The models were trained to a maximum context length of 4,096, making length 4,986 out-of-distribution.Source: §7.1 · pp. 7–8, Table 4 · p. 8
  • As context grows, sparse attention may underscore relevant tokens, weakening their value signals; query-key score scaling did not improve retrieval and worsened training and downstream performance.The authors present underscoring as a hypothesis for weaker retrieval and the larger model’s performance.Source: §7.1 · p. 8

7.2 Dead Attention Heads

  • 17.19% to 9.11% dead heads in the 340M model and 42.87% to 20.90% in the 1.8B model over training; the larger model retains more dead heads.At the 10^-2 threshold, the end-of-training percentages are 22.66% for 340M and 50.39% for 1.8B.Source: Figure 6 · p. 8, §7.2 · p. 8
  • Masking dead heads at inference damages the 340M model more than the 1.8B model, while changing the masking threshold from 10^-6 to 10^-2 has nearly negligible performance impact.The authors interpret this as support for using a higher threshold to identify dead heads.Source: §7.2 · p. 8

8 Discussion and Implications

Softpick’s suppression of activation outliers is associated with stronger accuracy retention under aggressive quantization, particularly at lower bit widths. Sparse attention may also support efficiency, pruning, and interpretability opportunities.

8 Discussion and Implications

  • Softpick models retain accuracy under aggressive quantization and benefit at lower bit widths, potentially allowing simpler quantization and low-precision training without specialized outlier-handling mechanisms.The discussion links this to eliminating attention sinks and suppressing extreme activations.Source: §8 · pp. 8–9
  • Sparse attention could enable inference acceleration with sparse kernels and reduce unnecessary attention-value computation; inactive-until-triggered heads may also be easier to identify and prune.The authors frame these as potential benefits, not demonstrated performance gains.Source: §8 · p. 9
  • More legible attention maps may benefit interpretability methods such as attention rollout, and softpick may generalize to vision, video, and multimodal transformers.The paper notes that attention sinks and activation outliers have also been observed beyond language models.Source: §8 · p. 9

9 Conclusion

Softpick is a rectified, non-sum-to-one drop-in replacement for softmax in transformer attention. Across 340M and 1.8B models, it eliminates attention sinks and suppresses activation outliers, but scaling and long-context behavior remain open challenges.

9 Conclusion

  • Softpick is a rectified, non-sum-to-one drop-in replacement for softmax in transformer attention.The conclusion presents it as an alternative with desirable properties for transformer language models.Source: §9 · p. 9
  • 0% sink rate across 340M and 1.8B models trained from scratch, alongside genuinely sparse attention maps and strongly suppressed massive activation outliers.The conclusion reports improved robustness under low-bit post-training quantization, especially at very low precisions.Source: §9 · p. 9
  • Scaling and long-context behavior remain open challenges, motivating further work on softmax alternatives for transformer language models.The conclusion reports 340M results as competitive with softmax.Source: §9 · p. 9

Limitations

The authors identify unreliable scaling under the current training recipe and no demonstrated improvement in long-context retrieval. They leave broader scaling studies for future work.

Limitations

  • Softpick is competitive with softmax at 340M but underperforms at 1.8B when the same training setup is reused.The authors hypothesize that underscoring at longer contexts or persistently inactive heads may contribute to the larger-model gap.Source: Limitations · p. 9
  • Softpick does not yet improve long-context retrieval in the reported experiments, and stronger value signals as context grows remain an open need.The authors cite passkey retrieval as an example and note that sparse attention maps do not by themselves ensure improved retrieval.Source: Limitations · p. 9
  • The fixed compute budget favored experimental breadth over exhaustive hyperparameter sweeps and substantially longer runs, leaving larger models, longer contexts, and training horizons for future study.This bounds the paper’s scaling evidence to its tested settings and training runs.Source: Limitations · p. 9

C Other Experiments

This section reports negative experiments with two Softpick-related variants, compares softmax and softpick attention maps on sample inputs, and catalogs full quantization comparisons for 340M and 1.8B models.

C Other Experiments

  • Two negative Softpick-variant experiments examine rectified-only softmax, which masks negative logits before normalization, and scalable-softpick, which adds sequence-length-dependent scaling.In the scalable variant, the implementation uses a vector assigning possibly different effective sequence lengths to token positions.Source: C Other Experiments · p. 15

D Full Results on Quantization of the 340M Models

  • Full 340M quantization comparisons cover BNB, GPTQ, and HQQ, with each table defining the delta as softpick minus softmax.The variants use the same setup and configuration as the main 340M models; training-loss and gradient-norm curves and downstream-task and passkey-retrieval results are also referenced.Source: C Other Experiments · p. 15, Figure 7 · p. 16, Table 8 · p. 17

E Full Results on Quantization of the 1.8B Models

  • Full 1.8B quantization comparisons cover BNB, GPTQ, and HQQ, with each table defining the delta as softpick minus softmax.The supplied table captions identify the comparisons but do not state their outcomes.Source: Table 11 · p. 20, Table 12 · p. 21, Table 13 · p. 22

F More Attention Maps

  • The attention-map comparison runs three sample text inputs through the models and extracts maps for softmax and softpick.The three inputs are a sequence-transduction passage, a passage about bees flying, and a passage about a phoenix character.Source: F More Attention Maps · p. 23

Related papers

Loading 2504.20966v4…