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

Authors: jonas  zausinger
lars  pennig
kacper  chlodny
vincent  limbach
anna  ketteler
thorben  prein
vishwa mohan singh
michael morris danziger
jannis  born
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-11 20:30:02

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 reasoning over quantities, especially arithmetics.
  This has particular relevance in scientific datasets where combinations of text and numerical data are abundant. 
  One fundamental limitation is the nature of the CE loss, which assumes a nominal (categorical) scale and thus cannot convey proximity between generated number tokens.
  As a remedy, we here present two versions of a number token loss. The first is based on an $L_p$ loss between the ground truth token value and the weighted sum of the predicted class probabilities. 
  The second loss minimizes the Wasserstein-1 distance between the distribution of the predicted output probabilities and the ground truth distribution.
  These regression-like losses can easily be added to any language model and extend the CE objective during training.
  We compare the proposed schemes on a mathematics dataset against existing tokenization, encoding, and decoding schemes for improving number representation in language models. Our results reveal a significant improvement in numerical accuracy when equipping a standard T5 model with the proposed loss schemes.

Main paper: 
 introduction: 
 As coined by Thawani et al. [14] , numbers in natural texts are ubiquitous and important, yet systematically neglected by language models (LMs). Even worse, while Transformers [15] were invented for NLP, they have permeated various scientific domains (chemistry, biology, etc [2, 8, 1] ), where tabular/numerical data is more prevalent than in NLP and often even fundamental for constructing task definitions: Molecules are labeled with drug efficacy, chemical reactions with yield, and synthesis procedures are natural text interspersed with quantities and times. Still, LMs notoriously struggle even with simple arithmetic tasks like three-digit multiplication [5] for multiple reasons:1. Tokenization: Standard subword tokenization splits numbers into arbitrary tokens, disrupting their structure. Mitigation strategies include scientific notation [18] or digit-level tokenization [6] , which may also preserve the decimal order of each digit [1] . 2. Embedding: Canonically, the model has to recover the structure of numbers from data because the embeddings of numerical tokens are learned like any other token. Countless flavors of numeracy-preserving word embeddings exist [13, 1, 7] , often akin to positional encodings. 
number token loss: 
 The idea of the Number Token Loss (NTL) is to add an additional loss term to the CE, which is only applied to number tokens and takes their numerical proximity into account. To achieve this, we propose two versions.Number Token Loss with Mean Squared Error (NTL-MSE) This loss compares the numerical value of the ground truth token with the weighted sum of the respective numerical token values, with the weights corresponding to the predicted class probabilities (cf. Figure 1 right ). Given a model f (•), input tokens x ≤i (where i ≤ N ), the numerical value ŷi of ground truth token y i and a vocab V consisting of tokens (with indices j, ..., k representing the number tokens), we compute NTL-MSE:EQUATIONInstead of a nominal-scale loss with regular CE, the NTL-MSE effectively conveys proximity between numbers. For example, if the label is [4] and the LM predicts a [5] instead of [9] , the loss will be lower, matching our intuitive expectation, unlike the CE which gives constant loss no matter the proximity of the number (cf. Figure 2 ). This is sufficient for the vast majority of cases, however, since the NTL is not injective (like CE), it can return spuriously low loss for incorrect predictions. Consider e.g., a label [4] with 50% of the mass on [0] and 50% on [8] token, then NTL will be zero (Figure 3 ). While such cases are rare due to the softmax emphasizing logit differences, combining NTL with CE loss helps correct spurious cases, as CE continues refining predictions without reducing its value in these instances. However, to address this non-injectiveness, we propose a second version based on the Wasserstein-1 distance.Number Token Loss with Wasserstein-1 distance (NTL-WAS) This loss calculates the Wasserstein-1 distance between the predicted probability distribution of the (sorted) number tokens and the ground truth probability distribution, which is 1 for the label token and 0 for all other tokens. Given the ground truth y i , a vocab V with number tokens ordered from indices j to k and the cumulative distribution function CDF(•), we compute NTL-WAS:EQUATIONAs one can see in Figure 2 , this version of the NTL not only conveys proximity between numbers correctly but also eliminates the non-injectiveness problem, shown in Figure 3 . Both versions of the NTL are scaled with λ (0.3 unless mentioned otherwise) and added to the regular CE loss:EQUATIONNote that both versions of the NTL shall be 0 for all non-numerical tokens. By changing the p-order in NTL-MSE, different L p -norm losses can be obtained (e.g., NTL-MAE). Huber loss is also compatible. In Appendix A.2, we provide pseudo-code for both versions of the NTL.We use a T5 backbone [11] (Appendix A.3) for our experiments and extend it with both versions of the NTL and the Regression-Transformer tokenization scheme [1] , due to its flexible encoder-decoder architecture and its success in various natural language processing tasks.Regression Transformer (RT). The Regression Transformer [1] tokenizes numbers on digit level, considering both the position and value of each digit. Since standard learnable embeddings may not adequately preserve the inherent structure of numbers, it leverages an inductive bias to account for the relative proximity of the numbers through numerical encodings, further explained in Appendix A.5.xVal encoding and decoding scheme. The xVal method [7] encodes real numbers using a single[NUM] token multiplied by its numerical value. For decoding (see Figure 1 ), a number head predicts the value while the token head outputs the sequence, replacing [NUM] during inference. However, this scheme is incompatible with T5 (see Appendix A.6). We thus use the xVal encoder and masked language modeling in our experiments.Integration of the Number Token Loss Both versions of our proposed NTL, depicted in the right panel of Figure 1 , can be integrated into any model that treats numbers as clearly separated tokens of single digits by applying it as an additional loss term. Therefore, we adapt the tokenization scheme of the standard T5 model to tokenize all numbers on the digit level to make it compatible with the NTL. As RT already tokenizes numbers on digit level by default, we can integrate the NTL without any changes. Integrating NTL into xVal is not feasible, as xVal encodes every number with the same token. Moreover, xVal already uses both MSE and CE loss. 
