REVIEW 5 major objections 5 minor 4 cited by
Scalable-Softmax Is Superior for Attention
T0 review · 5 major / 5 minor · reviewed 2026-08-09 · deepseek-v4-flash
Pith's one-line read The paper claims that replacing Softmax with SSMax—a one-line rescaling of attention logits by $\log n$—prevents attention fading, giving Transformers faster pretraining, better length generalization, and reliable key-information…
desk verdict SSMax is a simple, plausible attention tweak with consistent but under-supported experiments; the theoretical framing overreaches, but the idea deserves a referee. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the SSMax formula, $z_i \mapsto \frac{n^{s z_i}}{\sum_{j=1}^{n} n^{s z_j}}$, which is exactly Softmax applied to logits multiplied by $s\log n$; in an attention layer this is implemented by multiplying the query vector by $s\log n$ before the dot product. The argument that it works rests on a pair of bounds: for a fixed logit vector, the top Softmax output is at most $\frac{1}{(n-1)e^{z_{\min}-z_{\max}}+1}$, which decays to zero, while the top SSMax output lies between $\frac{1}{(n-1)n^{s(z_{\min}-z_{\max})}+1}$ and $\frac{1}{(n-1)n^{s(z_{\mathrm{2nd}}-z_{\max})}+1}$, so it saturates near 1 when $z_{\max}-z_{\mathrm{2nd}}>1/s$ and near 0 when $z_{\max}-z_{\min}<1/s$. A separate empirical fit, $p_n \approx a_1\log n + a_2$, measured on a 162M-parameter model, is what motivates the logarithmic scaling in the first place.
What would settle it
Run the same pretraining recipe with and without SSMax on 1,024-token sequences, then evaluate per-position test loss at context sizes up to 20,000 with RoPE $\theta$ scaled 50-fold; if the SSMax model does not stay below the standard model across contexts, the central claim fails. A more targeted check is to take a trained standard Transformer and measure the maximum attention probability for a fixed logit vector as $n$ grows; if that maximum does not trend to zero, attention fading is not the operative mechanism.
Extended reading notes
Core claim
The central claim is that attention fading is a real, trainable bottleneck: as the number of tokens $n$ in the context grows, the largest entry of a Softmax output shrinks toward zero because the denominator collects $n$ positive terms while the numerator stays bounded. The paper's proposed SSMax, defined by $z_i \mapsto \frac{n^{s z_i}}{\sum_{j=1}^n n^{s z_j}}$, replaces Softmax in every attention layer with a version whose logits are multiplied by $s\log n$ before the Softmax. This makes the attention distribution behave like a peaked selector: when the top logit exceeds the second-best by more than $1/s$, the maximum attention weight approaches 1 as $n$ grows; when all logits lie within $1/s$ of each other, attention stays diffuse. Experimentally the paper reports that SSMax-trained models achieve lower pretraining loss, retain low test loss at long contexts, and concentrate attention on key tokens in needle retrieval, whereas standard Softmax models fail on the same tasks.
Load-bearing premise
The whole design depends on the claim that the learned attention-scaling factor follows a logarithmic law in context length, a relation fitted on one 162M-parameter model and assumed to transfer to other sizes and architectures; if that law is not universal, the motivation for the specific $\log n$ form weakens.
Editorial extensions
If this is right
- A one-line implementation change—multiplying the query by $s\log n$—can be dropped into existing Transformer codebases without changing the architecture, loss, or data pipeline.
- SSMax-trained models should keep much of their long-context ability even when RoPE's $\theta$ is scaled 50-fold without further training, reducing the need for context-extension fine-tuning.
- Replacing Softmax with SSMax during the final phase of pretraining, or even after pretraining, gives partial length-generalization gains, so existing models can get some benefit without a full retrain.
- The learned scaling parameter $s$ controls the attention concentration threshold ($1/s$), giving model builders a direct, interpretable dial between peaked and diffuse attention.
- Omitting the bias parameter is important for long-context performance even though it slows pretraining slightly; the paper's comparisons suggest a clean design choice for production use.
Reading between the lines
- If the logarithmic scaling law is universal, one would expect that the optimal attention temperature for any Transformer shrinks like $1/\log n$, which suggests a direct relation between context length and entropy control that the paper does not explore.
- SSMax's threshold behavior implies it should help tasks that depend on a few key tokens (retrieval, factual recall) more than tasks that need broad attention; testing on summarization or translation could reveal a trade-off the paper leaves unmeasured.
- Because SSMax is a query-only rescaling, it can be composed with any positional encoding or attention mask; a natural extension is to check whether per-head learned $s$ values specialize, with some heads keeping near-uniform attention while others focus, rather than a single global temperature.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes replacing the Softmax function in Transformer attention with Scalable-Softmax (SSMax), defined as z_i -> n^{s z_i} / sum_j n^{s z_j} = Softmax((s log n) z_i). The authors argue that Softmax attention suffers from 'attention fading' as context size n grows, based on a fixed-logit upper bound, and claim that SSMax prevents this fading, improves pretraining efficiency, length generalization, and key-information retrieval. Experiments are conducted on a 162M-parameter Transformer pretrained on SlimPajama with sequence length 1024, using six configurations that vary whether and when SSMax replaces Softmax, and with or without a bias parameter. The reported results show lower training loss for SSMax variants, better per-position test loss up to context 20,000 after scaling RoPE's theta by 50x, higher needle-in-a-haystack retrieval accuracy at contexts up to roughly 10x the training length, and higher attention scores on key tokens.
Significance. If the findings hold, SSMax is an appealingly simple one-line change to attention normalization that could improve long-context behavior in Transformers. The paper has real strengths: the bounding argument in Section 2.2 is mathematically clean for a fixed logit vector; the evaluations cover multiple model variants, a length-generalization test, a retrieval benchmark, and an attention-allocation analysis, which is broader than a single loss-curve comparison; and the reported effects are directionally consistent across tasks. However, the evidence is limited to a single 162M-parameter model with no multiple seeds or error bars, no code or checkpoints, and a long-context evaluation protocol that confounds the attention modification with a drastic change in RoPE's theta. The central theoretical premise that trained attention logits are approximately fixed as context grows is not empirically verified, and no comparison is made to a simple length-dependent temperature baseline, which SSMax mathematically equals. The general claim that SSMax is 'superior for attention' across all Transformer-based LLMs is not yet established by the presented evidence.
major comments (5)
- [Section 2.2, Eq. (6)] The upper bound that motivates SSMax holds for a fixed logit vector z. In a trained Transformer, attention logits are not frozen as n grows; if the range z_max - z_min grows at least logarithmically with n, Softmax attention need not flatten at all. The paper does not report measured logit ranges, attention entropies, or attention concentration as a function of context length in the standard Softmax model, so the 'attention fading' premise is unverified. Because the paper generalizes to all Transformer-based LLMs and the entire mechanism rests on this premise, this is a load-bearing gap that needs to be addressed with direct measurements or a control experiment.
- [Sections 3.2 and 3.3] All long-context evaluations are performed after increasing RoPE's theta from 10,000 to 500,000 with no additional training. Figure 5 shows that the standard model (a) degrades even at short contexts after this change, so the comparison conflates the attention normalization with robustness to positional-encoding extrapolation. To isolate the effect of SSMax on length generalization, the authors should also report results at the original theta (or with a standard adaptation procedure such as NTK-aware scaling) and show the interaction between the attention change and the positional-encoding change.
- [Section 2.1, Figure 2] The log-law fit that motivates the SSMax functional form is not quantified: the fitted constants a1 and a2, the goodness of fit, and the run-to-run stability of p_n are not reported, and the fit comes from a single exploratory 162M-parameter model. Since the paper explicitly derives the design of SSMax from this fit, the missing quantitative support weakens the design rationale. At minimum, report the fitted values, R^2 or similar, and discuss whether the log law transfers to other model sizes and architectures.
- [Section 2.3, Eq. (11)] SSMax is exactly Softmax with logits multiplied by s log n, i.e., a length-dependent inverse temperature. The paper does not compare SSMax against a simple temperature-scaled Softmax baseline, such as Softmax((tau log n) z) with learned or fixed tau. Without this control, the observed improvements in long-context loss and retrieval could be due to temperature sharpening rather than to counteracting 'attention fading'. Adding such a baseline is essential to support the mechanistic claim and to distinguish SSMax from a temperature schedule.
- [Section 3] All quantitative results come from single runs of one 162M-parameter model, with no error bars, multiple seeds, or statistical significance tests, and no code or checkpoints are provided. Given the strength of the title and the conclusion that SSMax should become standard in Transformer-based LLMs, the evidence base is too narrow for the generality of the claim. The authors should provide at least a few random seeds (or an explicit statement of compute constraints), and ideally release code and model checkpoints to allow verification.
minor comments (5)
- [Abstract and Section 1] The statement that the maximum element of Softmax output 'approaches zero as the input vector size increases' should be qualified as 'for a fixed input vector z'; otherwise it reads as a general property of trained attention logits, which is precisely what is at issue.
- [Section 2.1] Please clarify the relationship between the shared p_n parameters in Eq. (3) and the per-layer/per-head learnable s in Eq. (2). If p_n is approximately a1 log n + a2, the constants a1 and a2 can be absorbed into the learnable s and b parameters, so it would be helpful to state this explicitly.
- [Section 3.2, Figure 5] The text reports per-position test loss on sequences of length 20,000 but does not state whether the reported loss is averaged over all positions or only over positions beyond 1024; please clarify, as this affects the interpretation of the gray dotted line.
- [Section 3.3, Figure 6] The figure caption mentions that colors indicate retrieval accuracy, but the text does not define the color scale; adding a colorbar or a textual description would improve readability.
- [Section 2.2] The footnote for s < 0 is terse; since s is a learnable parameter, it would be clearer to state explicitly that the threshold argument in Eqs. (8) and (9) holds for any nonzero s after replacing s with |s| and swapping the roles of z_max and z_min as appropriate.
Circularity Check
No significant circularity: the log n form of SSMax is empirically motivated from a learned per-length fit, but the length-generalization and retrieval results are independent evaluations, not consequences of that fit.
full rationale
Section 2.1 fits per-length scaling parameters p_n in a 162M-parameter model and observes p_n ≈ a1 log n + a2, then adopts the log n form in SSMax. This is an empirical design choice rather than a circular prediction: the fitted constants a1 and a2 are not used as targets in any later experiment, and the scaling parameter s in SSMax is independently learnable (or set to 1 in variant (c)). The bounds in Section 2.2 are mathematical properties of the proposed definition, not fitted results or self-citations, and they do not by themselves generate the empirical claims of lower pretraining loss, better long-context perplexity, or higher retrieval accuracy; those claims come from direct comparisons against the standard Softmax baseline on SlimPajama and SQuAD. Equation (11) shows that SSMax is exactly Softmax with logits multiplied by s log n, but this equivalence is a reparameterization, not a circular reduction of an evaluation result to an input. There are no load-bearing self-citations or imported uniqueness theorems. The reviewer's concern that Eq. (6) assumes fixed logits is a validity and evidence limitation, not a circularity, because the empirical comparisons are independent of that bound. Accordingly, no specific circular step can be exhibited.
Assumptions & free parameters
free parameters (4)
- s (SSMax scaling parameter) =
learned per layer and head, not reported
- a1 and a2 (log-fit constants) =
not reported
- p_n (per-length learned scaling in exploratory model) =
learned for n = 1..1024, not reported as a list
- RoPE theta multiplier in long-context evaluation =
50
assumptions (4)
- domain assumption For a fixed logit vector with bounded max-min gap, the maximum softmax output tends to 0 as n grows.
- ad hoc to paper The learned per-length scaling p_n in Section 2.1 is well approximated by a1 log n + a2, and this log law transfers to the SSMax design.
- domain assumption The 162M-parameter Llama-like architecture is representative of larger Transformer language models.
- standard math Standard Transformer components such as RoPE, RMSNorm, SwiGLU, and AdamW interact with SSMax in the same way as with Softmax.
Cite this review
Pith. "Pith review of Scalable-Softmax Is Superior for Attention." pith.science (2026). https://pith.science/paper/QQHBHC4H
@misc{pith2026250119399,
author = {Pith},
title = {Pith review of: Scalable-Softmax Is Superior for Attention},
year = {2026},
howpublished = {\url{https://pith.science/paper/QQHBHC4H}},
note = {Machine review of arXiv:2501.19399}
}
read the original abstract
The maximum element of the vector output by the Softmax function approaches zero as the input vector size increases. Transformer-based language models rely on Softmax to compute attention scores, causing the attention distribution to flatten as the context size grows. This reduces the model's ability to prioritize key information effectively and potentially limits its length generalization. To address this problem, we propose Scalable-Softmax (SSMax), which replaces Softmax in scenarios where the input vector size varies. SSMax can be seamlessly integrated into existing Transformer-based architectures. Experimental results in language modeling show that models using SSMax not only achieve faster loss reduction during pretraining but also significantly improve performance in long contexts and key information retrieval. Furthermore, an analysis of attention scores reveals that SSMax enables the model to focus attention on key information even in long contexts. Additionally, although models that use SSMax from the beginning of pretraining achieve better length generalization, those that have already started pretraining can still gain some of this ability by replacing Softmax in the attention layers with SSMax, either during or after pretraining.
Figures
Figures from the paper (3 more)
Forward citations
Cited by 4 Pith papers
-
Threshold Differential Attention for Sink-Free, Ultra-Sparse, and Non-Dispersive Language Modeling
TDA uses length-dependent row-wise thresholding plus differential subtraction to produce over 99% exact zeros in attention, eliminate sinks, and keep competitive performance on long-context tasks.
-
Critical attention scaling in long-context transformers
In a simplified attention model with normalized tokens, the phase boundary between token collapse and identity attention occurs when the attention-temperature scaling factor β_n is of order log n, with constant 1/(1−ρ).
-
AQUA: Attention via QUery mAgnitudes for Memory and Compute Efficient Inference in LLMs
A training-free method that prunes low-magnitude dimensions of projected query/key vectors in attention, cutting dot-product cost by 25% with small benchmark degradation.
-
On the Mathematical Impossibility of Safe Universal Approximators
This preprint argues that useful universal approximators must have dense catastrophic failures, but the proof rests on unproven measure convergence and circular use of adversarial examples.
Reference graph
Works this paper leans on
-
[1]
write newline
" write newline "" before.all 'output.state := FUNCTION n.dashify 't := "" t empty not t #1 #1 substring "-" = t #1 #2 substring "--" = not "--" * t #2 global.max substring 't := t #1 #1 substring "-" = "-" * t #2 global.max substring 't := while if t #1 #1 substring * t #2 global.max substring 't := if while FUNCTION format.date year duplicate empty "emp...
-
[2]
Etc: Encoding long and structured inputs in transformers
Ainslie, J., Ontanon, S., Alberti, C., Cvicek, V., Fisher, Z., Pham, P., Ravula, A., Sanghai, S., Wang, Q., and Yang, L. Etc: Encoding long and structured inputs in transformers. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp.\ 268--284, 2020
work page 2020
-
[3]
Needle in a haystack - pressure testing llms, 2023
Arize AI . Needle in a haystack - pressure testing llms, 2023. URL https://github.com/Arize-ai/LLMTest_NeedleInAHaystack2. Accessed on Jan 19, 2024
work page 2023
-
[4]
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020
arXiv 2004
-
[5]
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A.,...
work page 1901
-
[6]
Generating long sequences with sparse transformers
Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019
arXiv 1904
-
[7]
Redpajama: An open source recipe to reproduce llama training dataset, April 2023
Computer, T. Redpajama: An open source recipe to reproduce llama training dataset, April 2023. URL https://github.com/togethercomputer/RedPajama-Data
work page 2023
-
[8]
Gupta, A. and Berant, J. Gmat: Global memory augmentation for transformers. arXiv preprint arXiv:2006.03274, 2020
arXiv 2006
Show all 33 references
-
[9]
Needle in a haystack - pressure testing llms, 2023
Kamradt, G. Needle in a haystack - pressure testing llms, 2023. URL https://github.com/gkamradt/LLMTest_NeedleInAHaystack. Accessed on Jan 19, 2024
2023
-
[10]
The impact of positional encoding on length generalization in transformers
Kazemnejad, A., Padhi, I., Natesan Ramamurthy, K., Das, P., and Reddy, S. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36: 0 24892--24928, 2023
2023
-
[11]
Reformer: The efficient transformer
Kitaev, N., Kaiser, ., and Levskaya, A. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020
2001 arXiv
-
[12]
Gradient-based learning applied to document recognition
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86 0 (11): 0 2278--2324, 1998
1998
-
[13]
World model on million-length video and language with blockwise ringattention
Liu, H., Yan, W., Zaharia, M., and Abbeel, P. World model on million-length video and language with blockwise ringattention. arXiv preprint arXiv:2402.08268, 2024 a
2024 arXiv
-
[14]
Scaling laws of ro PE -based extrapolation
Liu, X., Yan, H., An, C., Qiu, X., and Lin, D. Scaling laws of ro PE -based extrapolation. The Twelfth International Conference on Learning Representations, 2024 b
2024
-
[15]
and Hutter, F
Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. International Conference on Learning Representations, 2019
2019
-
[16]
A., and Lewis, M
Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409, 2021
2021 arXiv
-
[17]
Language models are unsupervised multitask learners
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. OpenAI blog, 2019
2019
-
[18]
SQ u AD : 100,000+ questions for machine comprehension of text
Rajpurkar, P., Zhang, J., Lopyrev, K., and Liang, P. SQ u AD : 100,000+ questions for machine comprehension of text. Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp.\ 2383--2392, 2016
2016
-
[19]
Know what you don't know: Unanswerable questions for SQ u AD
Rajpurkar, P., Jia, R., and Liang, P. Know what you don't know: Unanswerable questions for SQ u AD . Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, 2: 0 784--789, 2018
2018
-
[20]
Ramachandran, P., Zoph, B., and Le, Q. V. Searching for activation functions. arXiv preprint arXiv:1710.05941, 2017
2017 arXiv
-
[21]
Efficient content-based sparse attention with routing transformers
Roy, A., Saffar, M., Vaswani, A., and Grangier, D. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9: 0 53--68, 2021
2021
-
[22]
Self-attention with relative position representations
Shaw, P., Uszkoreit, J., and Vaswani, A. Self-attention with relative position representations. North American Chapter of the Association for Computational Linguistics, pp.\ 464--468, 2018
2018
-
[23]
Glu variants improve transformer
Shazeer, N. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020
2002 arXiv
-
[24]
R., Hestness, J., and Dey, N
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, J. R., Hestness, J., and Dey, N. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama . https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama, June 2023. URL https...
2023
-
[25]
Roformer: Enhanced transformer with rotary position embedding
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 0 127063, 2024
2024
-
[26]
Adaptive attention span in transformers
Sukhbaatar, S., Grave, E., Bojanowski, P., and Joulin, A. Adaptive attention span in transformers. Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp.\ 331--335, 2019
2019
-
[27]
R., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D
Touvron, H., Martin, L., Stone, K. R., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D. M., Blecher, L., Cant \'o n Ferrer, C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami,...
2023 arXiv
-
[28]
N., Kaiser, ., and Polosukhin, I
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, ., and Polosukhin, I. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017
2017
-
[29]
Length generalization of causal transformers without position encoding
Wang, J., Ji, T., Wu, Y., Yan, H., Gui, T., Zhang, Q., Huang, X., and Wang, X. Length generalization of causal transformers without position encoding. arXiv preprint arXiv:2404.12224, 2024
2024 arXiv
-
[30]
V., and Zhou, D
Wei, J., Wang, X., Schuurmans, D., Bosma, M., ichter, b., Xia, F., Chi, E., Le, Q. V., and Zhou, D. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35: 0 24824--24837, 2022
2022
-
[31]
Differential transformer
Ye, T., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, F. Differential transformer. arXiv preprint arXiv:2410.05258, 2024
2024 arXiv
-
[32]
A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., and Ahmed, A
Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., and Ahmed, A. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33: 0 17283--17297, 2020
2020
-
[33]
and Sennrich, R
Zhang, B. and Sennrich, R. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019
2019
Reviewed August 9, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.