Regress, Don't Guess: A Regression-like Loss on Number Tokens for Language Models

We appreciate your constructive and positive review of our work!

(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).

(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.

(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):

Moreover, 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.

Additionally, 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) 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.

We 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.

The 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.

(4) Literature: Thanks for sharing the literature on modular arithmetics and cryptography. We will highlight these works in the final version of the paper.

Thanks for the constructive feedback. Below are detailed responses & new analyses:

1. Real-world task from physics: To 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$:.

   Loss	RMSE	$R^2$
   CE	$1.08 \pm 0.16$	$0.72 \pm 0.07$
   NTL	$0.91 \pm 0.07$	$0.80 \pm 0.03$

Our approach outperforms the baselines from the RT paper, including a RF (RMSE: 1.16) and XGBoost (RMSE: 1.05). Regarding 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.

2. “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.

3. 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). To 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.

4. Finer analysis of error cases: Please see the error analysis for R4/4obq

5. 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

6. 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

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

8. 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.

9. 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

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:

(1) Non-euclidean topology:

Thanks 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.

We 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. 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).

Moreover, 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.

(2) Extended explanation for xVal:

We have carefully revisited the corresponding section and would like to clarify several points.

1. xVal Incompatibility with T5
   As 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:

Hence, 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).

2. Limited Dynamic Range of xVal & Log-scaling

   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 under section 2 “Implicit normalization via layer-norm”. Therefore, xVal normalizes each value to [-5,5] prior to training to 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 a simple min-max normalisation to [-5,5] prior to training or inference is not really practical.

   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.
   Thus: Yes, the results in Table 1 do use this logarithmic mapping. We will clarify this explicitly in the table caption and experimental setup.

3. Experimental results on dataset from xVal paper
   For 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).

4. Accuracy of xVal Predictions
   We appreciate your request for explicit accuracy metrics and your idea of using the nearest neighbor for xVal predictions.

   When rounding xVal predictions to match the decimal places of our dataset, the resulting accuracies are quite low:

   • Interpolate: 0.052
   • Extrapolate: 0.018

   If we reduce precision by rounding to only two decimals, the accuracy improves somewhat, but remains modest:

   • Interpolate: 0.096
   • Extrapolate: 0.058

   These findings underscore that xVal struggles with larger numbers in particular.

(3) Hard math tasks:

While 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.

(4) Literature:

Thanks for sharing additional references about arithmetics & number representation in LMs. We will add those to the final version of the manuscript.

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):

(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:

NTL 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

We hope that this comprehensive experiment with 6 models & 2 backbones rules out your concern regarding generalizability to decoder-only models.

(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.

(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.

Interpolation test set

Extrapolation test set

(4) Additional questions

1. 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.

2. 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.

3. Continuous number embeddings: You mention that we lack such comparison where actually xVal does use continuous number embeddings (see R2/fuzE for details).

4. 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.

This 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