Research map / Learning theory and algorithms / Prediction and decision losses
Softmax is not Enough (for Sharp Size Generalisation)
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proves that softmax attention coefficients must disperse (decay as Θ(1/n)) as the number of input items grows with a finite token vocabulary, and proposes an inference-time adaptive temperature that sharpens softmax on larger inputs.
Problem
Softmax attention cannot robustly stay sharp (select a constant number of items) when the number of input items at test time exceeds training sizes, causing out-of-distribution size generalisation failures in reasoning tasks.
Approach
Shows theoretically that bounded logits force softmax coefficients to decay as O(1/n) for every item, and proposes an entropy-based adaptive temperature (a fitted inverse degree-4 polynomial mapping entropy to temperature, capped at 1) that sharpens softmax at inference time, streamable like Flash Attention.
Key results
- On max retrieval averaged over ten seeds, adaptive temperature raises accuracy at 16,384 items from 12.4% to 14.0% and at 1,024 items from 53.8% to 57.7%, with no change in-distribution (98.6%).
- Empirical logit spread in Gemma 2B attention heads on a ~4,000-token code sample averages 5.69±2.05, far below worst-case bounds.
- Fine-tuning Gemma 2B with adaptive temperature gives a significant advantage on nearly all of the thirty CLRS-Text algorithms, with underperformance on Heapsort, MST Kruskal and Bubble Sort.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply adaptive temperature at inference without changing learned parameters | single-head attention model · max retrieval, size 16 | accuracy: 98.6% -> 98.6% | No clear effect |
| Apply adaptive temperature at inference without changing learned parameters | single-head attention model · max retrieval, size 32 | accuracy: 97.1% -> 97.1% | No clear effect |
| Apply adaptive temperature at inference without changing learned parameters | single-head attention model · max retrieval, sizes 64–16,384 | accuracy: 94.3% -> 94.5%; 89.7% -> 89.9%; 81.3% -> 82.1%; 70.1% -> 72.5%; 53.8% -> 57.7%; 35.7% -> 39.4%; 22.6% -> 24.9%; 15.7% -> | Helped |
| Use adaptive-temperature softmax during fine-tuning and inference | Gemma 2B fine-tuned on CLRS-Text · CLRS-Text algorithmic tasks | accuracy | Mixed |
| Apply the learned adaptive-temperature function only at inference (ablation) | Gemma 2B on CLRS-Text · CLRS-Text | accuracy | Hurt |
Surprising findings
- Applying adaptive temperature only at inference to a pretrained Gemma 2B did not work well on CLRS-Text, because numbers span multiple tokens; the model had to be fine-tuned with the adaptive softmax for gains.
Limitations
- Adaptive temperature is ad-hoc and, by the authors' own admission, does not escape the theory's conclusions.
- Theory assumes continuous building blocks and finite token vocabularies; worst-case bounds are loose compared with empirical logit spreads.
- Inference-time temperature does not transfer well to tokenised floating-point inputs, requiring fine-tuning for CLRS-Text gains.
Experimental setup
Datasets: CLRS-Text synthetic max retrieval dataset
Benchmarks: max retrieval (synthetic, up to 16,384 items) CLRS-Text
Metrics: accuracy entropy paired t-test p-value
Baselines: Baseline softmax attention (no adaptive temperature)
Builds on: Transformer Flash Attention CLRS-Text
Closest papers by meaning
- Think Smarter not Harder: Adaptive Reasoning with Inference Aware Optimization (ICML 2025 · similarity 0.84)
- LLMs are Single-threaded Reasoners: Demystifying the Working Mechanism of Soft Thinking (ICLR 2026 · similarity 0.84)
- From Noise to Diversity: Random Embedding Injection in LLM Reasoning (NeurIPS 2026 · similarity 0.84)
- The Deterministic Horizon: When Extended Reasoning Fails and Tool Delegation Becomes Necessary (ICML 2026 · similarity 0.83)
- Neural Network Approximators for Marginal MAP in Probabilistic Circuits (AAAI 2024 · similarity 0.83)
- On the Hardness of Approximating Distributions with Tractable Probabilistic Models (NeurIPS 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.