Research map / Learning theory and algorithms / Matrix and tensor methods
Adaptive Sampling for Efficient Softmax Approximation
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.
Proposes AdaptiveSoftmax, a bandit-based adaptive sampling algorithm that estimates the top softmax output of a matrix-vector product with PAC guarantees, reducing cost by subsampling features rather than classes.
Problem
Computing the full softmax of a high-dimensional matrix-vector product is expensive when only the top few softmax outputs are needed, e.g., in final LLM/classifier layers.
Approach
Reduce softmax computation to statistical estimation: adaptively subsample the inner (feature) dimension of Ax with importance-weighted estimators, using best-arm identification to find the top index and a two-stage sampling scheme (guided by estimated optimal frequencies alpha and gamma = exp(beta*mu/2)) to estimate the partition function with PAC (epsilon, delta) guarantees.
Key results
- On Wikitext with Mistral-7B, AdaptiveSoftmax gives 32.65x sample complexity gain at delta=10% (89.08% success rate).
- On MNIST with a shallow CNN, gains of 8.95x at delta=10% (92.25% success) and 8.13x at delta=1% (99.38% success).
- On synthetic data, average gains are 3.953x for the all-ones query and 29.039x for the sign query at delta=10%, epsilon=30%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AdaptiveSoftmax instead of brute-force softmax computation | VGG-19 classifier · EuroSAT test set | sample complexity gain: 5.18x (80.62%); 5.16x (83.00%); 4.54x (98.37%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | Shallow CNN · MNIST test set | sample complexity gain: 8.95x (92.25%); 8.81x (93.75%); 8.13x (99.38%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | GPT-2 · Wikitext | sample complexity gain: 8.25x (88.94%); 7.80x (93.54%); 6.67x (98.26%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | Llama3-7B · Wikitext | sample complexity gain: 14.68x (91.44%); 11.43x (94.04%); 6.88x (99.38%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | Mistral7B · Wikitext | sample complexity gain: 32.65x (89.08%); 26.37x (91.20%); 17.71x (97.77%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | GPT-2 · Penn Treebank | sample complexity gain: 8.10x (81.68%); 7.50x (90.73%); 6.66x (96.79%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | Llama3-7B · Penn Treebank | sample complexity gain: 19.18x (87.82%); 16.57x (91.60%); 10.72x (97.81%) | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | AdaptiveSoftmax · Synthetic all-ones query, δ = 10%, ε = 30% | sample complexity gain: 3.953× | Helped |
| Use AdaptiveSoftmax instead of brute-force softmax computation | AdaptiveSoftmax · Synthetic sign query, δ = 10%, ε = 30% | sample complexity gain: 29.039× | Helped |
Surprising findings
- The authors note that estimating the normalization constant dominates sample complexity, so top-k identification adds minimal cost.
- Aggressive constant tuning on training data occasionally yielded success rates below target.
Limitations
- Benefits are modest when the inner dimension is small, and adaptivity conflicts with batched BLAS, so wall-clock gains are not demonstrated.
- Requires a known sub-Gaussian variance proxy bound (Assumption 1) for theory; practical version uses empirical variance and tuned constants.
- In LLMs the final softmax layer is often not computationally dominant, so end-to-end acceleration is limited.
Experimental setup
Datasets: MNIST EuroSAT Wikitext Penn Treebank
Benchmarks: MNIST EuroSAT Wikitext Penn Treebank synthetic all-ones and sign queries
Metrics: sample complexity gain over naive nd success rate at epsilon=30% multiplicative accuracy
Baselines: naive full softmax computation
Builds on: Bandit-Based Monte Carlo Optimization multi-round epsilon-arm best-arm identification Gumbel top-k sampling without replacement sampled softmax
Closest papers by meaning
- Density-Softmax: Efficient Test-time Model for Uncertainty Estimation and Robustness under Distribution Shifts (ICML 2024 · similarity 0.84)
- Rethinking Approximate Gaussian Inference in Classification (NeurIPS 2025 · similarity 0.84)
- Statistical Perspective of Top-K Sparse Softmax Gating Mixture of Experts (ICLR 2024 · similarity 0.83)
- High-Probability Minimax Adaptive Estimation in Besov Spaces via Online-to-Batch (NeurIPS 2026 · similarity 0.83)
- Understanding the Training Speedup from Sampling with Approximate Losses (ICML 2024 · similarity 0.83)
- $\epsilon$-Softmax: Approximating One-Hot Vectors for Mitigating Label Noise (NeurIPS 2024 · 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.