training objective:: 
 The standard 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 [3] instead of a [2] token will not generally induce lower loss than a [9] . This problem has been surprisingly neglected and is the focus of this work.Here, we aim to equip LMs with better inductive biases to handle combinations of textual and numerical data, such as math word problems or scientific datasets. In particular, we propose two versions of a regression loss on number tokens that respect numerical proximity (cf. Figure 1 right) and can be effectively combined with regular CE. The first version of this loss computes the Mean Squared Error (MSE) between the sum of the predicted class probabilities, weighted by their respective numerical token value, and the numerical token value of the label. 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. We integrate these improved training objectives with existing solutions for tokenization and embedding, in particular the Regression Transformer [1] . We evaluate all methods on a subset of the mathematical-question-answer dataset from DeepMind [12] .Prior art for joint language-number modeling suggested the use of verifiers [3, 10] , calculators (typically: Python interpreters), or chain-of-thought (CoT) reasoning [19] to yield improved performance in Large Language Models (LLMs). We argue that all such strategies avoid the fundamental, underlying problem (i.e., number representation in LMs is poor) by reformulating the task, trying to correct answers a posteriori with calculators, or using significantly more compute (CoR). Therefore, we herein intentionally attempt to improve a classic, relatively small encoder-decoder LM with up to 220M parameters, namely T5 [11] .2 MethodsTo test the mathematical capabilities of the methods, we use a dataset with more than 25 million samples from the mathematical Q&A dataset from DeepMind [12] . The dataset comes with two sets of tests: interpolation tests, one for each type of question occurring in the training set, and extrapolation tests, which measure generalization along various axes of difficulty beyond that seen during training. We provide more information about the dataset in Appendix A.4. We evaluate all five models on the two test sets of this dataset and report the accuracy (how often the model predicts the number exactly), as well as the Mean Absolute Error (MAE) and the R 2 -score. Since the dataset is skewed with some very high values, we perform a log 10 transformation on the predicted and ground truth numbers before calculating MAE and R 2 -score.All experiments except the one with xVal are built upon the T5 implementation and language modeling trainer based on the Hugging Face transformers library [16] . We use the T5-base model as a pretrained base for our respective models. All models were trained for approximately one million steps with a batch size of 32 over a period of approximately 3 days. More details on the models' training hyperparameters can be found in Appendix A.7. The results can be seen in Table 1 and Figure 4 . They show that vanilla T5 clearly benefits from both our loss variants. Indeed, accuracy 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. This confirms our hypothesis that number representation in LMs can be effectively improved through a minor, architecture-agnostic modification of the loss function. The RT consistently surpasses vanilla T5 on interpolation, however no further benefit was found by augmenting RT tokenization with NTL-MSE, potentially due to the custom number embeddings conveying numerical proximity. The limited performance of xVal [7] is explained by the extensive range of numbers in the used dataset. The dynamic range of xVal is limited due to the combination of its scaling of the number token embeddings and the pre-layer-norm in the backbone [17] . As a result, the effective number range of xVal is limited to [-5, 5] . To take this into account, we scale our dataset for xVal with log(1 + x). However, this means that large numbers can no longer be adequately distinguished by the model, as their embeddings become very similar. 
