Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment

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 by proposing 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.

Large language models often struggle when answers must be numerically precise. A key reason is that standard training treats numbers like unrelated word labels: predicting “4” instead of “3” is penalized much like predicting “9,” even though “4” is numerically closer. We introduce SMMD, a training method that helps language models respect numerical distance. It compares the model’s predicted number distribution with the correct answer using a similarity structure over numeric tokens, and encourages nearby numbers to behave consistently. SMMD requires no change to the model architecture and can be added during fine-tuning. Experiments on math reasoning, arithmetic calculation, clock-time recognition, and chart question answering show that SMMD improves numerical accuracy over standard training and recent number-aware methods. This can help make language models more reliable in settings where numerical mistakes matter.