Research map / Large language models / Language, safety and interpretability
Regress, Don't Guess: A Regression-like Loss on Number Tokens for Language Models
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 Number Token Loss (NTL), a regression-like loss on number-token probabilities (NTL-MSE/Lp and NTL-WAS Wasserstein) added to cross-entropy, improving arithmetic and regression performance for LMs without runtime overhead.
Problem
Cross-entropy loss treats number tokens as nominal classes, so language models fail to convey numerical proximity between generated numbers, hurting arithmetic and quantitative reasoning.
Approach
Compute a regression-style loss directly on the token distribution over number tokens: either an Lp loss between the probability-weighted numeric value and the label's value (NTL-MSE/MAE), or a Wasserstein-1 distance over number tokens with Euclidean cost (NTL-WAS), added to CE with weight lambda; it is zero for non-number tokens and needs only a token-to-value map.
Key results
- On the DeepMind math interpolation set, T5-base with NTL-WAS reaches 0.75 accuracy vs 0.64 for CE, with lower MAE (0.10 vs 0.13).
- On rJokes regression, T5 with NTL-WAS reaches RMSE 1.81 vs 2.01 for CE, on par with a T5 regression head (1.82).
- T5-3B on GSM8K improves top-1 accuracy from 13.5% (CE) to 17.7% with NTL-WAS.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add NTL-WAS to CE | T5-Base · Mathematics Dataset interpolated test | Accuracy: 0.64 -> 0.75 | Helped |
| Add NTL-WAS to CE | T5-Base · Mathematics Dataset extrapolated test | Accuracy: 0.367 -> 0.432 | Helped |
| Add NTL-WAS to CE | T5-3B · GSM8K | Accuracy: 13.5\% -> 17.7\% | Helped |
| Add NTL-WAS to CE | GPT-2 Small (125M) · Multiplication interpolation | MAPE: 0.55 -> 0.49 | Helped |
| Add NTL to CE | GPT-2 Medium (350M) · Multiplication extrapolation | MAPE: 0.82 -> 0.82 | No clear effect |
| Add NTL-WAS to CE | T5 · rJokes test | RMSE: 2.01_{\pm 0.01} -> 1.81_{\pm 0.01} | Helped |
| Add NTL-WAS to CE | T5 · MultiRC validation | BLEU: 0.13$_{\pm 0.00}$ -> ${0.13}_{\pm 0.00}$ | No clear effect |
| Use NTL-MAE instead of CE (ablation) | T5-Small · Arithmetic Q&A interpolated test | Accuracy: $0.34_{\pm 0.01}$ -> $\textbf{0.45}_{\pm 0.02}$ | Helped |
| Use NTL-MSE with λ=2.0 instead of λ=0.15 (ablation) | T5-Small · Arithmetic Q&A interpolated test | Accuracy: $0.44 _{\pm 0.02}$ -> $0.33_{\pm 0.01}$ | Hurt |
| Use NTL-WAS with λ=0.3 instead of λ=2.0 (ablation) | T5-Small · Arithmetic Q&A extrapolated test | Accuracy: $ \textbf{0.10}_{\pm 0.01}$ -> $\textbf{0.10}_{\pm 0.01}$ | No clear effect |
| Combine GCE σ=0.5 with NTL-WAS (ablation) | T5-Small · Arithmetic Q&A interpolated test | Accuracy: $0.43$ -> $\mathbf{0.48}$ | Helped |
| Combine GCE σ=0.5 with NTL-WAS (ablation) | T5-Small · Arithmetic Q&A extrapolated test | MAE: $58.18$ -> $66.97$ | Hurt |
Surprising findings
- NTL-Huber and NTL-MAE sometimes match or beat NTL-WAS in the ablation studies.
- xVal performs very poorly on the DeepMind math dataset (accuracy 0.10), attributed to the large numeric range.
Limitations
- NTL-Lp has non-unique minima that can give spuriously low loss for incorrect predictions.
- Digit-level tokens are weighted equally, ignoring decimal place significance.
- Multi-digit tokenizers with large-valued tokens can destabilize the loss without squashing.
Experimental setup
Datasets: DeepMind mathematics dataset rJokes MultiRC GSM8K integer multiplication dataset
Benchmarks: DeepMind mathematics dataset (interpolation and extrapolation) rJokes MultiRC GSM8K 3-digit multiplication
Metrics: accuracy MAE R2 MAPE RMSE Pearson Spearman BLEU ROUGE-1 token accuracy
Baselines: Regression Transformer xVal Gaussian Cross Entropy regression head (MSE) BERT RoBERTa XLNet LLaMA-2 7B Gemma-2B
Builds on: cross-entropy Gaussian label smoothing Regression Transformer xVal T5
Closest papers by meaning
- Efficient numeracy in language models through single-token number embeddings (ICML 2026 · similarity 0.90)
- Learning In-context $n$-grams with Transformers: Sub-$n$-grams Are Near-Stationary Points (ICML 2025 · similarity 0.87)
- Number Cookbook: Number Understanding of Language Models and How to Improve It (ICLR 2025 · similarity 0.87)
- Tokens-per-Parameter Coverage Is Critical for Robust LLM Scaling Law Extrapolation (NeurIPS 2026 · similarity 0.87)
- Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning (NeurIPS 2025 · similarity 0.87)
- Universal Cross-Tokenizer Distillation via Approximate Likelihood Matching (NeurIPS 2025 · similarity 0.87)
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.