conclusion: 
 We introduced the Number Token Loss (NTL) for LMs to enhance their ability to handle numerical data by considering the numerical proximity between tokens. Our experiments unambiguously demonstrate the effectiveness of the NTL-WAS loss. This confirms our hypothesis that number representation in LMs can be effectively improved through a minor, architecture-agnostic modification of the loss function. By augmenting the standard CE loss with NTL, we provide a simple yet effective method that integrates seamlessly into existing architectures without requiring additional computational overhead. Experiments on the DeepMind Mathematics Dataset demonstrated that NTL significantly improves numerical reasoning, especially in models without specialized numerical embeddings. This approach offers a practical solution for enhancing language models in numerically rich domains, paving the way for more accurate and reliable applications in mathematics and science. 
 
 Appendix: A. return ntl 16: end functionThe T5 model is built upon the Transformer architecture [15] , 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 outputs and the previously generated tokens. The model can be trained with both Masked Language Modelling (MLM) [9] and Causal/Auto-Regressive Language Modelling (CLM) [4] , whereby we chose to use CLM.To test the mathematical capabilities of the methods, we use a subset of the mathematical questionanswer dataset from DeepMind [12] . The dataset was generated synthetically and therefore contains limited linguistic variability, but is sufficient for our 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 of the model and still leaves us with a large enough dataset of ∼26 million samples. The Regression Transformer [1] preserves the inherent structure of numbers by inducing information on relative proximity through numerical encodings that are set deterministically for all tokens. For every combination of a decimal place and digit value, a 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 at dimension j for numerical token t v,p with value v and decimal place p can be determined by NE Float (v, p, j) = (-1) j • v • 10 p j + 1 . 


Decision: The reviewers agree that this paper presents a valuable contribution for the MATH-AI community.

Reviews: review: title: title: Novel Loss Function for Numerical Data
main_review: main_review: This paper introduces a new loss function that works specifically on numerical tokens while ignoring linguistic tokens. It is then combined with other standard loss functions to compute the total loss. This is evaluated on a question-answer mathematics dataset. The results show a clear benefit when added to a T5 model. This new loss function is unique and addresses various issues in LMs when dealing with numerical data. There is a clear benefit in that it addresses proximity between numbers not considered in other loss functions.

