Underscoring the Problem: Why Softpick Fails at Initialization

Aryan Sood; Jaikaran Singh; Ishaan Bansal

cs.LGOct 4, 2026

Underscoring the Problem: Why Softpick Fails at Initialization

Authors: Aryan Sood, Jaikaran Singh, Ishaan Bansal

Organizations: IIT Roorkee

Abstract

Softmax attention gives every token a nonzero weight, which in trained models concentrates into attention sinks and massive activations that widen the dynamic range low-precision inference must cover. Softpick removes this constraint by rectifying scores, eliminating sinks and lowering hidden-state kurtosis, but its advantage fades at scale. We reframe this failure as a normalization problem. Softpick's denominator splits into positive- and negative-shifted sums $D^+$ and $D^-$ , used identically in the forward and backward pass, preventing their roles from being isolated. We separate them into a family of operators that independently choose each denominator. The failure originates at initialization: every layer contains rows where $D^+$ is exactly zero, while near-dead rows produce gradient norms above $10^{12}$ regardless of the backward denominator. Only Softpick and a stop-gradient variant, which keeps $D^+ + D^-$ forward but backpropagates through $D^+$ alone, train from scratch. At 230M parameters, the stop-gradient operator matches Softpick on quantization, has fewer dead heads, and retrieves passkeys more reliably, trailing only on peak attention-weight kurtosis.

Figures & tables

Operator	Forward $D_{f}$	Backward $D_{b}$	Role
Softpick	$D^{+}+D^{-}$	$D^{+}+D^{-}$	Coupled baseline
Stop-gradient	$D^{+}+D^{-}$	$D^{+}$ (stop-grad on $D^{-}$ )	The survivor
Rectified-forward	$D^{+}$	$D^{+}+D^{-}$	Honest gradient, rectified forward
Partial ( $\lambda$ )	$D^{+}+\lambda D^{-}$	matched	$\lambda\in\{0.11,0.3\}$
Floored	$\max(D^{+},\delta)$	$D^{+}+D^{-}$	Hard forward floor

Table 1: The operator family. Each row sets a forward and a backward denominator independently.

Panel of Figure 1 — Figure 1: (a) Softpick keeps one attention spike per query row (max 0.0625); the $D^{+}$ -only forward scatters spurious spikes at a higher peak (0.4333). (b) The stop-gradient operator ends 230M training with 23 dead heads of 238 against softpick’s 29.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Metric	ctx $\sim$ 1024	ctx $\sim$ 2048	ctx $\sim$ 3200
Exact match	0%	0%	0%
Target digits appear (substring)	0%	0%	0%
First generated number correct	0%	0%	0%
Digit recall (LCS / target length)	23.4%	32.7%	0.0%
Mean normalized edit distance	0.77	0.67	1.00
Generations containing any digit	32.0%	43.0%	0.0%

Appendix

Table 2: Digit-level analysis of softpick’s 230M generations,

n=500

per context length. All strict digit metrics are

0\%

; the low recall is chance overlap, and most generations contain no digits.

Condition	ctx $\sim$ 1024	ctx $\sim$ 2048	ctx $\sim$ 3200
abs (softpick baseline)	96.7%	93.3%	86.7%
relu ( $D^{+}$ only)	0%	0%	0%
relu, rescaled	0%	0%	0%
damped $c=0.01$	0%	0%	0%
damped $c=0.05$	0%	0%	0%
damped $c=0.10$	0%	0%	0%

Appendix

Table 3: Full inference-only denominator sweep, 340M checkpoint,

n=30

per cell.

Checkpoint	$c$	ctx $\sim$ 1024	ctx $\sim$ 2048	ctx $\sim$ 3200
Baseline	1.000	96.7%	93.3%	86.7%
step-500	0.309	100.0%	53.3%	40.0%
step-700	0.111	96.7%	90.0%	63.3%
step-900	0.012	83.3%	80.0%	30.0%
step-1000	0.000	0%	0%	0%

Appendix

Table 4: Passkey retrieval along the anneal, 340M. Step-700 is the last checkpoint before collapse.

	130M		230M
Metric	Softpick	Stop-grad	Softpick	Stop-grad
Baseline loss (fp32)	2.9647	2.9729	2.9091	2.8569
Mean hidden-state kurtosis	229.06	345.81	700	695
Attn.-weight kurtosis (mean 130M / peak layer-3 230M)	36,363.89	31,837.27	7,894,308	30,128,298
2-bit RTN loss delta	+5.81	+6.16	+6.32	+6.52
4-bit RTN loss delta	+2.42	+4.62	+6.29	+6.13

Appendix

Table 5: Softpick vs. stop-gradient operator, full matched-scale comparison. Single checkpoint per operator per scale, no seeds.

Condition	ctx $\sim$ 1024	ctx $\sim$ 2048	ctx $\sim$ 3200
Softpick, seed offset 0	0.0% (0/500)	0.0% (0/500)	0.0% (0/500)
Softpick, seed offset 100,000	0.0% (0/500)	0.2% (1/500)	0.0% (0/500)
Stop-gradient, seed offset 0	100.0% (500/500)	77.6% (388/500)	0.0% (0/500)

