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

TL;DR

Softpick introduces a rectified, non-sum-to-one replacement for softmax in transformer attention to address attention sink and massive activations. It preserves much of softmax’s gradient behavior while enabling zeroed outputs and sparse attention, leading to 0% sink rate in tested models and markedly reduced activation outliers, with quantized Softpick outperforming softmax at low bit-widths. Across 340M and 1.8B parameter models, Softpick delivers competitive or improved downstream performance at 340M and shows robust quantization advantages, though scaling to 1.8B and long-context tasks remains a challenge due to underscoring and dead heads. The work suggests practical benefits for low-precision training, sparsity-driven efficiency, interpretability, and broader applicability to vision and multimodal transformers, while outlining avenues for future scaling and long-context improvements.

Abstract

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. Code: https://github.com/zaydzuhri/softpick-attention.

Figures (19)

• Figure 1: (Top) Comparison between the attention maps when using softmax vs softpick and overall sink rate of the 340M models. (Bottom) Largest hidden state activation per layer of the 340M models.
• Figure 2: Training loss and gradient norm during training of 340M models.
• Figure 3: Quantization results of softmax vs. softpick across model scales. We deliver results on 2, 3, 4, and 8-bit quantization. Full tabular results in Appendix \ref{['sec:quantization-appendix']} and \ref{['sec:quantization-appendix-1.8B']}.
• Figure 4: Box plots of the hidden state activations at some layers of the softmax (left, blue) and softpick (right, orange) models.
• Figure 5: Attention maps of softmax and softpick 340M models on 2 different input texts. Two heads are visualized: Head 1 of Layer 11 and Head 2 of Layer 21. See more attention maps in Appendix \ref{['sec:attention-maps-appendix']}.