It would have been nice to see more experiments done on other datasets. For example, the MATH dataset (https://openreview.net/forum?id=7Bywt2mQsCe) is quite common in this field, and assessing this new loss function on this dataset would be interesting as well. I would also like to understand more thoroughly the downside to adding NTL to the regression transformer. MAE actually increases when NTL is added, why do you think that is? I understand in line 113 the explanation for why RT does not benefit from NTL, but it seems to actively be getting worse on both datasets.

Overall, a very interesting paper and this shows a clear benefit of how LMs can process numerical data more efficiently.
title: title: Review on Augmenting Language Models with Regression Loss on Number Tokens for Arithmetic Reasoning
main_review: main_review: The study describes an important step in overcoming the well-known vulnerabilities of language models in processing numbers. Number Token Loss proposes such a task as well by introducing a loss function based on the distance between numerical predictions which may improve on tasks that require arithmetic reasoning but where the standard cross-entropy loss has proven ineffective. The experimental results obtained high evaluation scores on the DeepMind Mathematics Dataset leveraging better numerical performance. Nevertheless, the paper is in need of important modifications particularly with respect to the methodology and the technical aspects of the work. There are some shortcomings in this respect: the issues of NTL usage in various encoding-decoding architectures – Regression Transformer, xVal, etc. – are addressed too briefly. In addition, a broader experiment on other databases with a more thorough investigation of the benefits and weaknesses of the proposed solution should be conducted to support the paper arguments. However, enough of these concerns were raised and the main concept is still under-researched; with more workload it is conceivable that this work will be of great importance by its results.
title: title: The paper introduces Number Token Loss as an augmentation to Cross Entropy Loss to enable the feasibility of capturing number proximity relationship. The paper evaluates the method on T5 model and shows reasonable improvements
main_review: main_review: *Quality*   

The paper uses deepmind math dataset and T5 open source model. The evaluation is limited to methods on T5 rather than a comparative analysis across llm and also the metrics MAE, R2 and Acc might not be sufficient to cover all the different reasoning tasks in deepmind math dataset. Despite the limitation I still think the current evaluation is still very helpful to improvise and explore more in this direction so not penalizing this limitation

Though the model and dataset are open source and replicating Number Loss Token is feasible , reproducibility is dependent on the sampling of dataset

*Clarity*  

The paper is well written ,organized , structured and delightful to read. 

*Originality*  

Biggest missing piece was Related work. Though the number token loss mentioned in the paper is a new approach , there are various methods in which custom loss functions for numerical representation might have been tried before

*Significance of this work*  

This paper has very good learning and novelty to proceed in the future directions. Adapting the method in other LLM based models and evaluating would be more fruitful. Also its important to address issues at token level rather than post training like mentioned in the paper and this work would help give the community more insights and encouraging results to research further

*Pros*  

Reproducible with certain caveats  

Addresses a significant problem with LLM 

*Cons*  

Dataset Sample limitation  

Compute efficiency claim across any LLM is questionable  

Not evaluated on other LLM models like LLAMA to generalize if NTL might work
title: title: Promising Approach with Practical Applications, but Needs Improved Formatting and Justification of Results
main_review: main_review: ### Summary and Final Decision
**Rating: 6 (Marginally above acceptance threshold)**

This paper introduces Number Token Loss (NTL), a novel method aimed at improving numerical reasoning in language models by addressing limitations in cross-entropy loss for handling numerical tokens. The authors demonstrate that NTL improves performance, particularly in interpolation tasks, and the method is easy to integrate into existing models, making it highly practical for real-world applications like finance and scientific data processing. However, improvements in formatting (e.g., figure placement and reducing excessive baseline explanations) and including variance in the results to justify the performance gains could push this paper to a 7 or 8.


### Strengths:

1. **Simple Idea**: The introduction of the **Number Token Loss (NTL)** is an elegant and simple solution to a complex problem—handling numerical tokens more effectively in language models.
  
2. **Clear Objective**: The paper has a focused goal—improving numerical reasoning in LMs—which is well-defined from the outset and directly addressed through the proposed methodology.

3. **Significant Results**: The experimental outcomes clearly show that NTL improves arithmetic reasoning, particularly in tasks involving numerical interpolation. This demonstrates the practical benefits of the proposed method.

4. **Practical Application**: The method is easy to implement and integrates seamlessly into existing architectures. This makes it highly practical for real-world applications, particularly in domains where numerical accuracy is crucial (e.g., scientific and financial data processing).

---

### Areas of Improvement:

1. **Formatting**: **Figure 1** starts the paper abruptly, making the flow feel disjointed. Additionally, **Figures 2 and 3** take up too much space, and their contribution is minimal beyond explaining the non-injectiveness of NTL (i.e., cases where the predicted average can match the target despite errors), which is already well covered in the text.

2. **Failure Mode of NTL**: The authors claim that the **non-injectiveness** of NTL is rare, but this issue may be more prevalent, especially in larger datasets. However an angle which is not explored is that the combination of **NTL and cross-entropy** can help with early convergence, and cross-entropy can address spurious cases as the model refines its predictions. This balance should be discussed more explicitly.

3. **Lack of Edge Case Exploration**: The paper could explore more **edge cases**, such as predictions where token-level accuracy is close, but the numeric difference is large (e.g., predicting 129 instead of 123 or 133). In such cases, NTL may perform worse than cross-entropy. Expanding this discussion would provide a fuller understanding of NTL's limitations.

4. **Excessive Explanation of Baselines**: The detailed descriptions of methods like T5 and xVal take up too much valuable space. A brief summary with references would suffice, freeing up room for a deeper exploration of the novel contributions.



### Criteria Ratings:

1. **Relevance to Journal**: 5/5  
   - The paper aligns perfectly with the journal's focus on improving mathematical reasoning in LMs, particularly in the context of arithmetic and numerical tasks, which is highly relevant.

2. **Clarity of Language and Structure**: 3/5  
   - While the paper's objectives are clear, the structure suffers from some abrupt transitions and excessive details on baseline models. The formatting of figures and the placement of certain explanations could be improved for better readability.

3. **Experimentation Procedure**: 3/5  
   - The experimental setup is well-structured, but the results could be clearer. The performance differences are minimal, and without averaging across multiple runs, it’s difficult to determine if the improvements are real or just noise. A more thorough analysis would help distinguish meaningful gains, and **knowing the variance** in these results would also be beneficial to understand the consistency of the improvements.

4. **Innovation and Original Contribution**: 4/5  
   - The introduction of the **Number Token Loss (NTL)** is a novel and impactful contribution. However, the lack of exploration of edge cases and the challenges with non-injectiveness leaves some room for further innovation.

5. **Reproducibility**: 4/5  
   - The paper provides sufficient detail on model configurations, datasets, and hyperparameters, supporting solid reproducibility.




### Line-by-Line Review:

---

**Line 15 (Introduction)**:  
*Figure 1 is placed immediately after the introduction.*

**Comment**: The placement of **Figure 1** disrupts the flow of the introduction, starting the paper with a figure rather than guiding the reader through key concepts. This figure should be moved to a later section where the methodology is discussed, allowing the introduction to flow more naturally and provide proper context before jumping into figures.

---

**Lines 47**:  
*"We herein intentionally follow a simplistic scheme of attempting to improve a classic, small encoder-decoder LM, namely T5 [11]."*

**Comment**: Instead of using "small model," it would be clearer to specify the actual size of the model (e.g., "a 220 million parameter model"). This adds precision and avoids ambiguity, as "small" is relative in the context of language models.

---

**Line 51**:  
*"We compute the weighted sum of the class probabilities with the numerical value of each token."*

**Comment**: Use the term **"predictions"** instead of **"values"** to distinguish between the model's outputs and the ground truth labels. This will clarify the difference between the predicted and true values in the context of NTL.

---

**Line 53 (Equation 1)**:  
**Comment**: Be more specific that **"V"** refers to the **numerical tokens in the vocabulary** rather than embeddings. Clarifying this point will make it clear that the operation is being done on numeric tokens rather than embedded representations.

---

**Lines 59-60**:  
*"Even though such cases are empirically rare because the softmax emphasizes differences in the logits, in practice, we always combine NTL with a regular CE loss."*

**Comment**: Rephrase to acknowledge that such cases may be more frequent than implied. Consider:  
"While the authors claim these cases are rare, they may occur more frequently in practice. Combining NTL with cross-entropy loss helps mitigate this issue. NTL may accelerate early convergence by providing extra loss signals, but spurious cases (e.g., predicting an average value of 4 when the label is 4) can be corrected as the cross-entropy loss continues to refine predictions, since its value will not be reduced in such instances."

---

**Figures 2 and 3**:  
*Line 63 - Discussion of Figures 2 and 3*

- **Figure 2**:  
  **Comment**: This figure takes up too much space for a relatively small point that has already been explained in the text. It visually reinforces an obvious concept—the limitations of non-injectiveness in NTL—but doesn’t provide additional value. Reducing its size or removing it would improve the paper's readability.

- **Figure 3**:  
  **Comment**: This is the third reiteration of the same concept, and together with Figure 2, takes up about a huge chunk of the paper's space. While it's important to clarify NTL’s non-injectiveness issue, this level of detail is excessive. Compressing the discussion and illustrations of this concept would allow more room for novel contributions and further analysis.


---

**Lines 67-71**:  
*"For every combination of a decimal place and digit value, a corresponding numerical token is added to the vocabulary. For instance, the number 11.4 is tokenized to [1_1, 1_0, 4_-1]."*

**Comment**: The explanation of the Regression Transformer’s tokenization scheme is clear but takes up unnecessary space. Since this is a baseline method, a concise description and reference to the original work would be more appropriate. This would free up space for a deeper discussion of the novel contributions of the paper.  Alternatively consider moving this into a separate Related Work section.

---

**Lines 78-89**:  
*"The xVal method [7] introduces a specialized numerical encoding and decoding scheme that represents any real number using just a single [NUM] token..."*

**Comment**: The explanation of **xVal** is detailed but overly long. A more concise summary, combined with a reference to the original work, would allow more space for analysis of the limitations of xVal in this specific context, particularly concerning scaling issues in the T5 architecture. Alternatively consider moving this into a separate Related Work section.

---

**Line 110 (Table 1)**:  
*Performance metrics with accuracy, MAE, R2 scores.*

**Comment**: The differences in performance between models are minimal, making it difficult to discern meaningful improvements. Including variance or standard deviation in the results could strengthen the justification for improvements and demonstrate whether these improvements are consistent and significant or merely due to noise.

---

**Line 123 (Conclusion Section)**:  
*"This approach offers a practical solution for enhancing language models in numerically rich domains."*

**Comment**: The conclusion could be enhanced by discussing future work. For instance:  
"A deeper exploration of edge cases where token-level predictions are close but numerically different (e.g., predicting 129 vs. 123 or 133) would help clarify the limitations of NTL. Moreover, methods like the Regression Transformer (RT), which already handle numerical proximity well, could be integrated with NTL. Future work could explore combining RT’s strengths with NTL, potentially improving performance beyond what either method achieves alone."