Appendix

Table 6: Passkey retrieval, 230M, exact match,

n{=}500

per cell unless noted, matched output-length budget.

Context	Softpick	Stop-gradient
$\sim$ 1024	99.8% (499/500)	85.6% (428/500)
$\sim$ 2048	65.8% (329/500)	55.2% (276/500)

Appendix

Table 7: Passkey retrieval, 130M, exact match,

n{=}500

per cell.

Operator	Forward	Backward	Outcome (from scratch)
Softpick	$D^{+}+D^{-}$	$D^{+}+D^{-}$	Trains stably to 28k+ steps
Stop-gradient	$D^{+}+D^{-}$	$D^{+}$	Trains stably to 28k+ steps
Rectified-forward	$D^{+}$	$D^{+}+D^{-}$	Loss flat $\sim$ 10.39, grad norm $10^{12}$ – $10^{14}$
Partial (0.11, 0.3)	$D^{+}+\lambda D^{-}$	matched	NaN before 1000 steps
Floored	$\max(D^{+},\delta)$	$D^{+}+D^{-}$	Learns, oscillates $\sim$ 4 nats, grad norm spikes
Matched-rectified	$D^{+}$	$D^{+}$	Not run (Section J )

Appendix

Table 8: Every operator trained from scratch, and how it failed. Softpick and the stop-gradient operator are the only two that train.

Figure 2: Loss (left) and gradient norm (right) over training at 230M, for softpick, the stop-gradient operator, and the floored operator. The floored operator learns instead of pinning at random-guess loss, then oscillates well above the two stable operators and spikes to nearly $6\times 10^{9}$ in gradient norm near step 19,000.

Figure 3: Attention-weight surface at layer 2, companion to Figure 1 . Softpick and the $D^{+}$ -only forward reach the same maximum weight, 0.1875, but the $D^{+}$ -only forward redistributes it across several spurious peaks.

Figure 4: Full training panel for the $c{:}1\to 0$ anneal at 340M: loss, gradient norm (log scale), learning rate, and the $c$ -schedule, all against step.

	130M	230M	340M
Parameters	130M	228.9M	340M
Hidden size	768	896	1024
Layers	12	17	24
Attention heads (total)	144 (12/layer)	238 (14/layer)	384 (16/layer)
Head dim	64	64	64
MLP	SwiGLU, ratio 4	SwiGLU, ratio 4	SwiGLU, ratio 4

Appendix

Table 9: Training configuration, read from our run configurations and training scripts and the publicly released 340M checkpoint.

Explore similar work

Pretraining Transformers with Quantized Softmax in Attention
Sep 27, 2026Shangzhen Zhu, Muyan Hu, Tomasz KozlowskiTransformer Attention Gumbel-Softmax Relaxation
University of Illinois Urbana-Champaign
EFQ-Softmax: Exp-Free Quantization for Softmax
Sep 9, 2026Haohui Han, Yuming Wan, Hongni Wang +4Gumbel-Softmax Relaxation Block Sparse Flash Attention
Xi'an Jiaotong University · Xi’an Jiaotong University · Huawei Technologies Co., Ltd +1
From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers
Sep 28, 2026Nafiseh HosseinpourFardi, Negar Alihadi, Mahmoudreza Babaei +2Post-Training Quantization Mixed-Precision Quantization
Sharif University of Technology, Tehran, Iran · GISMA University of Applied Sciences, Potsdam, Germany · BRAINS, Brandenburg Research Center for Applied Intelligent Systems, Potsdam, Germany +3

Sep 27, 2026, cs.LGJ/K move · Enter open · S save

Pretraining Transformers with Quantized Softmax in Attention

Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski

University of Illinois Urbana-Champaign

Transformer Attention Gumbel-Softmax Relaxation Generative Pretrained Transformers Transformer Architectures Matrix Multiplication Activation Functions

Open paper AskPDF arXiv

Low-precision Transformer systems increasingly quantize attention matrix multiplications, while softmax often remains at higher precision. During pretraining, an approximate softmax changes the gradients that train the model as well as its forward computation. We study this interaction with K-interval attention, which approximates the exponential using K+1 grid values. We vary per-row grid calibration, interpolation versus hard rounding, and the placement of a straight-through surrogate relative to normalization. We derive the corresponding backward rules, including calibration derivatives, and compare these choices in pretraining experiments matched on model, data, and optimizer. Detaching the row extrema leaves the forward computation unchanged but produces a delayed increase in validation loss. With hard rounding at K=4, min-max calibration and a pre-normalization surrogate incur a large loss gap; changing either choice substantially reduces it. At 124M parameters and 2.5B training tokens, fixed-window calibration with a post-normalization surrogate yields a validation loss gap of +0.019 nats relative to softmax at K=4, and with a pre-normalization surrogate yields +0.004 nats at K=16.

Abstract

Explore similar work

Pretraining Transformers with Quantized Softmax in Attention

EFQ-Softmax: Exp-Free Quantization for Softmax

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers