Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment
Zhuo Zuo, Li Yue, Wenhao Zheng, Chenpeng Wang, Xianggen Liu
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 crossentropy 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 valuedistance kernels over numeric tokens and graphbased 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 crossentropy 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 .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on6
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- PAL: Program-aided Language ModelsLuyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon et al.ICML 2023 · 700 citations
- FoNE: Precise Single-Token Number Embeddings via Fourier FeaturesTianyi Zhou, Deqing Fu, Mahdi Soltanolkotabi, Robin Jia et al.ICLR 2026 · 24 citations
- Injecting Numerical Reasoning Skills into Language ModelsMor Geva, Ankit Gupta, Jonathan BerantACL 2020 · 12 citations
- Teaching Metric Distance to Discrete Autoregressive Language ModelsJiwan Chung, Saejin Kim, Yongrae Jo, Jaewoo Park et al.ICLR 2026 · 5 citations
Related papers
- Regress, Don't Guess: A Regression-like Loss on Number Tokens for Language ModelsJonas Zausinger, Lars Pennig, Anamarija Kozina, Sean Sdahl et al.ICML 2025
- Enhancing Numerical Prediction of MLLMS With Soft LabelingPei Wang, Zhaowei Cai, Hao Yang, Davide Modolo et al.ICCV 2025 · 1 citation
- RNED: Rotary Number Encoding and Decoding for Quantitative Medical VLM AnalysisFengbei Liu, Sunwoo Kwak, Nusrat Binta Nizam, Ilan Richter et al.CVPR 2026
- GeoNum: Bridging Numerical Continuity and Language Semantics via Geometric EmbeddingShengkai Jin, Tianyu Chen, Chonghan Gao, Jun HanAAAI 2026
- Methods for Numeracy-Preserving Word EmbeddingsDhanasekar Sundararaman, Shijing Si, Vivek Subramanian, Guoyin Wang et al.EMNLP 2020 · 28 citations
