Regress, Don't Guess: A Regression-like Loss on Number Tokens for Language Models
Authors: Jonas Zausinger
Lars Pennig
Anamarija Kozina
Sean Sdahl
Julian Sikora
Adrian Dendorfer
Timofey Kuznetsov
Mohamad Hagog
Nina Wiedemann
Kacper Chlodny
Vincent Limbach
Anna Ketteler
Thorben Prein
Vishwa Mohan Singh
Michael Danziger
Jannis Born
Venue: ICML 2025 poster (2025)
Keywords: language models
mathematical reasoning
arithmetics
number representation
number token loss
Primary area: deep_learning->large_language_models
TL;DR: The Number Token Loss (NTL) is a regression-like loss on number tokens that augments cross-entropy to improve Language Models in numerical tasks.

Abstract: While language models have exceptional capabilities at text generation, they lack a natural inductive bias for emitting numbers and thus struggle in tasks involving quantitative reasoning, especially arithmetic. One fundamental limitation is the nature of the cross-entropy (CE) loss, which assumes a nominal scale and thus cannot convey proximity between generated number tokens. In response, we here present a regression-like loss that operates purely on token level. Our proposed Number Token Loss (NTL) comes in two flavors and minimizes either the  $\mathcal{L}_p$ norm or the Wasserstein distance between the numerical values of the real and predicted number tokens. NTL can easily be added to any language model and extend the CE objective during training without runtime overhead. We evaluate the proposed scheme on various mathematical datasets and find that it consistently improves performance in math-related tasks. In a direct comparison on a regression task, we find that NTL can match the performance of a regression head, despite operating on token level. Finally, we scale NTL up to 3B parameter models and observe improved performance, demonstrating its potential for seamless integration into LLMs. 
We hope to inspire LLM developers to improve their pretraining objectives and distribute NTL as a minimalistic and lightweight PyPI package *ntloss*: https://ibm.biz/ntl-pypi-repo. Development code for full paper reproduction is available separately.

Regress, Don't Guess -
A Regression-like Loss on Number Tokens for Language Models
Jonas Zausinger" 12 Lars Pennig" 1234 Anamarija Kozina 12 Sean Sdahl 12
Julian Sikora"
Adrian Dendorfer 12 Timofey Kuznetsov 12 Mohamad Hagog 15 Nina Wiedemann
Kacper Chlodny 12
Vincent Limbach 12
Anna Ketteler 12 Thorben Prein 12 Vishwa Mohan Singh" 15
Michael Danziger?
Jannis Born?
Abstract
Language Model
Text tokons
Number tokena
While language models have exceptional capabil- ities at text generation, they lack natural induc- tive bias for emitting numbers and thus struggle in tasks involving quantitative reasoning, espe- cially arithmetic. One fundamental limitation the nature of the cross-entropy (CE) loss, which assumes nominal scale and thus cannot convey proximity between generated number tokens. In response, here present regression-like loss that operates purely on token level. Our proposed Number Token Loss (NTL) comes in two flavors and minimizes either the Tp norm or the Wasser- stein distance between the numerical values of the real and predicted number tokens. NTL can easily be added to any language model and extend the CE objective during training without runtime overhead. We evaluate the proposed scheme on various mathematical datasets and find that it con- sistently improves performance in math-related tasks. In a direct comparison on a regression task, we find that NTL can match the performance of a regression head, despite operating on token level. Finally, scale NTL up 3B parameter mod- els and observe improved performance, demon- strating its potential for seamless integration into LLMs. We hope to inspire LLM developers to improve their pretraining objectives and distribute NTL as a minimalistic and lightweight PyPI pack- age nososs: https://github.. hutp:://tthub.om//ai//dunuu com/ai4sd/number- token-loss. Development code for full paper re- production is available separately.
Lies Lyg Lann
Head
Logits
trials
Vicah
Vocah
Option
L WIILWAS
0.0
Option 2
0
0.1
0.1 0.551
0.442
W
2
MSE(I, 1.3)
(a) The two flavors of the proposed Number Token Loss (NTL).
Loss as a function predicted number token
Cross Entropy
Number Token Loss
Loss
Correct token: 4
3
8 9
Predicted value
(b) Comparison of NTL to CE.
Figure I. (a) Number Token Loss (NTL). The NTL allows com- puting regression like loss directly a language modeling head. We propose two schemes: LNTLASE uses a dot product of number token values and their probabilities; CxTL was the Wasserstein- distance between label and prediction distributions. cominal-scalle and thus assigns equal loss all incorrect predic- tions whereas NTL increases with distance from ground truth like regression loss.
Equal contribution aUMM. "Technical University of Munich "Helmholtz Munich "Munich Center for Machine Learning (MCML) 'LMU Munich "ETH Zurich 'IBM Research Europe. Correspondence to: Jannis Born <jab@ www.ii.ib..cooo..........
(b) CE
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

NTL The Number Token Loss
1. Introduction
In this paper, we aim equip LMs with better inductive biases to handle combinations of textual and numerical data, such as math word problems or scientific datasets. We pro- pose two versions of regression loss on number tokens that respect numerical proximity and can be effectively com- bined with standard CE (Figure Ia). The first version of the number token loss (NTL) computes the Mean Squared Error (MSE) between the numerical value of the label and the predicted class probabilities weighted by their numerical token value. The second version computes the Wasserstein distance between the distribution of the predicted number probabilities and the ground truth distribution, which is the one-hot encoding of the label.
As coined by Thawani et al. (2021), numbers in natural texts are ubiquitous and important, yet systematically ne- glected by language models (LMs). Even worse, while Transformers (Vaswani et al., 2017) were invented for NLP, they have permeated various scientific domains (Jumper et al.. 2021; Boiko et al., 2023), where tabular/numerical data is more prevalent than in NLP and often governs task definitions: Molecules are labeled with drug efficacy, chem- ical reactions with yield, and synthesis procedures are natu- ral text interspersed with quantities and times. Prior art for joint language- number modeling suggested the use of veri- fiers (Cobbe et al., 2021; Li et al., 2023), chain-of-thought (CoT) reasoning (Zhong et al., 2024; Wei et al., 2022; Lee & Kim, 2023), calculators or code interpreters (Qu et al., 2025; Gao et al., 2023) to yield improved performance in LMs. On a model-level, customized ideas such as reverting num- bers (Zhang-Li et al., 2024), right-to-left tokenization (Singh & Strouse, 2024; Lee al., 2023), joint prediction of multiple number tokens have been proposed (Bachmann & Nagarajan, 2024). Still, LMs notoriously struggle even with simple arithmetic tasks like three-digit multiplication (Dziri et al., 2024). We argue that all these strategies - reformu- lating the task, trying to correct answers a posteriori with calculators, or using significantly more compute (CoR) avoid the fundamental, underlying problem: number repre- sentation in LMs is poor, for multiple reasons:
In both cases, the NTL. has the following key characteris- tics. (1) NTL is model agnostic. It can be applied to any language model (Transformer, Mamba, etc.) in any style (encoder-decoder, decoder-only etc). (2) NTL is plug-and- play for any LM because it makes minimal assumptions about the vocabulary. It only requires a map between to- kens (strings) and their numerical value (float) and is thus compatible with digit-level as well as multi-digit tokeniza- tion'. (3) NTL does not add computational overhead. The most capable NTL version (NTL-WAS) slows down the loss calculation only by 1% compared to cross-entropy. This difference becomes negligible over the whole training step (see later in Figure 3).
Our empirical analysis further shows that: (4) NTL con- sistently improves performance in mathematical tasks across architectures, compared to cross-entropy alone. (5) On real regression task, a LM with NTL performs on- par with a regression head and improves by 10% over a LM with standard CE. (6) NTL does not hamper per- formance on text-only tasks and (7) NTL scales well to billion-size models. Therefore nothing speaks against using it during LLM pretraining.
Standard subword tokenization splits numbers into arbitrary tokens, disrupting their structure and losing a significant amount of numerical informa- tion (Wallace et al., 2019). Mitigation strategies include scientific notation (Zhang et al., 2020) or digit-level tok- enization (Geva et al., 2020), which may also preserve the decimal order of each digit (Born & Manica, 2023). 2. Embedding: Canonically, LMs have to recover the struc- ture of numbers from data because number token em- beddings are learned like any other. To mitigate that, many flavors of numeracy-pres prreserryyy word embeddings exist (Sundararaman et al., 2020; Born& & Manica, 2023; Golkar et al., 2023), often akin to positional encodings and sometimes adapted for special cases like angular em- beddings for modular arithmetics (Stevens et al., 2024; Saxena et al., 2024).
1. Tokenization:
2. Number Token Loss (NTL)
Inevitably, some tokens of any LM correspond to digits or numbers, not to text. Our basic idea is to combine the log- its of these number tokens with their numerical values to compute a loss. Such a loss can effectively augment cross- entropy by taking numerical proximity of number tokens into account. Note that outside NLP, some systematic stud- ies of "regression by classification" exist (Shah & Aamodt, 2023; Stewart et al., 2023), yet a principled approach that facilitates seamless integration into LLMs lacks entirely.
3. Sequential prediction: Token-wise decoding fails to account for the greater significance of higher-order digits in a number. Zhang-Li al. (2024) find that reversing the digit order and human-inspired, step-by-step calculations for addition and multiplication improve accuracy.
digit-level, e.g.: L.aMA2 (Touvron et al., 2023), DeepSeek- V2 (DeepSeck-Al et al., 2024) and PaI.M (Chowdhery et al., 2022); multi-digit, e.g. L.aMA3 (Grattafiori et al., 2024) or TS (Raffel et al., 2020).
4. Training objective: Cross-entropy (CE) loss assumes a nominal scale, thus it fails to convey the proximity between numbers, effectively inducing a semi-supervised setting. For example, predicting a instead of a 2 will not generally induce lower loss than a 9.
2

---

NTL The Number Token Loss
NTL with Cp norm (e.g., NTL-MSE) This loss maps the predicted token probabilities to a real-valued output by calculating the dot product of the probabilities and their corresponding numerical token value (cf. Figure la). This transformation enables the application of standard regres- sion losses, such as the MSE, to compare the predicted output against the numerical value of the ground truth token. Assume model f()), imput tokens x (i S NN and a vo- cabulary V. Let y; f(xg:) be the predicted probability distribution, y; the one-hot label with yı as the numerical value of the ground truth token and V V + IR map to convert tokens (strings) to their numerical values (floats), with the index range s...t representing the number tokens. Then we compute NTL-MSE:
by defining the cost function e as the Euclidean distance between the numerical token values:
min (PM >ETly"" 9191)
P V (3)
j= A=
Note on Cost Function: Equation 3 is very flexible by changing the cost function c, NTL can be applied even when the number tokens lie in arbitrary, non -on-Eucli spaces. The cost function does not need to adhere to the definition of distance metric and can instead be specified as pairwise matrix between number tokens. NTL could thus find utility also in exotic cases such as modular arith- metics (Charton, 2023; Kera et al., 2024; Gromov, 2023). In this paper, we mainly assume standard Euclidean distance due its general applicability, A practical example where non-euclidean cost is useful are multi-digit tokenizers which often contain individual tokens for very large number (Section 4.5 contains explicit results on non-euclidean cost).
1 N 1(u Sy" Dostyz
CNIL-MSE
(1)
Notably, instead of a nominal-scale loss like CE, the NTL- MSE effectively conveys proximity between numbers. For example, if the label is and the LM predicts 5 instead of 9, the loss will be lower, matching our intuitive expec- tation. This is in contrast to the CE which gives constant loss regardless of the proximity of the number due to its nominal nature (ef. Figure lb). By changing the p-order in NTL-MSE, different -nom losses can be obtained (e.g., NTL-MAE). Huber loss is also compatible. While NTL. describes family of losses, we note that NTL-MSE has been proposed already in concurrent work by Lukasik al. (2025) where it is called RAFT and has shown improved performance of LLMs on pure regression tasks with up to five number digits. However, NTL-Lp has a non-unique minimum and can thus return spuriously low loss for incor- rect predictions. Consider e.g., a sample with label where 50% of the mass is on 0 and 50% on token 8, then NTL will be zero (see Figure 2). While such cases are rare due to the softmax emphasizing logit differences, combining NTL with CE loss helps correct spurious cases. But to address the non-unique minima more systematically we propose a second, refined version.
Computing the quantity in Equation 3 requires solving a linear program and is thus not generally differentiable. Ap- proximations could be made for the general case via en- tropic regularization (Cuturi, 2013), but our application of We entails two special cases that enable its efficient and differentiable computation: First, if y, is one-hot, CVIL-WAS coincides with the weighted sum of the absolute differences from the label and the numerical token values:
N
1 N
ANTL-WAS =
vllu 17
(4)
1
Second, if the number token indices s...t are sorted by nu- merical value you and the values are equidistantly spaced, the Wasserstein-1-dies 1-distance can be computed using the cu- mulative distribution function CDF(-):
N Σ 2m1
LNTL N CDF (y.*t) - CDF (s:t)) (5)
The advantage of LNTI. CAS is that it does not require the label y; to be one-hot instead it supports arbitrary target distributions obtained e.g.. via label smoothing on number tokens (Wang et al., 2025). In practice, unless mentioned otherwise, we compute LNTL.WAS via Equation 4 because it is 230x faster than Equation 5 and does not re- quire sorting or equidistant spacing and is thus applicable beyond digit-level tokenization. As one see in Figure Ib, LNTL not only conveys proximity between numbers cor- rectly but also eliminates the problem of non-unique minima in CNIL-MSE (Figure 2). Both versions of the NTL are scaled with À (default 0.3) and added to the standard CE loss:
NTL with Wasserstein-1 Distance (NTL-WAS) To mea- sure the similarity of predicted probability distribution of number tokens and a ground truth distribution, we take inspiration from Optimal Transport (OT) and leverage the discrete W distance, generally defined as:
We (P.QQ)) min ΣΣ kyclu - 0z)
TEr(P.Q)
where i; and tj are points, P and Q their associated proba- bilities respectively, and ca function specifying the transport costs. 7 is a coupling between P and Q where Yij specifies the mass to be transported from Un to Uy. We apply OT to measure the difference between the distributions y; and y;
L -= -= + XLNTL.
(6)
Note that both versions of the NTL yield zero loss for all non-numerical tokens (see Appendix A.2 for pseudo-code).
3

---

NTL - The Number Token Loss
Intuition
1.0
Cross Entropy
1.0
NTL.MEE
1.0
NTLWWAS
ing; Wang et al. (2025)) which addresses the limitations of standard CE for continuous-valued tokens. Instead of one- hot encoded labels, Gaussian distribution is centered on the true numeric label. In this scheme, nearby numeric tokens receive partial probability mass, reflecting the continuous nature of numerical values. Given a numeric label yi its one-hot vector yi is replaced with a Gaussian distribution:
Token (5)
0.5
0.5
0.5
0.5
1.0
0.0
0.5
1.0
0.0
9
0
0.0
0.0
0.9
Taken [3]
Taken (3]
Token [3]
Teken (3)
Figure 2. The heatmap plot shows the respective loss for a given combination of the class probabilities for tokens 3 and 5, where the ground truth is token 1. The behavior of the NTL-WAS is closest to the intuitive desired behavior of the loss function, while the NTL-MSE does not have a unique minimum.
J where a governs the smoothing width. The objective re- mains the cross-entropy, but now between the modified labels y, and the predicted distribution y;:
exp
(V' (V--))") 20
Vj E (s..)
(7)
V 2T 0²
3. Experiment Setup
N Σ N
1
LOCE
logy/.
(8)
3.1. Backbone T5 and Model Variants
We use a T5 backbone (Raffel et al., 2020) (Appendix A.3) for most of our experiments, unless stated otherwise, due to its flexible encoder various natural language processing tasks. We built upon the T5 implementation and language modeling trainer from t s formers (Wolf et al., 2020). For the multitask mathematics experiments in Section 4.1 we use T5-Base with 220M parameters. For the scaling experiment in Sec- tion 4.7 we use T5-3B. For ablation studies and other exper- iments, we use T5-Small with 60M parameters. To illustrate the model-agnostic nature of NTL, we evaluate it on both GPT-2 and IBM Granite model variants. We always initial- ize the models with their pre-trained weights. Building up on prior work (Geva et al., 2020), we adapt the tokenization scheme to tokenize all numbers on the digit level, while leaving all other tokens unchanged. Although both NTL versions are compatible with any tokenizer, single-digit tok- enization improves performance for both CE and NTL on mathematical tasks in our experiments (see subsection 4.5).
The obvious disadvantage is that the smoothing is entirely fabricated: there no real uncertainty in the label. While this technically provides incorrect labels, it ensures that the inherent similarity among nearby numeric tokens is cap- tured and continuity in numeric space is preserved. Wang et al. (2025) finds that this does not alter classification per- formance on strictly categorical tokens.
architecture and its success in
4. Experimental Results
Across all experiments. we always optimize NTL jointly with CE (cf. Equation 6; default À = 0.3). For brevity refer to this as just NTL-MSE or NAL.
4.1. NTL Improves Performance in Arithmetics
To test the mathematical capabilities of our models, more than 25 million samples from the mathematical Q&A dataset from DeepMind (Saxton et al., 2019). The dataset comes with two sets of tests: interpolation tests, one for each type of question occurring in the training set, and extrapola- tion tests, which measure generalization along various axes of difficulty beyond that seen during training. We provide more information about the dataset in Appendix 1.
3.2. Baselines
We compare to three related methods.
First, the Regression Transformer (Born & Manica, 2023) which tokenizes numbers digit level (considering both the position and value of each digit) and leverages a fixed number embedding that preserves relative proximity of the numbers (details in Appendix A.4).
4.1.1. MULTITASK MATHEMATICS DATASET
We evaluated all models on the two test sets of this dataset and report the accuracy (the proportion of exactly predicted numbers), as well as the Mean Absolute Error (MAE) and the R²-score. Since the dataset skewed with some very high values, we perform a log1o transformation on the pre- dicted and ground truth numbers before calculating MAE and R--score. All models were trained for ~ 1M steps with batch size of 32. More details on the models' training hyperparameters can be found in Appendix A.6.
Second, Val (Golkar et al., 2023) which encodes real num- bers using a single [NUM] token multiplied by their numer- ical value. For decoding, number/regression head predicts the value while the token head outputs the sequence, re- placing [NUM] during inference. However, this scheme is incompatible with T5 (see Appendix A.5) and problematic for beam search decoding. We thus use the Xal encoder and masked language modeling in our experiments.
The results in Table and Figure Al show that vanilla T5 clearly benefits from both NTL variants. In fact, the accu-
Third, Gaussian Cross Entropy (or Gaussian label smooth-
4

---

NTL The Number Token Loss
MAE and NTL-Huber. Training of these models was done subtask of the original dataset (100k samples), consist- ing of basic integer arithmetic Q&A pairs, where the input is short natural language question containing an arithmetic expression, and the output is a single integer. A comparison of their performances on the interpolation and extrapolation test sets is shown in Table 2.
Table Evaluation on Mathematics Dataset
Table I. (a) Interpolated Test Data
Model
Loss
Accuracy
MAE
R²
TS
CE
0.64
0.13
0.97
T5
NTL-MSE
0.72
0.11
0.97
T5
NTL-WAS
0.75
0.10
0.98
T5
NTL-MAE
0.74
0.10
0.97
Table 2. Evaluation NTL Loss Variants
RT
CE
0.71
0.11
0.97
XVal
MSE
0.10
0.26
0.97
Table 2. (a) Interpolated Test Data
Table I. (b) Extrapolated Test Data
Loss
À
Accuracy
MAE
R²
CE
0.34
2. 2.15
0.95
Model
Loss
Accuracy
MAE
R²
0.15
0.44-u.u2
0.92+0.05
0.99
T5
CE
0.367
0.785
0.913
0.3
0.41
0.99 1.02
T5
NTL-MSE
0.428
0.779
0.909
NTL-MSE
0.8
0.37
1.29
0.99++0.0
T5
NTL-WAS
0.432
0.744
0.913
2.0
0.33
1.67 200.0
0.97
T5
NTL-MAE
0.427
0.792
0.906
RT
CE
0.404
0.987
0.738
0.15
0.44
0.001
0.99
XVal
MSE
0.058
0.826
0.819
0.3
0.43
0.91
0.99
NTL-WAS
0.8
0.44
0.94
0.99
2.0
0.41
1.01
0.99
racy increases by more than 10% for NTL-WAS compared to vanilla T5 in the interpolation tasks. The NTL-WAS was found to have the best performance across all three metrics and both interpolation and extrapolation tasks. Moreover, NTL consistently outperforms competing methods such as the RT and x Val. This confirms our hypothesis that number representation LMs can be effectively improved through a minor, architectur--anostic loss modification.
NTL-Huber
0.3
0 0.44
0.89 +0.03
1.00
NTL-MAE
0.3
0.45
0.89
0.99
Table (b) Extrapolated Test Data
Loss
X
Accuracy
MAE
R²
CE
0.0520.00
61.91
0.61
0.15
0.09
57.13 1.51
0.68
Since x Val is equipped with regression head, numbers never predicted exactly, the accuracy value for x Val refers to the predictions rounded to two decimal places. The poor performance of Val can likely be explained by the extensive range of numbers in the used dataset. The effec- tive number range of x Val is limited to [-5,5] 5] due to the combination of its scaling of the number token embeddings and the pre-layer-norm in the backbone (Xiong al., 2020; Golkar et al., 2023). To take this into account, we scale our dataset for Val with log + 2²). However, this means that large numbers can no longer be adequately distinguished by the model, as their embeddings become very similar. For a direct comparison without any modifications to the x xVal processing, we repeated the 3-digit multiplication experi- ment from Golkar al. (2023). Again, our method beats xVal. Details can be found in Appendix A.5.
0.3
60.99 60.99+1.35 0.09
0. 0.65
NTL-MSE
0.8
0.08
58.35 .06
0.68
2.0
0.07
19.48422.7
0.66
0.15
57.31 2.09
0.68
0.3
0.10.0.01
58.18 89
0.68
NTL-WAS
0.8
0.10
59, 59.46+1.55
0.66
2.0
0.10 0.10
62.29
0.64
NTI--Huber
0.3
0.10=0.01
48.81
0.67
NTL-MAE
0.3
0.10
57.99+1.48
0.67
First, it evident that all flavors of NTL are consistently better than a standard CE loss. Mean and standard deviation were calculated over different training runs. Details and additional metrics can be found in Appendix A.6.2. The ablation studies in Table 2 support our finding that both variants of the NTL improve arithmetic performance com- pared to using the CE loss alone. Additionally, NTL-WAS mostly outperforms NTL-MSE across both interpolation and extrapolation test sets. Among the tested weights for the NTI., 0.3 yields the best performance for both loss
4.1.2. ABLATION STUDIES
In order to investigate the impact of variations in loss func- tions, we conducted extensive experiments testing nine dif- ferent configurations: standard cross entropy loss, NTL- MSE and NTL-WAS (both with À E {0.3, 0..,0..,220 2)), NTL.
5

---

NTL The Number Token Loss
variants. Interestingly, find that NTL-Huber achieves competitive and NTL-MAE sometimes even superior results to NTL-WAS, Across all experiments, the R² achieves high values, seemingly in contrast to the accuracy. A closer in- vestigation of the error distribution revealed that this can be attributed to many predictions being numerically very close to the ground truth, despite not being strictly accurate. These smaller errors have a very limited impact on R², while still affecting accuracy.
4.2. Training with NTL Is Efficient
To assess potential computational overhead, we performed comprehensive benchmarking across four distinct loss con- figurations: CE loss serving as the baseline, CE with NTL- MSE, CE with NTL-Wasserstein implemented according to Equation (NTL-WAS-CDE), and CE with its optimized variant (NTL- -ASS Equation 4). We evaluated these loss functions under two different scenarios to quantify both the standalone computational overhead of the loss calculation and its impact within a complete training step, including the forward pass, backpropagation and optimizer step. We standardized the number of tokens in the input for both for- ward passes and training steps, along with maintaining a consistent proportion of number tokens across all experi- ments, although their positions were randomized. Each loss function was evaluated through 100 iterations on a GPU, The results in Figure 3 demonstrate that NTL-MSE and NTL-WAS introduce only marginal computational overhead across both benchmark scenarios, as each benchmark in- cludes the runtime of the CE loss combined with the runtime of the respective NTL variant.
4.1.3. GAUSSIAN CROSS ENTROPY
The Gaussian Cross Entropy (GCE) proposed by Wang et al. (2025) constitutes an alternative, yet complimentary means to convey proximity between number tokens. We argue that is less principled because it essentially blurs the labels and assigns mass to technically incorrect tokens. Yet, assessed its performance by replacing the vanilla CE objective with GCE objective. We found a 0.5 to yield the best results (Table A6).
The experiments on both the interpolation and extrapola- tion sets suggest that while both NTL and GCE improve model performance, NTL is generally preferable, as it leads to larger improvements. Furthermore, we investigated the
Standalone
Loss NTL-WAS Loss NTL-WAS-CDF CE Losa + NTL-MSE
10
20
25
30
Table 3. Gaussian Cross Entropy (GCE). Standalone and com- binatory effect GCE and NTL. Means and standard deviations across four runs.
Time (ms)
Training step
CE Loss + NTL-WAS Loas NTL-WAS-CDF
GCE NTL Accuracy MAE
R2
Loss NTL-MSE GE Loss
Interpolation
50
100
150
200
250
300
Time (ms)
X
0.34
2.15
0.95
X
0.43
0.91
0.99
Figure 3. Runtime Comparison. Benchmarking of the four loss functions in two contigurations: standalone execution (top) and complete training step (bottom), with a mean number token proportion of 80 percent, highlighting computational overheads.
V
X
0.42
0.95
0.99
0.48
0.76
0.99
Extrapolation
X
X
0.05
61.92
0.61
X
V
0.10
58.18
0.68
Notably, computing NTL-WAS alone is around 125x faster than computing CE alone. This is achieved because only tiny fraction of tokens in the vocab are number tokens. When assessing only the loss computation time, augmenting the CE loss with NTL introduces computational overhead ranging from a 1% ENTL-WAST over 2.9% (NTL-MSE) and up substantial 286% for the less efficient NTL-WAS- CDF. However, during the full training step, the relative overhead becomes negligible as the overall computationa effort consisting of the forward pass, backpropagation and optimizer step becomes dominant. For instance, the relative overhead of NTL-WAS LLL--WAS-DD compared to the CE decreases to 4.4%.
V
0.10
58.55
0.65
V
0.10
66.67
0.59
performance of using NTL and GCE in combination. Since the labels are longer one-hot encodings in this setting, we employed the NTL-WAS- WFF formulation (Fquation 5) which allows for training with arbitrary target distributions. On the interpolation set, the combination of NTL and GCE yields the best results. These experiments suggest that NTL and GCE can be mutually beneficial for certain tasks.
6

---

NTL The Number Token Loss
4.3. NTL Can Match Regression Models
et al., 2023). This tasks consists of multiplying two num- bers with & and I digits, with k I E [1,.-,5] in training and k. E [I,.., 6] in testing, We report the mean absolute percentage error (MAPE) separately for unseen interpola- tion samples, which involve multiplying numbers with up to 5 x 5 digits, and extrapolation samples, which involve multiplications with at least one six-digit factor. The re- sults presented in Table 5 demonstrate the effectiveness of training decoder-only models with NTL,, as it consistently improves performance on the multiplication task in terms of MAPE
Ideally, LMs should be able to solve even tasks that are focused solely on predicting numerical values, such timating property of a molecule. Such tasks are often approached with regression heads on LMs which allow to use appropriate loss functions such as the MSE. To test the ability of NTL to solve a regression task solely boken-based, we train on the rJokes dataset (Weller & Seppi, 2020) that poses the challenge to predict the level of humor in a joke. A train-test split is provided with the dataset.
In Table 4, we show the results of the T5 model with CE, NTL-WAS and combined with a regression head, as well as the results from the rJokes benchmark ( (Weller & Seppi, 2020). We evaluated all flavors of the T5 model over three independent runs to ensure robust metrics. NTL signifi- cantly improves over standard CE and performs on-par with the regression-head model. This is remarkable considering that it operates on token-level and can still be utilized for non-numeric tasks, unlike models that rely on a dedi- cated regression head. Even though BERT, RoBERTA, and XLNet with a regression head outperform T5 trained with NTL-WAA the performance gap is not significant consider- ing the difference the number of parameters. We attribute this to the fact that these models have nearly twice as many parameters as T5-small model. Details and additional metrics can be found in Appendix A.6.3.
Interpolation CE
Extrapolation
Model
Size
NTL
CE
NTL
GPT2 Small
125M
0.55
0.49
1.11
1.00
GPT2 Medium
350M
0.43
0.42
0.82
0.82
GPT2 Large
774M
0.39
0.37
0.76
0.75
GPT2 XI.
1.5B
0.43
0.40
0.83
0.82
Granite 3.2
2B
0.35
0.21
0.60
0.42
Granite 3.1 MOE
1B
0.28
0.15
0.68
0.23
Table S. NTL Performance for Decoder-only Models. MAPE for training GPT-2 and IBM Granite variants with CE and NTL on the multiplication task, All values are error percentages.
The positive impact of NTL is evident across all model scales. Notably, the improvements from NTL are more pro- nounced in the extrapolation setting than in interpolation, indicating that training with NTL enhances model's ability to generalize beyond the training distribution. This is partic- ularly relevant due to the poor extrapolation capabilities in mathematical tasks (Razeghi et al., 2022). Some models, such as Granite 3.1, benefit particularly from training with NTL, showing the largest performance gains among the evaluated models (full evaluation in Appendix Figure A2).
Table NTL Performance on a Regression Task. NTL performs on-par with regression-head model. Mcan and standard devia- tions across five runs. BERT, ROBERTA and XLNNet values are taken from Weller& & Seppi (2020).
Model
Loss
Reg, Head
RMSE
Pearson
TS
CE
X
2.01 0.41+0.00
TS
NTL-WAS
X
1..0
0.40
TS
MSE
V
1.82+0.01
0.45+0.00
0.0
0.4
0.6
0.6
0.7
0.7
0.0
0.4
0.5
0.6
0.5
0.7
BERT
MSE
1.62
0.47
3.9
3.9
4.0
4
2.6
2.7
2.8
3.2
3.0
ROBERTA
MSE
<
1.61
0.47
4.1
5.0
4.8
5.1
3
2.8
3.2
3.5
3.7
3
XLNet
MSE
1.74
0.46
0.6
4.1
4.8
4.6
4.8
5.0
2
0.5
2.9
3.5
3.3
3.6
3.8
2
4.9
3.2
3.6
1
4.4. NTL Is Model-Agnostic
8
0.7
4.0
4.9
5.3
4.9
5.1
in
0.6
3.2
3.7
3.5
4.4
3.6
As a loss function, NTL. can be applied to train arbitrary models, including LMs beyond Transformers like Mamba. The experiments described so far were conducted with the T5 model, which is based on an encoder-decoder Trans- former architecture. In this section, we demonstrate the ef- fectiveness of training the decoder-only models GPT-2 (Rad- ford et al., 2019) and IBM Granite (IBM Research, 2024) with NTL. We test model sizes from 125M to 2B.
Number digits
Number digits
(a) Cross entropy
(b) NTL
Figure 4. Sample Effficienc.. Number of epachs needed reach MAPE helow 0.5 on the multiplication task (GPT2 Small). NTL. particularly improves sample efficiency for samples with longer factors.
Furthermore, we analyzed the sample efficiency of training with NTL on the multiplication task compared to using CE only. As revealed by Figure 4, GPT2 Small trained with
To this end, we constructed an arithmetic multiplication task, much akin to the length generalization task (Jelassi
7

---

NTL The Number Token Loss
NTL requires significantly fewer epochs to achieve a MAPE below 0.5. This effect is more pronounced for challenging samples consisting of larger factors.
single-digit tokenizer performs slightly better without NTL- WAS. Since the extrapolation task and the rlokes task more closely resemble real-world scenarios, experiments further reinforce the overall benefits of the NTL, particularly when combined with single-digit tokenization.
argue that these
Similarly, we observe that the average number output dis- tributions are more sharply centered on the correct token when using NTL (see Appendix Figure A4).
Furthermore, we observed a practical issue with models that tokenize numbers into multiple digits: some tokens have large numerical values, which disproportionately affect the loss, especially when the number tokens are irregularly spaced. For digit-level tokenization and NTL with euclidean distance (as used thus far) the loss for the farthest incorrect number token cannot exceed 9x of the loss for the nearest incorrect number token". But a single, large number token in the vocabulary, say 1001, can distort this ratio 1000x and may even destabilize training. NTL can intrinsically mitigate such edge cases because it is not limited us- ing Euclidean distances between numbers (cf. Equation 2). This can be accomplished through an optional rescaling fac- tor that squashes the distances between tokens. Note that, with digit-level tokenization, the implicit squashing factor is while setting to effectively recovers standard cross- entropy. Thus squashing allows smoothly interpolate between vanilla NTL and CE. On the integer multiplication dataset (Jelassi et al., 2023), we found that GPT2- with a squashing factor of 3 improves over CE but performs worse than NTL as measured by MAPE (Table A7). Overall, we recommend to enforce digit-level tokenization as this ensures NTL well-behaved.
4.5. NTL Is Effective for Different Tokenizations
Our standard implementation of NTL-WAS for the T5 model relies on a custom single-digit tokenizer derived from the standard T5 tokenizer, which is based on SentenceP i ece and includes some tokens consisting of multiple-digits. We conducted additional ablation studies to disentangle the effect of the custom tokenizer and the loss function. To this end, we performed further experiments on the rJokes and the mathematical Q&& datasets, in which we also evaluated the effect of using the single-digit tok- enizer without the NTL-WAS and a more general version of NTL-WAS that supports multi-digit tokens. The results on the rJokes dataset in Table 6 confirm the effectiveness the NTL-WAS loss across both tokenization schemes, with the combination of single-digit tokenization and NTL-WAS performing best.
Table 6. Tokenizer Comparison on rJokes Datset. T5-small used across all experiments. Mean and standard deviation across four runs.
Loss
Custom Tokenizer
RMSE
Pearson
4.6. NTL Does Not Hamper Text I Learning
To ensure that training with NTL does not introduce draw- backs for normal text generation tasks, we conduct exper- iments on the MultiRC dataset (Khashabi et al., 2018), a question-answering dataset testing reading comprehension. We reformat the dataset to train the model to generate the answer based on the question, rather than evaluating it as a multiple-choiee task. To this end, we preprocess the dataset by extracting passages along with their corresponding ques- tions and concatenating them to construct the input. The answer field consists of all correct answers concatenated with a delimiter. The NTL was only applied to the number tokens in the answers, which account for only 1.5% of all answer tokens propotion we believe to be roughly rep- resentative of real-world text datasets. In Table 8 the mean and standard deviation over three runs evaluated on the vali- dation set are shown. Token accuracy as well as BLEU and ROUGE-1 score remain unaffected by the addition of the NTL and remain on the same level as when only using CE, These results demonstrate that training with NTL does not hamper performance on text comprehension in task where
CE
x
2.01 10.01
1.400
WTL-WAS
x
1.97 0.41
CE
V
2.02
0.41
NTL-WAS
1.81 1.81 10.01
0.44
MSE
X
1.82+0.01
0.45
A similar trend is observed in the extrapolation test set of the mathematical Q&A task in Table 7, where the NTL- WAS improves accuracy for both tokenization strategies. In the interpolation test set in Table A5, the NTL-WAS
Table 7. Extrapolation Performance for Different Models and Tokenizers
Custom Tokenizer
Accuracy (Extrapolate)
Pearson (Extrapolate)
Loss
CE
X
0.05
0.81
NTL-WAS
X
0.06
0.76
CE
V
0.09
0.87
NTL-WAS
V
0.10+0.01
0.88
Forr ground truth of 0, the nearest/farthest incorrect tokens and 9.
enhances performance for the standard tokenizer, while the
8

---

NTL The Number Token Loss
experiments unambiguously demonstrate the effectiveness of the NTL for encoder-decoder and decoder-only architec- tures. Experiments several datasets related to mathemat- ics, arithmetics and reasoning showed significant improve- ments in numerical reasoning, especially in models without specialized numerical embeddings. Furthermore, our re- sults on regression datasets revealed that training with NTL allows LM with a voken-head to compete with models de- signed specifically for regression tasks. We also verified that NTL does not compromise capabilities on standard text-only tasks like reading comprehension. Finally, the scalability of NTL was demonstrated by successfully applying it to 3B- 3--aarameter model on the challenging GSM8k dataset, opening the door for integration into LLMs.
Table NTL on a Standard Language Modeling Task. NTL does not impact performance on a pure textual task.
Loss
Token Acc.
BLEU
ROUGE-1
CE
0.36+0.00
0. 0.13
0.000
NTL-WAS
0.36+0.00
0.13+0.00
0.32
its benefits for numerical data are limited or absent. These findings underline that NTL can be integrated as standard enhancement in LM training, without compromising general text understanding and generation.
4.7. NTL Scales Well to LLM-Size
From the two proposed NTL flavors (NTL-WAS and NTL- Ep) the NTL-WAS is preferable loss since it avoids the undesired local minima of NTI-Cpp caused by errors can- celing out. However, unlike in NTL-WAS, the dot product inside NTL-Cp yields a quantity in the original number space which can be useful to compute more complex losses incorporating multiple tokens. For example, a current lim- itation of NTL is treat each digit-level token as equally important, thereby ignoring the higher relevance of digits with higher decimal places. This limitation could be ad- dressed by computing the loss at the number level rather than at the digit level, e.g., by naive scaling by decimal place or by leveraging the dot-product idea of Cp to calculate a number-level prediction. Furthermore, future research should focus on exploring synergies with addi- tional arithnetii-specifie training adaptations. Particularly exciting will be the fine-tuning of mid- and large-scale open- source models, such as LLaMA or Gemma, to further assess its scalability and generalization potential. Overall, this approach offers a practical solution for enhancing language models in numerically rich domains, contributing towards more accurate and reliable applications of LMs in mathe- matics and science.
To demonstrate the potential to integrate NTL into LLMs, we train the 3B parameter version of the T5 model on the GSM8K dataset (Cobbe et al., 2021), a benchmark for math- ematical reasoning based on school math problems. This dataset poses significant challenge for standard LLMs; for instance, 6B- parammeterr models initially achieved only 22% accuracy after fine-tuning (Cobbe et al., 2021). In our exper- iments, the T5-3B model trained on GSM8K with CE loss attains a top- l-acuurccy of 13.5% on GSM8K (see Table 9). Training with the NTL loss improves accuracy to 17.7%, matching the performance of LLMs that are more suitable for mathematical reasoning, such as Gemma-2B (Team et al., 2024), and surpassing larger models such as LLaMA-2 7B (14.6%). A qualitative example of GSM8K task and pre- dicted solutions is provided in Appendix Example 1.
Table 9. Results with T5-3B on GSM8K
Loss
Accuracy
Pearson
CE
13.5%
0.67
NTL-WAS
17.7%
0.72
Further Resources
Notably, specialized training approaches can significantly boost performance. For example, tailored fine-tuning strate- gies have enabled even a 1.3B model to reach 81.5% ac- curacy (Liu et al., 2023). Future work should explore how NTL can further enhance such specialized models.
Code
1. Usage: NTL ships via PyPI. For minimalistic integration in your work use github.com/AIII htt../////AAsSmmmmbbreeoootee
2. Reproduction: Full development code for reproduction of all experiments:
5. Conclusion
We introduced the Number Token Loss (NTL) for Language Models as novel approach enhance their intrinsic ability to handle numerical data by considering the numerical prox- imity between tokens. By augmenting the standard cross entropy loss with NTL, we provide a simple yet effective method that integrates seamlessly into existing architectures without requiring additional computational overhead. Our
Learn more
1. Landing page: Project overview site
2. Demo: Streamlit app on HuggingFace spaces
3. Video: 5min paper walkthrough on YouTube
9

---

NTL The Number Token Loss
Impact Statement
References
The introduction of the NTL enhances the ability of LMs handle numerical data and can improve performance in tasks involving quantitative reasoning and arithmetic. This advancement has potential implications across fields like scientific research, finance, and engineering, where precise numerical outputs are essential for data-driven decision- making and automated analysis. However, this capability also raises potential risks, including misuse of generated numerical data or biased predictions sensitive areas, such as healthcare economics. As with any AI methods, care- ful implementation and human oversight are necessary to ensure model transparency and accuracy.
Bachmann, G. and Nagarajan, V. The pitfalls of next-token prediction, 2024. URL
2403.06963
Boiko, A., MacKnight, R., Kline, B., and Gomes, tonomous chemical research with large language models. Nature, 624(7992):570-578 2023.
Born, J. and Manica, M. Regression transformer enables concurrent sequence regression and generation for molec- ular language modelling. Nature Machine Intelligence, (4):432-444, 2023.
Charton, F. Learning the greatest common divisor: explain- ing transformer predictions. In The 3rd Workshop Mathematical Reasoning and AI at NeurlPS'23, 2023.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A.. Sepassi, R.. Dohan, D., Agrawal, S.. Omernick, M., Dai, A. M., Pillai, T. S.. Pellat, M., Lewkowyez, A., Moreira, E., Child, R., Polo- Zov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J.. Meier-Hellsterre K., Eck, D., Dean, J., Petrov, S., and Fiedel, N. Palm: Scaling language modeling with pathways, 2022. URL https: / www.a..vvvrrrooo abs 2204 1.0440400 02311.
Cobbe, K.,, Kosaraju, V., Bavarian, M., Chen, M.. Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 168, 2021.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Dai, A. M. and Le, Q. V. Semi- Semi-supervised sequence learning. Advances in newral information processing systems, 28, 2015.
DeepSeek-AN AI, Liu, A., Feng, B., Wang, B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D., Guo. D., Yang, D., Chen, D., Ji, D., Li, E., Lin, F., Luo, F., Hao, G., Chen, G., Li, G., Zhang, H., Xu, H., Yang, H., Zhang, H., Ding, H., Xin, H., Gao, H., and et al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024.
10

---

NTL The Number Token Loss
Dziri, N., Lu, X., Sclar, M., Li, X. L.,, Jiang, L.,, Lin, B. Y., Welleck, S., West, P., Bhagavatula, C.,, Le Bras, R., al. Faith and fate: Limits of transformers compositionality. Advances in Neural Information Processing Systems, 36, 2024.
Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pp. 252-262, 2018.
Lee, N., Sreenivasan, K., Lee, J. D., Lee, K., and Papail- iopoulos, D. Teaching arithmetic to small transform- ers, 2023, URL https: /atxiviivvvvv/////v/ org/ 2337 03381.
Gao, L., Madaan, A., Zhou, S., Alon, U., Liu, P., Yang. Y., Callan, J., and Neubig, G. Pal: Program-aided lan- guage models. In International Conference on Machine Learning, pp. 10764-10799. PMLR, 2023.
Lee, S. and Kim, G. Recursion of thought: A divide-and- conquer approach to multi-context reasoning with lan- guage models. arXiv preprint (rXiv:23060068// 2023.
Geva, M., Gupta, A., and Berant, J. Injecting numerical reasoning skills into language models. In Proceedings of the 58th Annual Meeting of the Association for Com- putational Linguistics. Association for Computational Linguistics, 2020.
Li, Y., Lin, Z., Zhang. S., Fu, Q., Chen, B., Lou, J.-G., and Chen, W. Making language models better reasoners with step-aware verifier. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 5315-5333, 2023.
Golkar, S., Pettee, M., Eickenberg, M., Bietti, A., Cranmer, M., Krawezik, G., Lanusse, F., McCabe, M., Ohana, R., Parker, L., et al. xval: A continuous number encoding for large language models. arXiv preprint 2000v22310.02989, 2023.
Liu, B., Bubeck. S., Eldan, R., Kulkarni. J., Li, Y., Nguyen, A., Ward, R., and Zhang, Y. Tinygsm: achieving; 80% on gsm8k with small language models. arXiv preprint arXiv:2312.0924 2023.
Grattafiori, A., Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Vaughan, A., Yang, A., Fan, A., and et al. The Llama 3 herd of models, 2024.
Lukasik, M., Meng. Z., Narasimhan, H., Chang, Y.W.,, Menon, A. K., Yu, F., and Kumar, S. Better autoregressive regression with Ilms via regression-aware fine-tuning. In The Thirteenth International Conference Learning Representations, 2025.
Gromov, A. Grokking modular arithmetic. arXiv preprint arXiv:2301.02679 2023.
IBM Research. Granite 3.0 language models, 2024.
Qu, C., Dai, S., Wei, X., Cai, H., Wang, S., Yin, D., Xu, J., and Wen, J.-R. Tool learning with large language models: A survey. Frontiers of Computer Science, 1988)43,333,, 2025.
Jelassi, S., d' Ascolii S., Komingo-Enrichh C., Wu, Y., Li, Y., and Charton, F. Length generalization in arithmetic trans- formers, 2023. URL https: // /arxivvvvvv./////////ax///aaa/l/ org/ abs/ 2306. 15000.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. Technical report, OpenAI, 2019. Technical report.
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Židek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583 2021.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W. and Liu, J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (1 (140):1-67, 2020.
Kenton, J. D. M.-W. C. and Toutanova, L. K. Bert: Pre- training of deep bidirectional transformers for language understanding. In Proceedings of naL. volume 1, pp. 2019.
Razeghi, Y., IV, R. L. L.., Gardner, M., and Singh, S. Impact of pretraining term frequencies on few- shot numerical reasoning. In Goldberg, Y.. Kozareva, Z. and Zhang, Y. (eds.), Findings of the Association for Computational Linguisties: EMNLP 2022, Abu Dhabi, United Arab Emirates, December 7-11, 2022, pp. 840-854. Association for Computational Linguis- tics, 2022. doi: 10.
Kera, H., Ishihara, Y., Kambe, Y., Vaccon, T., and Yokoyama, K. Learning to compute gröbner bases. Ad- vances in Neural Information Processing Systems, 37: 2024.
Khashabi, D.,. Chaturvedi, S., Roth, M., Upadhyay, S., and Roth, D. Looking beyond the surface: challenge set for reading comprehension over multiple sentences. In Pro- ceedings of the 2018 Conference of the North American
59. URL httpp:://aa..cccr10.1185////////
2022. f findi

---

NTL The Number Token Loss
Saxena, E., Alfarano, A., Wenger, E., and Lauter, K. Teach- ing transformers modular arithmetic scale, 2024. URL https: /atp:///axivvvv org/ //////// abs/ 20s/24400 22104.00... 03569.
A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I,, Zhang, Y., Fan, A., Kambadur, M., Narang. S., Rodriguez, A., Sto- jnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models, 2023. URL https: / /osss//arree htt///axvvvoooo/ wog.orr/////.... aaaaab//23300 2307.09288.
Saxton, D., Grefenstette, E., Hill, F., and Kohli, P. Analysing mathematical reasonnn abilities of neural models, 2019.
URL httpt: arxiv...vvvv. org/ abs/ 1904. 01557.
Vaswani, A., Shazeer, N.,, Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., et al. Attention is all you need. Ad- vances in neural information processing systems, 30(1): 261-272, 2017.
Shah, D. and Aamodt, T. M. Learning label encodings for deep regression. In The Eleventh International Confer- ence on Learning Representations, 2023, URL https: //openreview net/ forum?
Wallace, E., Wang, Y., Li, S., Singh, S., and Gardner, M. Do NLP models know numbers? probing numeracy in embeddings, 2019.
Singh, A. K. and Strouse, D. Tokenization counts: the impact of tokenization on arithmetic in frontier Ilms, 2024. URL https:tp///// /tt//aaxivi/// /btggg/bbb////// 14903.
Wang. J., Qin, R., Wang, M., Fang, M., Zhang, Y., Zhu, Y., Su, Q.. Gou, Q. Shen, C., Zhang, O., Wu, Z., Jiang, D., Zhang, X., Zhao, H., Ge. J., Wu, Z., Kang, Y., Hsieh, C.-Y,, and Hou, T. Token-mol 1.0: tokenized drug design with large language models. Nature Commumications, 16 (1):4416, 2025. doi: 10. 10.1038/s41467-000 59628-y.
Stevens, S., Wenger, E., Li, C.,, Nolte, N., Saxena, E., Char- ton, F., and Lauter, K. Salsa fresca: Angular embed- dings and pre-training for ml attacks on learning with er- rors, 2024. URL https: httppps://axxv..///eeeeasstsmmr/ org/ wttgg////////////// 01082.
Stewart, L., Bach, F., Berthet, Q., and Vert, J.-P. Regression as classification: Influence of task formulation neural network features. In International Conference Artifi- cial Intelligence and Statistics, pp. 11563-11582 PMLR, 2023.
Wei, J., Wang. X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V., Zhou, D., et al. Sfain-of-thhoulee prompting elicits reasoning in large language models. Advances in neural information processing systems, 01-24824-24831 2022.
Sundararaman, D., Si, S., Subramanian, V., Wang, G., Haz- arika, D., and Carin, L. Methods for cummeracy-preserving word embeddings. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 4742 4753, 2020.
Weller, O. and Seppi. K. The rjokes dataset: a large scale humor collection. In Proceedings of the Twelfth Language Resources and Evaluation Conference, pp. 3136-6141, 2020.
Team, G., Mesnard, T., Hardin, C., Dadashi, R., Bhupatiraju, S., Pathak, S., Sifre, L., Rivière, M., Kale, M. S., Love, J., al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.0822 2024.
Wolf, T., Debut, I.., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M.. et al. Transformers: State-of state-of-the-ant art natural language processing. In Proceedings of the 2020 conference on em- pirical methods natural language processing: system demonstrations, pp. 38-45, 2020.
Thawani, A., Pujara, J., Ilievski, F., and Szekely. Rep- resenting numbers in nlp: survey and vision. In Pro- ceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 644-656, 2021.
Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T. On layer normalization in the transformer architecture. In Inter- national Conference on Machine Learning, pp. 10524- 10533. PMLR, 2020.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A.. Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L.., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami. V., Goyal, N., Hartshorn, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Moly- bog, I., Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subra- manian, R., Tan, X. E., Tang, B., Taylor, R., Williams,
Zhang, X., Ramachandran, D., Tenney, I., Elazar, Y., and Roth, D. Do language embeddings cap- ture scales? In Cohn, T., He, Y., and Liu, Y. (eds.), Findings of the Association for Computa- tional Linguistics: EMNLP 2020, pp. 4899. On- line, November 2020. Association for Computational Linguistics. doi:
439. URL https: httpt::ttacalaeeeeee /aaa/aaantttt //aalaaaacnloggooo..... org/2020. SEndings-emnlp.. 439.
12

---

NTL The Number Token Loss
Zhang-Li, D., Lin, N., Yu, J., Zhang, Z., Yao, Z.,, Zhang, X., Hou, I.., Zhang. J., and Li, J. Reverse that number! decoding order matters in arithmetic learning, arXiv preprint arXiv:2403.05845, 2024.
Zhong, Q.. Wang, K., Xu, Z., Liu, J., Ding, L., Du, B., and Tao, D. Achieving > 97% on GSM8K: Deeply under- standing the problems makes LLMs perfect reasoners. arXiv preprint arXiv:2404. 14963, 2024.
13

---

NTL The Number Token Loss
A. Appendix
A.I. Statement on code
The code for this paper is available at https://gtttabooomm//ummm/uumbbe
A.2. Algorithm for the Number Token Loss
Algorithm Al Pseudo-code compute NTL-MSE
{
int(vocab(i)) if vocabli] E R
Initialize: n vocab
NaN
otherwise
i=1
2:
3: function FORWARD(logits E BXXXV labels E BxT)) Float
4: ntl +- 0
5: n_logits + logits|:; vo.vobbbissnan))
Ignore non-number tokens
6: probs Softmax (logit)
7:
E; ;r n_vocab
ÿ B x T y B T Can be any regression loss
8: nlvooab(labels]
9: ntl f- MSE(v.
10: return ntl
11: end function
Algorithm Pseudo-code to compute NTL-WAS
int(vocab(i))
if vocab[i] E R otherwise
1: Initialize: n vocab
NaN
i=1
if order numbers True then
3: Sort the numbers in n_vocab by their numerical values
4: end if
5:
6: function FORwARD(logits E BxXxxV labels E BBXT)) Float
7: nogigs + logits|;; nn.vocab.isnan))
Di Ignore non-number tokens
8: n.probs + Softmax(loging
9: y < n_vocab[labe [labls]
Retrieve true numerical values
if using CDF version of NTL-WAS then
11:
y distr(b, t| one (b, t), num classes=len(n vocab))
One hot encode y
beaibeesttin_iistaneee (b, VV (CDF(n.probsjb, f)) CDF(C_ distr(b, [)] end if
m=1
14:
if using absolute difference version of NTL-WWAS then
15:
wasserstein_distanee = + Σ n-probs|b, v] |n_vocaffe] y
16: end if
17: ntl +- Mean(wassersstenn_diis (itpssttttiniissttneemyyysaaoo) (-y.isnan()
18: return ntl
19:
end function
A.3. T5 Architecture
The T5 model is built upon the Transformer architecture (Vaswani et al., 2017), consisting of stacked self-attention and feed-forward layers in both the encoder and decoder. The encoder processes the input tokens to create contextualized representations, while the decoder generates the output tokens autoregressively. attending to both the encoder's 's outputs and the previously generated tokens. The model can be trained with both Masked Language Modelling (MLM) (Kenton Toutanova, 2019) and Causal/Auto- Aut--Reesssivve Language Modelling (CLM) (Dai & Le, 2015), whereby we chose to use CLM.
14

---

NTL The Number Token Loss
A.4. Regression Transformer
The Regression Transformer (Born & Manica, 2023) preserves the inherent structure of numbers by inducing information on relative proximity through numerical encodings that set deterministically for all tokens. For every combination of decimal place and digit value, corresponding numerical token is added to the vocabulary. For instance, the number 11.4 is tokenized to [1-1, 1_0, 4.-1].
Non-numeric tokens are set to zero vectors. The numerical encodings are designed so that their pairwise distances are symmetric and monotonically decreasing with the float value. The final encoding of the input tokens is obtained by summing over numerical and regular word encodings. The Regression Transformer numerical encodings NE dimension j for numerical token hy.p with value and decimal place p can be determined by
NE (v. p. (-1)² V. 10P
i+1
A.5. Challenges with Integrating Val in Transformer Models like T5
In transformer models like T5, integrating numerical encoding schemes like Val presents challenges. Val multiplyies the [NUM] token embedding X by the number value a. In T5, however, a per-sample pre-layer normalization is applied immediately after the embedding, which effectively removes the scaling by a. Specifically:
aX - E|aX
aX - aE Va²E(X²) - a² (
X-E E(X
(10)
((a)
(X)
Hence, under T5's architecture, all numbers collapse to the same embedding, making x Val incompatible with T5.
Even in the original x Val urchitecture, the range of values Val can process meaningfully is limited by the layer normalization that follows the positional embedding step. Please see Golkar al. (2023) under "Implicit normalization via layer-norm" for more information. Therefore, XVal normalizes each value to [-5, 5] prior to training mitigate this issue.
We argue that this approach cannot be applied in practice, since in real texts the range of numbers is not known in advance, and thus simple min-max normalisation to [-5,5] 5| prior to training inference is not really practical.
Therefore we opted for a simpler approach in our experiments: Applying a signed log + x)) transformation to all numeric inputs. This avoids the overhead of parsing and re-scaling each number to 5-5, 5] prior to training, but it also has the drawback that large numbers are squashed the embedding space, making fine-grained distinctions difficult for the model.
For direct comparison without any modifications the x Val processing, we repeated the 3-digit multiplication experiment from Golkar al. (2023). The results can be in Table A1.
Encoding
R² Value
P10
0.9989
P1000
0.9997
B199
0.9998
FP15
0.7119
VVal
0.9986
T5 CE
0.999934
TS NTL.-WAS
0.999997
T5 regression head
0,999891
Table A1. R² scores for various number encoding methods on the 3-digit multiplication experiment from Golkar et al. (2023).
A.6. Experiments
For all trainings, we use trans (Wolf et al., 2020) 4.42.4. We train with a batch size of 32, a learning rate of le-4 and a weight decay of 0.01. All models were trained on single graphics processing units (GPUs) of type NVIDIA RTX A6000, NVIDIA A100 or NYIDIA A40.
15

---

NTL The Number Token Loss
A. MULTITASK MATHEMATICS DATASET
To test the mathematical capabilities of the methods, use subset of the mathematical question-answen dataset from DeepMind (Saxton al., 2019). The dataset was generated synthetically and therefore contains limited linguistic variability, but is sufficient for purposes to compare the mathematical capabilities of the different methods.
The dataset contains different modules and difficulty levels. For training and testing the models, we chose all difficulty levels but excluded modules where the answer contains complex fractions or variables. This allows us to focus on purely numeric answers to simplify the evaluation the model and still leaves us with a large enough dataset of ~26 million samples.
For training, validation, and interpolation tests, we selected the following modules from the DeepMind mathematical question-answer dataset:
algebra.lliear.dd.css
alggbra_Inmeerrddccmmoddddd
at
algbza__lieea.22ccommeeddd..
algebsa_sggeenneeeexxttmemtttt
arithmeticc..dd.oo..bbbkktt
arichmetic.mull.txt
numberss_ddiy.reemiiddee...tttt
numberss.plaace.vallee.tttt
numberss-roondd.nubbbr.....
For extrapolation tests, we selected the following modules:
axttmattic..addlo...uu.........
arithmettic.eedd.sb.... sub longer.txt
ariChmetic..niveedloggeee.tt
arithmettic_mml..iig..kk
numberss-plaee..aalue.........
numbers_roond.nubbbr.b....g
This resulted in a training dataset of 25,986,948 samples, a validation dataset of 13,026 samples, an interpolation test set of 130,000 samples, and an extrapolation test set 70,000 samples.
We train each model for 1050000 iterations. For these experiments, we used the TS-base architecture (220M parameters). For the Number Token Loss, trained with the hyperparameter À set to 0.3. The results can be in Table and Figure A1.
A.6.2. ABLATION STUDIES AND NTL-WAS-CDE WITH GAUSSIAN SMOOTHING
For training and validation, we selected a subset of the DeepMind mathematical Q&A dataset, in all three diffi- culty levels: arithmetic add sub multiple txt. Similarly, interpolation and extrapolation tests were done on ari trith ic aidiii_dddd__uueii wtt.ssuuuslltiseeessuee.llssl multiple.. txt and arithmetic add.heettc..ddddeeeeeeee add_sub mm welliipeellle haple_longer... txt, respectively. This resulted in training set 100,000 examples, a validation set 3,000 examples, and two test sets of 10,000 examples each.
16

---

NTL
The Number Token Loss
Accuracy
Accuracy
T5
T5
T5 + NTL MSE
T5 NTL MSE
T5 NTL WAS
0.75
+ NTL WAS
0.432
RT
RT
xVal
XVal
MAE
MAE
T5
NTL MSE
T5 NTL
T5 NTL WAS
0.10
NTL
WAS
0.744
RT
xVal
xVal
0.30
0.25
0.4
0.6
1.2
R²
R²
0.913
T5 NTL MSE
T5 NTL
MSE
T5 NTL WAS
0.98
T5 NTL WAS
0.913
RT
RT
xVal
xVal
0.0
0.2
0.4
0.6
0.8
1.0
1.2
0.0
0.2
0.4
0.6
0.5
1.0
(a) Evaluation metries interpolated test data.
(b) Evaluation metrics on extrapolated test data.
Figure AI. Comparison of evaluation metrics interpolated and extrapolated test data.
For these experiments, we used the T5-small architecture (60M parameters). The mean and standard deviation were taken over 4 different runs for the ablation studies and 3 different runs for the studies on GCE. We trained each run for 2.5M iterations.
In addition to accuracy, mean absolute and R² given in Table and Table 3, the achieved Pearson correlation and Spearman's is rank correlation coefficient are shown in Table A2.
A.6.3. NTL FOR REGRESSION AND NORMAL NLP TASK
We trained and evaluated the NTL on two tasks: a regression task using the rJokes dataset (Weller & Seppi, 2020) and test generation task using the transformed MultiRC dataset (Khashabi et al., 2018). In both cases, each model was trained for 2.5M iterations. The mean and standard deviation were taken over different runs. When using T5 with NTL-WAS, we set the hyperparameter À to 2.0.
In addition to RMSE and Pearson correlation given in Table 6, mean squared error, accuracy and Spearman's 's rank correlation coefficient are shown in Table A3,
A.6.4. NTL IS EFFECTIVE FOR DIFFERENT TOKENIZATIONS
To evaluate the effect of the tokenizer we trained T5-small again for 2.5M iterations. For the NTL-WAS loss we set the hyperparameter to 2.0. A comparison of the achieved number accuracies and Pearson correlation the models is shown in Table A5, Table and Table A3.
17

---

NTL The Number Token Loss
Table A2. Pearson and Spearman Correlation Coefficients on the Interpolation and Extrapolation Test Sets
Loss
Pearson
Spearman
Interpolation
Extrapolation
Interpolation
Extrapolation
CE
0.98
0..0
0.98
0.86
0.3
1.00
0.86
1.00
0.8940.00
NTL-MSE
0.8
0.99 19991..00
0.87 0.87
10.00
1.00
0.9010.00
2.0
0. 0.99
0.8640.01
0.99
0.8940.01
0.3
1.00
0.88
1.00
0.91
NTL-WAS
0.8
1.00
0.86
1.00
0.8940.00
2.0
0.99
0.85
1.00
0.88+0.01
NTL-MAE
0.3
1.00
0.87
1.00
0.90
NTL-Huber
0.3
1.00
0.87
1.00
0.90
GS CE
0.49
10.38
0.20.1..22
1.00
0.90
GS NTL-WAS
0.3
1.00+00000
0.87+0.01
1.00g0.00
0.90+0.01
Table A3. Full Evaluation on the rJokes Dataset. Evaluation metries (mean + std).
Model
MSE
Number accuracy
Pearson
Spearman
Standard T5 (CE)
4.03 + 0.03
0.30 + 0.00
0.41 + 0.00
0.35 + 0.00
T5 NTL
3.27+0.04
0.30 + 0.00
0.44
+
0.40
10.00
T5 Regression Head (MSE)
3.32 = 0.04
0.00 0.00
0.45 #0.00
0.40
10.00
T5 NTL + Standard tokenizer
3.88
0.30 1 0.00
0,41 0.00
0.35 1.00
Standard T5 (CE) Custom Tokenizer
4.09 + 0.04
0.30 + 0.00
0.41 + 0.00
0.37 0.00
A.6.5. ERROR ANALYIS
we conducted a detailed error analysis on the GSM8K dataset to examine predictions for numbers ending with specific digits (0-9), comparing CF and NTL. The error histograms shown in Figure A3 reveal a consistent pattern across all digit groups: NTL error distributions are narrower and concentrated around zero, confirming improved numerical reasoning compared to CE.
We further analyzed model predictions for numbers near digit boundaries, focusing particularly numbers ending with O 9 well as powers of 10. Table A4 reveals how often predictions overestimations, underestimations and exact matches. The results show, that NTL achieves more balanced error distribution. The exact match accuracy is consistently higher, but particularly so for samples ending with the 9 token, implying that NTL leads to better handling on those digit boundaries.
18

---

NTL The Number Token Loss
Table Error Rates. Error Rates by Digit Boundary Sample Type on GSM8K.
Sample Type
Metric
CE (%)
NTL (%)
Overestimation Rate
28.4
29.4
Ends with O
Underestimation Rate
56.5
51.1
Exact Match Rate
15.0
19.4
Overestimation Rate
54.4
49.1
Power of 10
Underestimation Rate
24.6
24.6
Exact Match Rate
21.1
26.3
Overestimation Rate
33.3
38.8
Ends with 9
Underestimation Rate
46.7
32.6
Exact Match Rate
20.0
28.6
Table A5. Interpolation Performance for Different Models and Tokenizers
Custom Tokenizer
Accuracy (Interpolate)
Pearson (Interpolate)
Loss
CE
X
0.34
0.98+0.00
NTL-WAS
X
0.39
0.9640.00
CE
V
0.45+0.00
1.00
NTL-WAS
V
0.43
1.00
Table A6. Gaussian Cross Entropy (GCE). Standalone and combinatory effect of GCE and NTL..
GCE a NTL Accuracy MAE
R2
Interpolation test set
X
0.34
2.15
0.95
V
0.43
0.91
0.99
X
0.19
4.34
0.95
V
0.34
3.14
0.96
0.5
X
0.42
0.95
0.99
0.5
V 0.48
0.76
0.99
Extrapolation test
X
X
0.05
61.92
0.61
V
0.10
58.18
0.68
X
0.03
126.74
0.08
V
0.06
111.09
0.25
0.5
X
0.10
58.55
0.65
0.5
V
0.10
66.97
0.59
Loss
МАPE
CE
0.502%
NTL-Squash-2
0.491%
NTL
0.485%
Table A7. Squashing NTL Loss. MAPE on the multiplication task using the GPT-2 model trained with a squashing NTL variiat.
19

---

NTL The Number Token Loss
Integer multiplication: Train on up to 5x5 digits, test on up to 6x6
Cross Entropy
Number Token Loss (NTL)
NTL Improvement
366
100
100
100
1.0%
5.0%
0.0%
Granite 3.1 1B Metric: Accuracy
Dights murippies
386
3CC
54.7%
50.7%
01.8%
1%
300
300
4%
05, 4%
6%
Digts muripleer
1,0%
7.5%
3.4
2.4%
25
25
366
35%
60. 3%
04
21.68 6/
4.0%
36C
3%
04: 4%
2/4
50 3/4
12.24
1.0%
1.0%
1.6%
10:44
20. 1/
1/mh
386
54.8%
66.4%
0.
.00%
D22
300
4%
67.2%
12:
90%
1.0%
1.0%
20 P5
1.0%
1.4%
20%
20%
multio icand
Diolts multlolicand
Dialts muttolicand
2.5
215
105
3.5%
1/5
105
1/5
105
3.5
0.22%
204
CPA
Granite 3.1 1B Metric: MAPE
Digits mutipler
3.4%
304
13%
31%
11%
135
4.6
3.0%
3.5%
135
13%
135
1.5
4.5%
mutilles
N-
0.33%
0.52%
0.04%
0.2
3.03%
8.4
326
1.4
15
8:00
1/m
23%
3.00
3:10
1%
2:2
2.5%
0.92%
1.%%
0.22%
0.24
222
0.00%
2.5%
2.5%
60%
3.0%
34%
1.5%
12%
5.
1.4%
3.0%
2.1%
1.0%
111%
6.%%
0.30%
0.33%
0.7%
033%
177%
4.19%
Digits
0.03%
4.%%
4.%%
0.996
0.7%
0.5%
Digits
Digits mulippicaad
Digits mutiilicand
306
166
ICG
100
BCG
96.9%%
16C
1:0
160
L00
3.0%
1.0%
0.0%
6.0%
5.0%
6%%
Metric: Accuracy
80%
30%
GPT2 Small
Digits muhiplee
16.2%
91.4%
9%
30C
50.3%
51.%%
Digits mubippee
1.0%
1.0%
1.2%
2.4%
6.6%
20.0%%
16.9%
50%
4%.2%
62.6%
50%
1,0%
-0%
.00m
Digii
-0%
12.09
2.6%
0.
12.45
2.6%
1.0%
..44%
2.0%
8.0%
33%
0.0%
Digits multip icand
Digits in multipiccad
Digits in multplicand
4.%%
4.%%
4440
GPT2 Small Metric: MAPE
Digits multppier
33%
3.0%
Digits multiller
3.4%
32%
1.5%
15%
3.4%
1.3%
1.4%
0.30%
0.%
3.5%
9/4
2.5
P/
PM
1.5%
115
2/1
pm
A/S
AM
+.5%%
0.28%
077
0.77
0.%
35%
1.3%
TM
7%
1/5%
1.24
3.0%
3.4%
1.0%
P%
1.0%
3.5%
1.2%%
A-
0.54%
0.44%
033%
4.9%%
037%
3.4%%
0.%
3.5
1.3%
90%
190
03%
2.45
1.1%
Pm
1/m
-55%
5.1%
9.%%
0.28%
927
3.1%
4.40%
4.2%
Digits in multip icand
Digits in multiplicad
Digits In mutiplicand
Figure A2, Performance improvement the multiplication task using NTL, evaluated GPT-2 and Granite 3.1.
20

---

NTL The Number Token Loss
2011
0L00 MEN
Figure A3, Error histograms for all number tokens for the GSM8k task.
21

---

NTL The Number Token Loss
Average Number Token Distributions in Early Training
Moan hoad baar Prcturi by
(E))
10-
Cal esokustions lngget Lay!t Misel Vue)
03
Ev cothee Saree Digi Digh Medel WWSS
Prikarians Trrger Eagery Modee TUT1, VI661
1.0
Dag Mdde
Figure A4. Average number token distributions for simple arithmetic tasks with target token 2 and 3 after training on the arithmetic task for 50k steps.
22

---

NTL The Number Token Loss
Example I: GSM8K sample NTL vs. CE
Question:
Roberta wants to have dinner party centered around soufflés. Each savory souftie calls for 8 eggs cach and the dessert ones call for eggs each. She wants make savory soufflés and dessert souftlés for the party. How many eggs will she need?
Correct answer:
The savory soufflés need eggs each and she wants to make of these so she needs 8 24 eggs The dessert souffiés need eggs each and she wants to make 5 of these so she needs 6*5 = 30 eggs For the dinner party she will need 24+30 54 eggs in total #### 54
T5-CE:
3=8=24 eggs. She wants to make dessert soufflés that are 6 eggs each for a total of 5"6=30 eggs. In total, she will need 24+30=62 eggs, # 62
Roberta wants make 3 savory soufflés are eggs
T5-NTL:
Roberta wants to make savory soufflés that use 8 eggs each so that's ' 3 = 24 eggs She wants to make dessert soufflés that use 6 eggs each that's 5 30 eggs All total, she needs 24 30 54 eggs #### 54
23


Reviews:
[{"review_id": "34plYOUQzh", "reviewer": "Reviewer_p8jQ", "summary": "This paper introduces Number Token Loss (NTL), a loss function designed to improve numerical reasoning in language models (LMs). The core idea is that standard cross-entropy (CE) loss treats numbers as categorical variables, disregarding numerical proximity. NTL aims to address this by incorporating numerical relationships into training, proposing two variants:\n1) NTL-MSE: Uses Mean Squared Error (MSE) between predicted number token distributions and the ground truth.\n2) NTL-WAS: Uses the Wasserstein-1 distance to align predicted and target number distributions.\n\nThe authors argue that NTL is model-agnostic, computationally efficient, and improves numerical reasoning while maintaining performance on text-based tasks. Experiments on mathematical tasks and a real-world regression dataset suggest that LMs trained with NTL perform comparably to those with dedicated regression heads.\n\n=======\n\n*update after rebuttal.*\n\nI appreciate the authors for additional experiments on a simple arithmetic baseline and on decoder-only models. \nAlthough I am not fully convinced about the results on the effectiveness of the claimed methods, I would like to increase my score to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the constructive feedback. We are delighted that you agree that NTL shows consistent improvements. To clarify your questions, we ran additional experiments with decoder-only models on “simpler arithmetic tasks” (multiplication):\n\n(1) **Decoder-only**: NTL is simply a loss function and can be used with any model, also beyond Transformers. To prove our point, we constructed an arithmetic task, much akin to the length generalization task by Jeleassi et al (2023): multiplications of 2 numbers with $k$ and $l$ digits, with $k, l \\in [1..5]$ in training and $k, l \\in [1..6]$ in evaluation. The model was trained to answer the question: “What is the result of multiplying x with y?”. We tested **2 different decoder-only models (GPT2 and IBM’s Granite)** both on varying sizes from 125M to 2B parameters. We report the mean absolut percentage error (MAPE) separately for unseen interpolation (up to $5 \\times 5$ digits) and extrapolation ($m \\times 6$ digits) samples:\n\n| Model | Size | Interpolation CE | Interpolation NTL | Extrapolation CE | Extrapolation NTL |\n| --- | --- | --- | --- | --- | --- |\n| GPT2 Small | 125M | 0.55% | **0.49%** | 1.11% | **1.00%** |\n| GPT2 Medium | 350M | 0.43% | **0.42%** | 0.82% | 0.82% |\n| GPT2 Large | 774M | 0.39% | **0.37%** | 0.76% | **0.75%** |\n| GPT2 XL | 1.5B | 0.43% | **0.40%** | 0.83% | **0.82%** |\n| Granite 3.2 | 2B | 0.35% | **0.21%** | 0.60% | **0.42%** |\n| Granite 3.1 | 1B | 0.28% | **0.15%** | 0.68% | **0.23%** |\n\nNTL works consistently better for 2 different decoder-only architectures across all model sizes. Length generalization is also improved for NTL, see detailed results (digit by digit and extra metrics): http://bit.ly/41Sx8jp\n\nWe hope that this comprehensive experiment with 6 models & 2 backbones rules out your concern regarding generalizability to decoder-only models.\n\n(2) **Simpler baseline** **task & limited baselines:** We had already done experiments on the (simple) 3-digit multiplication problem from xVal. But since this dataset is too easy, we didnt show it in the paper, in favor of more challenging experiments. This compares also to other number encoding strategies than xVal. We will add this Table to the final Appendix. \n\n| Encoding | R2 Value |\n| --- | --- |\n| P10 | 0.9989 |\n| P1000 | 0.9997 |\n| B1999 | 0.9998 |\n| FP15 | 0.7119 |\n| xVal | 0.9986 |\n| T5 CE | 0.999934 |\n| T5 NTL-WAS | **0.999997** |\n| T5 regression head | 0.999891 |\n\n(3) **Comparison with** **Gaussian label smoothing:** Note that we already reported results for Gaussian label smoothing (GCE) in the submitted version (Table 3). We add a more direct comparison below. The results on the arithmetic dataset show that GCE combined with NTL yields best performance, showing that both methods are complementary.\n\nInterpolation test set\n| GCE  | σ  | NTL  | Accuracy | MAE  | R²  |\n|------|----|------|----------|------|------|\n| X    | -  | X    | 0.34     | 2.15 | 0.95 |\n| X    | -  | O    | 0.43     | 0.91 | 0.99 |\n| O    | 0.5| X    | 0.42     | 0.95 | **0.99** |\n| O    | 0.5| O    | **0.48** | **0.76** | **0.99** |\n\nExtrapolation test set\n| GCE  | σ  | NTL  | Accuracy | MAE  | R²  |\n|------|----|------|----------|------|------|\n| X    | -  | X    | 0.05     | 61.92 | 0.61 |\n| X    | -  | O    | **0.10** | **58.18** | **0.68** |\n| O    | 0.5| X    | **0.10** | 58.55 | 0.65 |\n| O    | 0.5| O    | **0.10** | 66.97 | 0.59 | \n\n(4) **Additional questions**\n\n1. **Single-digit edge cases (19 vs 20 vs 21)**: Yes, with single-digit tokenization, NTL would give a higher loss for 19 than for 21 since 9 is far from the ground truth. However, cross entropy has the same problem: The loss for 19 is higher than for 21, since two tokens are incorrect. However, NTL also works with multi-digit tokens, where such a case would not occur.  Nevertheless, we looked into the frequency of such cases. Please see the response to R4 aka 4obq for details.\n\n2. **Number embedding:** A PCA of number token embeddings of CE and NTL did not show significant differences, embeddings were indeed roughly circular in both cases. We will add the plots to the final appendix.\n3. **Continuous number embeddings:** You mention that we lack such comparison where actually xVal does use continuous number embeddings (see R2/fuzE for details).\n4. **Sample efficiency:** As suggested, we analyze the sample efficiency for the decoder-only digit multiplication task described above. We show the evolution of the MAPE during training. As expected, the errors decrease much faster for NTL than for CE loss. \n\n| MAPE @Epoch | 1 | 5 | 10 | 20 | 40 | 100 |\n| --- | --- | --- | --- | --- | --- | --- |\n| CE | 3.1% | 2.0% | 1.6% | 0.6% | 0.7% | 0.4% |\n| NTL | 2.7% | 1.3% | 1.1% | 0.9% | 0.5% | 0.3% |\n\nThis corresponds to $3.43$ epochs on average to achieve a MAPE $<0.5\\%$ with CE loss, compared to $2.55$ epochs with NTL. This difference is pronounced for more difficult multiplications, in case of interest see Figure at: [bit.ly/4lc6wkY](https://bit.ly/4lc6wkY)", "claims_and_evidence": "Their claim that NTL improves performance in arithmetics is clear and convincing. \n1. NTL improves performance on arithmetic tasks\n- Supported by empirical results from the DeepMind Mathematics dataset, showing improved accuracy and lower error rates compared to CE.\n- NTL shows consistent improvements in both interpolation and extrapolation tests.\n- However, a straightforward baseline using simple arithmetic tasks is missing.\n\n2. NTL does not degrade text generation performance\n- MultiRC dataset experiments confirm that NTL does not negatively impact text-based tasks.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The paper does not present formal theoretical claims.", "experimental_designs_or_analyses": "Yes, I assessed the validity of experiement designs.\n\nThe experiments are well-designed, with strengths including:\n- Model scale comparisons: Results on both small-scale (T5-small) and larger models (T5-3B) suggest scalability.\n- Computational efficiency analysis: Confirms that NTL adds minimal overhead (<2%).\n\nHowever, some gaps remain:\n- Limited alternative baselines: While xVal and Regression Transformer are tested, other numeracy-enhancing strategies (e.g., GCE, continuous number embeddings) are not considered.\n- No decoder-only architecture evaluation: The method is only tested on an encoder-decoder (T5) model. It is unclear whether NTL applies effectively to decoder-only models like GPT-style transformers.", "relation_to_broader_scientific_literature": "The paper situates itself within work on:\n- Numeric representations in LMs (Geva et al., 2020; Golkar et al., 2023)\n- Mathematical reasoning and arithmetic in LMs (Cobbe et al., 2021; Dziri et al., 2024)\n- Tokenization strategies for numbers (Born & Manica, 2023)", "essential_references_not_discussed": "Some relevant references are not discussed:\nNumber embeddings: \n- Do NLP Models Know Numbers? Probing Numeracy in Embeddings (Wallace et al., 2019)\n- Tokenization counts: the impact of tokenization on arithmetic in frontier LLMs (Sing and Strouse, 2024)\n\nSequential Predictions, Training Objective:\n- Teaching Arithmetic to Small Transformers (Lee et al., 2023) \n- The Pitfalls of Next-Token Prediction (Bachmann and Nagarajan, 2024)\n- Length Generalization in Arithmetic Transformers (Jeleassi et al., 2023)", "other_strengths_and_weaknesses": "### Strengths\n- The method is straightforward and easy to implement.\n- The writing is clear and well-organized.\n- The empirical results show consistent improvements across tasks.\n\n### Weaknesses\n- The method is only evaluated on an encoder-decoder model (T5). It is unclear if NTL applies to decoder-only architectures.\n- The connection between NTL and learned numerical representations is not analyzed.\n- The paper lacks a direct comparison with Gaussian label smoothing, which is another method for handling numeric proximity.", "other_comments_or_suggestions": "- Synthetic arithmetic tasks: Simple arithmetic tasks, such as those from Length Generalization in Arithmetic Transformers (Jeleassi et al., 2023), could serve as a baseline. This would allow comparison of NTL’s performance, sample efficiency, and training efficiency.\n- Applicability to decoder-only models: The paper should test whether NTL works for autoregressive models like GPT.", "questions_for_authors": "1. For digit-level tokenization, how does NTL handle edge cases where predictions differ by a single digit?\nFor example, if the ground truth is \"20\" but the model predicts \"19\", does NTL penalize this more than it should because \"9\" and \"0\" are numerically distant tokens? Can the authors analyze error cases to determine whether this issue occurs frequently?\n\n2. How does NTL affect learned number embeddings?\nA visualization of how number tokens cluster in embedding space (e.g., using PCA) would help illustrate whether NTL improves numerical representation learning. For example, it is known that training on CE loss leads to the embeddings of number tokens to form a circular shape with PCA. \n\n3. Can NTL be used in decoder-only architectures?\nSince the method is only tested on T5, it is unclear if it generalizes to decoder-only transformers like GPT.\n\n4. Does NTL improve sample efficiency for learning arithmetic tasks?\nDoes it allow models to learn numerical relationships with fewer training samples compared to CE?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s9HsBaMwAu", "reviewer": "Reviewer_fuzE", "summary": "This study proposes Number Token Loss (NTL), a new regression-like loss to better handle numbers in texts. With the NTL loss, the prediction of a number token is determined by the weighted average of numbers with their softmax probability from logits, and the NTL loss measures MSE between the weighted average and the ground truth. This approach is advantageous over prior methods in that a) it is model agnostic with minimal assumptions about the vocabulary and b) it comes with minor computational overhead. The experiments demonstrate its superiority in accuracy over the standard cross-entropy based approach and several prior methods including Gaussian Cross Entropy, Regression Transformer, and xVal, particularly in classification tasks.\n\n**update after rebuttal.**\n\nI appreciate the author's full elaborations and answers to my concerns, which greatly deepened my understanding of their work. The explanation and additional results address my concerns, and I suppose that this work is worth being presented in the main conference.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable feedback! We appreciate your recognition of NTL’s effectiveness and practicality, considering its compatibility with standard pipelines. To address your questions:\n\n**(1) Non-euclidean topology:**\n\nThanks a lot for raising this interesting point regarding transferability beyond Euclidean topology. We agree that our specific loss (Equation 4) makes this assumption but note that the general NTL-WAS formulation (Equation 2) uses an arbitrary cost function $c$ defining the pairwise cost between tokens. This allows to cover more general relationships between the numbers. \n\nWe confirmed this experimentally. Instead of using euclidean distances, we squashed the distances so that for a label 0, predicting 9 is not exactly 9x worse than predicting 1 but only 2x.\n Results showed performance better than CE but slightly worse than vanilla NTL for multiplying integers with up to 6 digits (GPT2, measured by MAPE; details in response **Error Analysis** for R4/4obq).\n\nMoreover, regarding modular arithmetics: Consider the modular addition task, where $y(n,m) = n+m\\ mod\\ p$ as described in “Grokking modular arithmetic” (Gromov, 2023). Here, a reasonable cost function could be found in $c(y_1, y_2) = min(abs(y_1-y_2), m-abs(y_1-y_2))$, which considers the “warping” property / periodicity of numbers in modular arithmetics. \n\n**(2) Extended explanation for xVal:**\n\nWe have carefully revisited the corresponding section and would like to clarify several points.\n\n1. **xVal Incompatibility with T5**  \nAs noted in Appendix A.5 of our paper, xVal multiplyies the [NUM] token embedding $X$ by the number value $a$. In T5, however, a per-sample pre-layer normalization is applied immediately after the embedding, which effectively removes the scaling by $a$. Specifically:\n$$\n\\frac{aX-E[aX]}{\\sigma(aX)}=\\frac{aX-aE(X)}{\\sqrt{a^2E(X^2)-a^2E(X)^2}}=\\frac{X-E(X)}{\\sigma(X)}\n$$\nHence, under T5’s architecture, **all numbers collapse to the same embedding**, making xVal incompatible with T5. Consequently, we do **not** use xVal with T5 in our experiments. Instead, we follow the **original** xVal encoder architecture (as in the xVal paper).\n\n2. **Limited Dynamic Range of xVal & Log-scaling**  \n\n    Even in the original xVal architecture, the range of values xVal can process meaningfully is limited by the layer normalization that follows the positional embedding step. For further Information on that, please see [xVal](https://openreview.net/pdf?id=KHDMZtoF4i) under section 2 “Implicit normalization via layer-norm”. Therefore, xVal normalizes each value to [-5,5] prior to training to mitigate this issue.\n\n    We argue that this approach cannot be applied in practice, since in real texts the range of numbers is not known in advance, and thus a simple min-max normalisation to [-5,5] prior to training or inference is not really practical.\n    \n    Therefore we opted for a **simpler** approach in our experiments: applying a signed log⁡(1+x) transformation to all numeric inputs. This avoids the overhead of parsing and re-scaling each number to [-5,5] prior to training, but it also has the drawback that **large numbers** are squashed in the embedding space, making fine-grained distinctions difficult for the model.  \n    Thus: Yes, the results in Table 1 **do** use this logarithmic mapping. We will clarify this explicitly in the table caption and experimental setup.\n\n3. **Experimental results on dataset from xVal paper**  \nFor a direct comparison without any modifications to the xVal processing, we repeated the 3-digit multiplication experiment from the  xVAl paper. Again, our model beats xVal (see response **Simpler baseline** for R1/p8jq).\n4. **Accuracy of xVal Predictions**  \nWe appreciate your request for explicit accuracy metrics and your idea of using the nearest neighbor for xVal predictions.\n\n    When rounding xVal predictions to match the decimal places of our dataset, the resulting accuracies are quite low:\n\n    - Interpolate: 0.052\n    - Extrapolate: 0.018\n     \n    If we reduce precision by rounding to only two decimals, the accuracy improves somewhat, but remains modest:\n\n    - Interpolate: 0.096\n    - Extrapolate: 0.058\n\n    These findings underscore that xVal struggles with larger numbers in particular.\n\n(3) **Hard math tasks:**\n\nWhile the paper indeed only covers tasks related to euclidean topology, we respectfully disagree that it does not include “hard math tasks”. E.g., the extrapolation task of DeepMind’s math dataset is very difficult (→ maximal accuracy of 10%, see Table 3). It includes a large variety of tasks, including not only arithmetics but also algebra, number conversions and polynomials. Furthermore, the GSM8k dataset (see Table 8) is an accepted benchmark for reasoning and even commercial LLMs cannot solve it perfectly so far.\n\n(4) **Literature:**  \n\nThanks for sharing additional references about arithmetics & number representation in LMs. We will add those to the final version of the manuscript.", "claims_and_evidence": "The claims are generally well tested by experiments. The experiments cover classification and regression tasks, and they examine the variants of the NTL loss, the scaling characteristics on model size, and training speed.", "methods_and_evaluation_criteria": "The NTL loss naturally implements the idea that incorrect predictions on the number tokens should be evaluated with some distance metric, so that the model can know whether the predictions were close or not. While this idea itself has been tested by several methods including the Gaussian Cross Entropy and xVal, the NTL is more advantageous in that it can work within the framework of classification. This allows users to introduce minimal modification in their codes. The experiments show comparisons of these methods with standard metrics including prediction accuracy, MAE, and R2 scores.\n\nA caveat of the proposed method is that it assumes the Euclidean topology or its variants. Namely, it assumes that 2 and 3 are closer than 2 and 16. However, this is not always the case, this should be carefully discussed. For example, when one performs modular arithmetic, or some operations that involve periodicity, these numbers do not equip any distance (it violates triangular inequality). Thus, in such a case, the injection of the Euclidean topology should be harmful to the learning, and the cross-entropy loss should perform better.\n\nIt agrees that the NTL is useful, but the authors should make readers aware of this.", "theoretical_claims": "This study does not contain any theoretical claims.", "experimental_designs_or_analyses": "I checked the setup of experiments. Overall, the experiments investigate the proposed method and variants well. There are several unclear points.\n> [l.190, right] this scheme (of xVal) is incompatible with T5.\nAppendix A.5 explains that the xVal is not compatible with T5 model, but then the authors should test it on other reasonable models. The experiments can underestimate the performance of xVal since it does not use xVal as it is.\n\n> [l.212, left] The effective number range of xVal is limited to [-5, 5] due to the combination of its scaling of the number token embedding and the layer-norm in the backbone.\nOver the experiments, the performance of the xVal is very poor. This part needs more explanation. Further, Tables 1(a) and 1(b) do not show the accuracy by xVal. While the xVal does not directly predict discrete tokens, the accuracy can be computed simply by nearest neighbor.\n> To take this into account, we scale our dataset for xVal with log(1+x)\nAre the results in Table 1 using this logarithm map, or not? I assume not, but then where can I find the results?", "relation_to_broader_scientific_literature": "The key contributions of this study relate to the studies that address number tokens. This study addresses general documents and basic math tasks, but it also relates to the literature on hard mathematical problems as given in the next cell.", "essential_references_not_discussed": "Relevant to the earlier comment on the modular arithmetic, I encourage the authors to include the literature on modular arithmetic, or more broadly arithmetic/symbolic computation tasks.\n\n1. Learning the greatest common divisor: explaining transformer predictions, Francois Charton, ICLR'24\n2. Impact of Pretraining Term Frequencies on Few-Shot Numerical Reasoning, Yasaman Razeghi, Robert L. Logan IV, Matt Gardner, Sameer Singh, ACL'22\n3. Learning to compute Gröbner bases, Hiroshi Kera, Yuki Ishihara, Yuta Kambe, Tristan Vaccon, Kazuhiro Yokoyama, NeurIPS'24\n4. Learning Label Encodings for Deep Regression, Deval Shah & Tor M. Aamodt, ICLR'23\n\nThe first one discusses the dependency of the successful GCD calculation on the base to represent numbers. The second one provides an insight in the number embedding from the the frequency perspective. The third introduces a regression head to regress the coefficients in polynomials (similar to xVal) and observes that it particularly performs poor on finite field coefficients (i.e., the case involving modular arithmetic). The last one is not about Transformer, but regression by classificaion has been widely studies, so covering this literature makes this study richer.", "other_strengths_and_weaknesses": "Most of major strengths have been mentioned above, but I appreciate the proposed method in terms of the compatibility of the standard classification-based pipeline. It is natural to introduce regression to provide richer supervisory signals to incorrect number token prediction, but simple introduction of regression head requires additional modification in the codes such as in auto-regressive generation. The proposed method does not require this.", "other_comments_or_suggestions": "Nothing. The paper is well-written and easy to follow.", "questions_for_authors": "Please refer to the other cells, but major concerns from me are: a) Are baseline methods (e.g., xVal) implemented and compared in a fair manner? b) The applicability of the proposed method to modular arithmetic and other hard math tasks is not tested and discussed.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ichcVHDDYW", "reviewer": "Reviewer_bnfs", "summary": "The paper introduces a Number Token Loss (NTL), a regression-like loss function designed to improve the numerical reasoning capabilities of Language Models (LMs). The core contribution is twofold:\n\nNTL-MSE: A loss function that computes the Mean Squared Error (MSE) between the numerical value of the label and the predicted probability distribution of number tokens.\n\nNTL-WAS: A loss function based on the Wasserstein-1 distance, measuring the difference between predicted and true distributions of number tokens.\n\nThe authors claim that NTL improves arithmetic tasks, can match regression heads, and scales well to large models while adding minimal computational overhead. Empirical results on mathematical reasoning tasks, regression benchmarks, and large-scale models (T5-3B on GSM8K) support these claims.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the constructive feedback. Below are detailed responses & new analyses:\n\n1. **Real-world task from physics:**\nTo demonstrate applicability to scientific data, we evaluate NTL on estimating molecular solubility — as studied by the Regression Transformer (RT). Each molecule is a SMILES string & the goal is to predict its solubility as a float. Following the RT paper we report performance over 3 random splits. Our results show, again, that training with NTL significantly improves performance over standard CE loss, reducing RMSE and increasing $R^2$:.\n    \n    | Loss     | RMSE                       | $R^2$                |\n    |:------------ |:------------------    |:------------------    |\n    | CE               | $1.08 \\pm 0.16$    | $0.72 \\pm 0.07$  |\n    | NTL             | $0.91 \\pm 0.07$   | $0.80 \\pm 0.03$ |\n    \nOur approach outperforms the baselines from the RT paper, including a RF (RMSE: 1.16) and XGBoost (RMSE: 1.05).\nRegarding time series: this can be modeled with NTL in principle but time series transformers do not typically employ token-based generation like language models do. Since NTL is an improvement to cross entropy, we believe that studying utility for time series forecasting goes beyond the scope of this paper.\n\n2. “**The claim that NTL-WAS is always preferable to NTL-MSE is not fully substantiated, as it depends on task properties”:**  Indeed from a theoretical perspective, NTL-WAS is always preferable since it has a unique minimum as shown in Figure 3. However, there is one advantage of NTL-MSE: It explicitly computes a numerical value from the logits (via dot product), during loss calculation. This float can be combined with arbitrary loss functions (MSE, MAE etc), or even be used at inference time for decoding numbers as dot product of all logits (rather than via beam search). In the future we aim to test an expression-loss where all parts of a mathematical expression are transformed to numbers in this way & the loss penalizes the consistency of the equation. Thus, in some usecases NTL-MSE can be advantageous, although the theoretical properties are weaker. We will clarify this in the manuscript.\n\n3. **Extrapolation benefits are more limited:** Note that the extrapolation performance with NTL **doubles** (from 5% to 10%) on the arithmetics dataset. In relative terms, the extrapolation benefits are stronger than interpolation benefits (Table 2).\nTo strengthen the point, we ran new experiments with 2 decoder-only models (GPT-2, IBM Granite) on a multiplication dataset. NTL again outperforms CE in both interpolation & extrapolation but particularly for extrapolation tasks. For details, see response “Decoder-only” to R1/p8jQ.\n\n4. **Finer analysis of error cases:** Please see the error analysis for R4/4obq\n\n5. **Effect of NTL-WAS on confidence**: We analyzed the logit distributions over all number tokens for simple arithmetic tasks in throughout training: NTL increases the model’s confidence in its numerical prediction, particularly in early training. **NTL-WAS** yields logits that are more sharply centered around the correct number compared to CE. See plot at: http://bit.ly/445tlAD\n\n6. **Tendency to predict the highest-probability digit slightly below the true number for NTL-MSE?** Since NTL-MSE computes the dot product, or *weighted* sum, there is no reason why it should systematically underestimate the value. We also confirmed this empirically, see analysis right above: logits are more centered for NTL, plot at: http://bit.ly/445tlAD\n\n7. **Extracting of numbers in training data & handling of phone numbers:** In practice, we dont extract numbers explicitly but use an indicator vector from the tokenizer to index the logits of the number tokens to compute NTL. NTL is not suitable in cases where numerical proximity is irrelevant (e.g., phone numbers). However, such tokens could either be excluded from NTL or a squashing transformation could be applied (for details see “Error Analysis” to R4/4obq).  Such cases are exceptions rather than the norm. Also note that we still minimize CE for all tokens, NTL is just an extra loss term. For many common number types—such as years, quantities, and measurements—numerical proximity matters and NTL is more meaningful.\n\n8. **Unclear results for NTL+Multi-digit tokenization (MDT):** We already ran this experiment on a regression task (Table 5). Single-digit tokenization (SDT) and NTL yield complementary benefit:  NTL+MDT is better than just MDT, but SDT + NTL is the best.\n\n9. **Projection of last hidden state to linear layer:** We already ran this experiment, see Section 4.3 (”NTL can match regression models”). We even used a more complex regression head than just a simple linear layer (i.e., 2 linear layers with dropout). Nevertheless, NTL performs on-par with the model trained with this regression head (see Table 4). This shows its competitiveness, considering that it is a LM that can still be used on non-numeric tasks", "claims_and_evidence": "NTL improves arithmetic performance compared to standard CE loss. The results on the DeepMind Mathematics Dataset (Table 1) show that both NTL variants improve accuracy and reduce MAE. The improvement is most pronounced in interpolation tasks, while extrapolation benefits are more limited.\n\nOn the rJokes dataset, NTL matches a regression head in RMSE and Pearson correlation (Table 4).\n\nNTL does not degrade performance on non-numerical tasks. The MultiRC dataset results (Table 7) confirm that adding NTL does not harm text generation.\n\nNTL scales well to large models. The authors apply NTL to a 3B parameter T5 model on GSM8K, improving accuracy from 13.5% to 17.7% (Table 8). The computational overhead is minimal (Figure 4).", "methods_and_evaluation_criteria": "The experimental design is mostly rigorous but has a few gaps:\n\nAppropriate Choice of Tasks: The use of arithmetic benchmarks (DeepMind Mathematics, GSM8K) is well-justified. However, evaluating real-world numerical tasks (e.g., time series. finance, physics) would strengthen applicability.\n\nEvaluation Metrics: Accuracy, MAE, and R² are appropriate but could be complemented by a finer analysis of failure cases, such as systematic bias towards certain number magnitudes.", "theoretical_claims": "The theoretical justification for NTL is mostly sound:\n\nCross-entropy's failure in numerical tasks is well-motivated. The issue that CE treats numbers as categorical rather than ordinal is a widely recognized problem.\n\nNTL-MSE's non-uniqueness issue (where a sum of probabilities can approximate the correct number without a peaked distribution) is correctly identified, and NTL-WAS addresses this.\n\nWasserstein-1 distance as a better loss function is conceptually strong. However, the claim that NTL-WAS is always preferable to NTL-MSE is not fully substantiated, as it depends on task properties.", "experimental_designs_or_analyses": "Tokenization Choices: While single-digit tokenization improves performance, it is unclear whether models trained with multi-digit tokenization + NTL can reach comparable results.", "relation_to_broader_scientific_literature": "The paper positions itself well within the literature on numeracy in LMs", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "Strengths:\nNovel and simple loss function that is easy to integrate.\nComprehensive experiments covering multiple architectures and tasks.\nSolid theoretical motivation for why CE is suboptimal for numerical tasks.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "For NTL-MSE, does the model tend to predict the highest-probability digit slightly below the true number?\n\nIf the true number is 9, does the model predict 8 with a probability mass summing to 9?\n\nHow does NTL-WAS affect model confidence in its numerical predictions?\n\nDoes the Wasserstein loss lead to over-smoothing, making models less confident?\n\nIn practice you need to extract all numbers in the training data and add your NTL, what happens to numbers like phone number, years?\n\nIf you directly project the last hidden state to number using a linear layer, will your method be better? This should be a baseline", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ih8XsU5yF6", "reviewer": "Reviewer_4obq", "summary": "Authors start with a clear objective in mind: rather than relying on strategies to fix the number token behaviour, it would be better to really treat numbers as special class and impose a new loss which allows the model to be strongly penalized if the prediction is far off. So authors introduced a new loss, called NTL, to improve model performances on arithmetic tasks. \n\nAuthors also introduced this line work to be model agnostic and can be added without overhead.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your constructive and positive review of our work!\n\n(1) **Larger models**: We completely agree that evaluating NTL on larger models beyond 3B would provide further evidence regarding its scalability. We aim to do this in future work, however, at present, our computational resources are limited to a standard academic setting and running experiment on models that dont fit on single GPUs is beyond our capacity for this rebuttal. However, please be aware of the additional experiments on integer multiplication with multiple decoder-only architectures, namely GPT2 (up to 1.5B) and Granite (up to 3B), that confirm the previous results from our paper (see response to R1 aka p8jQ for details).\n\n(2) **Floats**: Regarding your comment on unclear performance on floating numbers: First note our existing experiment on a proper regression task (rJokes dataset, Table 4 and also Table 5) where labels are in $\\in [0, 11]$. In this experiment NTL **matches** the performance of a regression head, whereas CE performs substantially worse. Secondly, we added an experiment on a real-world task from chemistry, where the labels are floats with $\\sim$ 2 digits precision, and NTL outperformed CE (see response to R3 / bnfs). Third, note that DeepMind’s Mathematics dataset (see Table 1) contains almost 5M samples (18.7% of all 25M training samples) where labels are floats and not integers. This is particularly pronounced in the extrapolation test (28%) where we see a substantial improvement with NTL over CE. Together, these findings make us confident that the benefits of NTL extends well to floats.\n\n(3) **Error analysis:** One practical issue for models that tokenize numbers in multiple digits is that some tokens have large numerical value which disproportionally affects the loss (even if the logit is low), especially if the number tokens are not regularly spaced. We recommend to enforce digit-level tokenization as this ensure NTL is well-behaved. As a more conceptual mitigation strategy we note that NTL is not limited to using Euclidean distance between numbers. We will update Eq (4) in the paper to reflect this and emphasize that distances can be defined in a fully flexible manner. For example one can squash the distances so that, for a label 0, predicting 9 is not exactly 9x worse than predicting 1 but only 2x. We have done an experiment on this that confirms that it works still better than CE but less good than vanilla NTL on multiplying integers with up to 6 digits with GPT2-Large as measured by Mean Absolute Percentag Error (MAPE):\n\n| Loss | MAPE |\n| --- | --- |\n| CE | 0.502% |\n| NTL-Squash-2 | 0.491% |\n| NTL | 0.485% |\nFor further details on this experiment see our response \"decoder-only\" to R1/p8jQ. \n\nMoreover, note that this transformation does not need to adhere to the mathematical definition of a distance, the user could provide any pairwise distance matrix between number tokens. This allows to handle even exotic cases like modular arithmetics (for details see comment from R2 aka fuzE). We will clarify this in the final manuscript.\n\nAdditionally, we conducted a detailed error analysis on the GSM8K dataset to examine predictions for numbers ending with specific digits (0–9), comparing CE and NTL. The error histograms (see [last_digit_vs_distance_histogram.png](https://anonymous.4open.science/r/number-token-loss-5137/resources/last_digit_vs_distance_histogram.png)) reveal a consistent pattern across all digit groups: NTL error distributions are narrower and concentrated around zero, confirming improved numerical reasoning and lower systematic biases compared to CE.\n\nWe further investigated errors specifically at digit boundaries (e.g., numbers ending in 0 or 9) on the GSM8K dataset. The table below breaks down model errors at specific digit boundaries and  highlights how often predictions are overestimations, underestimations, and exact matches.\n\n| **Sample Type** | **Metric** | **CE** | **NTL** |\n| --- | --- | --- | --- |\n| **Ends with 0** | Overestimation Rate | 28.4% | 29.4% |\n|  | Underestimation Rate | 56.5% | 51.1% |\n|  | Exact Match Rate | 15% | 19.4% |\n| **Power of 10** | Overestimation Rate | 54.4% | 49.1% |\n|  | Underestimation Rate | 24.6% | 24.6% |\n|  | Exact Match Rate | 21.1% | 26.3% |\n| **Ends with 9** | Overestimation Rate | 33.3% | 38.8% |\n|  | Underestimation Rate | 46.7% | 32.6% |\n|  | Exact Match Rate | 20% | 28.6% |\n\nThe results show that NTL achieves a more balanced error distribution. The exact match accuracy is consistently higher, but particularly so for samples ending with the 9 token, implying that NTL leads to a better handling on those digit boundaries.\n\n(4) **Literature:**  Thanks for sharing the literature on modular arithmetics and cryptography. We will highlight these works in the final version of the paper.", "claims_and_evidence": "Claims are pretty rigorous. Authors provide extensive experiments across different datasets and show that NTL outperforms standard CE loss and additional baselines in numerical task. \nAnother critical point shown by the authors is the fact that NTL does not degrade performance in text-only task, which is a strong point in favor on author's claim of being easy to plug an play.", "methods_and_evaluation_criteria": "The evaluation criteria (accuracy, MAE, R²) are appropriate for the tasks being evaluated, and the authors provide a thorough comparison with baseline methods.", "theoretical_claims": "No theoretical claim has been made by the authors, and for future work I strongly suggest to understand if theoretical guarantees can be made on this new loss.", "experimental_designs_or_analyses": "The experimental design makes sense with a good ablation study to validate their proposed approach. They compare NTL with multiple baselines on different tasks and model size.", "relation_to_broader_scientific_literature": "The paper is within the literature on improving numerical abilities from the language models. Authors cite relevant prior work, like Regression Transformer, CoT and verifiers. I can suggest to add symbolic reasoning methods or program aided language models, so we can compare NTL with these techniques as well.", "essential_references_not_discussed": "Although this comes from a different area, https://arxiv.org/pdf/2402.01082 provides a new loss to treat numerical tokens. Similarly https://arxiv.org/abs/2410.03569 shows how to use this loss on very hard mathematical problem.", "other_strengths_and_weaknesses": "The paper is particularly original and significative. Before integrating to any large scale pretraining, it would be beneficial if authors can train a larger model (say 7B or 20B params) to be more conclusive on the pretraining part.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "How does NTL perform on even larger models (say 7B or 20B)?\n\nIt's relatively unclear how the model performance on floating number. \n\nCan you add a breakdown of where the model fails (maybe numerical ranges or specific types of arithmetic operations) to add more intuition of where NTL can be improved in the future?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes two new losses for language modeling that are meant to encourage better predictions on numerical tasks, by leveraging the ordered structure of numbers, either through an expectation of predicted numbers with an MSE loss, or through a specific choice of transport costs with a Wasserstein loss. The authors obtain encouraging results on mathematical and scientific datasets.

The reviewers were generally positive, particularly after the author response, thus I am leaning towards acceptance. That said, some additional experiments on decoder-only models would be good to include, and a study of sample/compute efficiency and scaling laws would be helpful here to assess if the proposed losses should actually be preferred to discrete tokenization in different regimes. The authors should include the additional results and datasets provided during the rebuttal in the final version, if the paper is accepted.
Decision: Accept (poster)