Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment¶

🕒 Published (v1): 2026-06-26 05:25 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

SMMD adapts kernel Maximum Mean Discrepancy (MMD) to supervise numeric token prediction in LLMs/VLMs by aligning predicted and target distributions in an RKHS using a value-distance kernel over numeric tokens, plus a graph-smoothness regularizer on the prediction-target residual. It consistently improves exact-match accuracy across mathematical reasoning, arithmetic calculation, clock-time recognition, and chart QA relative to cross-entropy and prior numeric-aware losses (NTL, NTIL, GCE).

Problem¶

Standard cross-entropy treats numeric tokens as unstructured categorical labels, ignoring ordinal and distance information — confusing "3" with "4" is penalized identically to "3" with "7". This objective mismatch makes LLMs unreliable when outputs must be numerically precise. Prior EMD-based losses (NTL, NTIL) incorporate value distance into the penalty but do not explicitly encourage local smoothness of the error profile, so per-token error signals can vary unevenly across neighboring numeric tokens.

Method¶

Numeric sub-vocabulary. Precompute $V_{\text{num}} \subseteq V$, the set of tokens parseable as real values ($N = 10$ for digit-level tokenization, up to $N=1000$ for multi-digit integer tokens). At training positions where the ground-truth token $y \in V_{\text{num}}$, restrict logits to $V_{\text{num}}$ and renormalize to obtain the numeric distribution $p = \text{softmax}(\ell[V_{\text{num}}]) \in \Delta^N$; the target is a one-hot $q = e_{\pi(y)}$.

Distance-induced kernel. Build an RBF kernel over pairwise value distances: $$K_{ij} = \frac{1}{|\Sigma|}\sum_{\sigma \in \Sigma} \exp\left(-\frac{|v_i - v_j|^2}{2\sigma^2}\right)$$ so numerically closer tokens have higher similarity. $K \in \mathbb{R}^{N \times N}$ is symmetric PSD.

Kernel MMD alignment. With residual $r = p - q$, squared MMD on the finite domain reduces to a quadratic form: $$L_{\text{MMD}} = r^\top K r$$ When $K$ is positive definite, $L_{\text{MMD}} = 0$ iff $p = q$.

Smoothness regularizer. Treat $V_{\text{num}}$ as a weighted graph with edge weights $K_{ij}$ and penalize the Dirichlet energy of the residual: $$L_{\text{smooth}} = \frac{1}{2}\sum_{i,j} K_{ij}(r_i - r_j)^2 = r^\top L r$$ where $L = D - K$ is the graph Laplacian. Smoothness is applied to the residual (not $p$) so it vanishes when $p=q$. The strength $\alpha$ is set automatically via degree-based normalization: $\alpha = 1/(2\bar{d})$.

Total objective. $L = L_{\text{CE}} + \lambda (L_{\text{MMD}} + \alpha L_{\text{smooth}})$, with default $\lambda=3.0$, $\sigma=2.0$. No architectural changes; the kernel is precomputed once per tokenizer.

Key Contributions¶

First to use kernel distribution matching (MMD) to supervise numeric token prediction in LLMs
Smoothness regularizer on the prediction–target residual via Dirichlet energy over the numeric kernel graph, encouraging locally coherent error profiles
Demonstrated complementary effects: MMD contributes most on language-only reasoning (GSM8K/SVAMP), smoothness contributes most on arithmetic and clock-time where local numeric consistency matters
Consistent gains across 4 task categories and 8 backbones (0.5B–11B parameters, both LLMs and VLMs)

Results¶

GSM8K: Qwen2.5-1.5B improves from 54.99% → 57.77% Acc, surpassing NTL (56.17%) and NTIL (55.19%)
SVAMP (cross-dataset, no additional training): Llama3-8B from 62.20% → 65.30% Acc
Arithmetic (DeepMind-Math): SMMD achieves the best Acc on all 4 backbones; SmolLM3-3B from 61.94% → 64.06% Acc with lowest MAE (1.88) and highest $R^2$ (0.82)
Clock-Time: Ministral-3-3B from 82.36% → 97.56% Acc, Time Gap from 37.36 → 8.07 min; Qwen2.5-VL-7B from 65.69% → 78.26% Acc
ChartQA: best or tied-best Acc on 3 of 4 backbones (Ministral-3: 71.64% → 73.21%; Qwen2.5-VL-7B: 77.02% → 78.38%)
Digit-distribution probe shows SMMD produces the sharpest, most target-aligned distribution over digit tokens compared to baselines

Limitations¶

ChartQA gains are modest — performance on this benchmark appears bottlenecked by visual grounding and value extraction rather than the numeric training signal
Single default ($\sigma$=2.0, $\lambda$=3.0) used across all experiments; optimal values vary per dataset (though sensitivity analysis shows stability)
On Clock-Time, adding MMD on top of smoothness alone slightly reduces Acc, suggesting a task-specific tradeoff where sharp bucket decisions are favored
Evaluation limited to open-weight models up to 11B; not tested on larger proprietary models
Only four task categories studied; real-world scientific/financial numerical prediction is not directly evaluated

Relevance to Vision-Language Models¶

Many VLM tasks — chart QA, clock reading, scientific figure interpretation — require precise numerical outputs grounded in visual inputs, yet standard CE ignores the metric structure of numbers. SMMD provides a lightweight, architecture-agnostic auxiliary loss that consistently improves numeric accuracy across diverse VLM backbones (Qwen2.5-VL, Ministral-3, Llama-3.2-Vision). The gains are orthogonal to visual improvements, meaning SMMD can be stacked on top of better visual grounding or representation-level numeracy enhancements.