Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment

2026-06-26

AI Summary

Large language models often underperform on numerically sensitive tasks because the standard cross-entropy loss treats numerical values as unstructured categorical tokens, disregarding their intrinsic metric relationships. To address this limitation, this work proposes Smooth MMD, a novel approach that integrates distance-aware kernel functions with graph-based smoothing mechanisms to align predicted and target distributions at the subword level. By leveraging residual smoothing over a kernel-induced graph, the method enhances local consistency among numerically proximate tokens. Evaluated across four diverse benchmarks—mathematical reasoning, arithmetic computation, clock reading, and chart-based question answering—Smooth MMD consistently outperforms cross-entropy and other state-of-the-art numerical loss functions, effectively mitigating the structural insensitivity inherent in conventional training objectives.

Abstract

Despite their strong general capabilities, large language models (LLMs) often remain unreliable when outputs must be numerically precise. A key reason is the training objective: standard cross-entropy treats numeric tokens as unstructured categories and ignores the metric structure of their values. We address this mismatch with Smooth Maximum Mean Discrepancy (SMMD), which builds on the classic MMD by incorporating value-distance kernels over numeric tokens and graph-based smoothness. With this kernel defined over a numeric sub-vocabulary, SMMD aligns the predicted numeric distribution to the target via kernel matching and smooths the prediction-target residual over the induced kernel graph to encourage local consistency. We evaluate SMMD on four numeric-target tasks: mathematical reasoning, arithmetic calculation, clock-time recognition, and chart question answering, across multiple open-weight LLM and VLM backbones. SMMD consistently improves accuracy over both cross-entropy and recent numeric-target losses; analyses show complementary effects between MMD and smoothness and underscore the importance of distance-based kernel design. Code is available at https://github.com/Zuozhuo/smmd-loss.

Problem

• numerical prediction
• large language models
• metric structure
• numeric tokens
• value-distance

Innovation

• Smooth MMD
• numerical prediction
• value-distance kernel
• distribution alignment
• LLM numerical reasoning