Language models improve numeric data understanding with two stage tuning

MoMaTune: Model, Match, Then Tune Language Models for Numerical Records

Artificial Intelligence

Summary

Language models often struggle to understand numbers properly when they appear in data records. The authors found that usual training methods don’t teach models about the importance of digits or how different numbers in the same record relate to each other. They developed MoMaTune, a two-step training approach that first teaches the model to accurately predict missing numbers and keep internal consistency, then fine-tunes it for specific tasks. This method made models better at filling in missing numbers and understanding digit changes across many datasets.

What this means in practice

  • For financial data teams: Improve data-driven models that fill in or predict missing numeric financial record values with better accuracy and numerical consistency.
  • For health informatics engineers: Build more reliable systems for handling patient numerical records by training models to better understand digit-level changes and data dependencies.

Authors

Weimin Lyu, Jingfen Zhang, Lingdao Sha, Peng sun, Xiaowei Chen, Vincent Gao, Rui Song

Abstract

Despite remarkable progress in language understanding and generation, decoder language models continue to exhibit striking failures on numerical data. Numerical examples are typically serialized as text by concatenating feature names with their values (e.g., fat:3.4), while conventional task adaptation applies loss only to the downstream prediction target. This leaves two structures weakly supervised: within-value sensitivity, where digit position determines numerical magnitude, because token cross-entropy contains no explicit numerical metric; and across-field structure, where measurements in the same record are statistically dependent, because target-only tuning provides no direct loss on the measurements or their dependencies. We introduce MoMaTune (Model, Match, then Tune), a two-stage post-training strategy organized around these supervision gaps. Stage 1 applies Model, ordinary next-token prediction densely over label-free serialized records, and optionally adds Match, a same-record consistency objective that aligns representations of different partial views of the same example. Stage 2 applies Tune, supervised fine-tuning for downstream target prediction. We evaluate MoMaTune on 169,285 records from eight public datasets, with 339 of 340 input fields numerical. At 1.5B parameters, Model raises the eight-dataset masked-value imputation aggregate from 0.106 to 0.401, while the three-dataset correct-context gain rises from 0.058 to 0.167. On diagnostics, Model also increases sensitivity to numerically consequential digit changes, while Match improves robustness to informative partial views. Adding Tune yields higher target-supervised embedding aggregates for each tested Stage-1 parent, and the two-stage recipes obtain higher aggregate scores than Tune alone. The same pattern is observed at 4B, 9.6B, and 27B parameters when model embeddings are used for downstream prediction.