Interleaving Text and Number Embeddings to Solve Mathemathics Problems

Authors: marvin  alberts
gianmarco  gabrieli
irina espejo morales
ibm  research
Venue: NeurIPS_cc_2024_Workshop_MATH-AI
Published: 2024-09-24 11:30:38

Abstract: Integrating text and numbers effectively is a crucial step towards enhancing Large Language Models (LLMs) capabilities in assisting in scientific tasks. While most current approaches rely on discrete tokenization of numbers, for instance, conversion to scientific notation or base 10-decomposition, a recent approach proposed a continuous numerical encoding as an inductive bias. In this paper, we build upon this approach by introducing more expressive numerical embeddings. Our method addresses key shortcomings, including the elimination of numerical artefacts and the ability to handle a wide range of magnitudes without clipping. Our work presents two key contributions. First, we employ an MLP to assign distinct directions in the embedding space to different numbers. Our second contribution is the introduction of a routing layer that differentiates between numerical and text embeddings. We hypothesise that this combined approach enables the model to distinguish between text and number distributions while maintaining its capacity for arithmetic operations. Using only a 45 M parameter encoder-decoder architecture our method achieves a R2=0.9988 over a wide range of magnitude (10−3, 10^8). In addition, we empirically observe a reduction of the numerical artefacts and biases observed compared to the baselines

Main paper: 
 introduction: 
 Transformer-based autoregressive language models have revolutionized machine learning over the past years [1, 2, 3] . Originally designed for neural machine translation, these models have found applications across diverse domains, including natural language processing, image analysis, and mathematical problem-solving [4, 5, 6, 7] . Despite their versatility in natural language tasks, accurate mathematical reasoning remains a challenge [8] , especially when Chain-of-Thought prompting or a scratchpad are not employed [9] . The limits of the transformer architecture to perform mathematical computations and predictions from first principles have been extensively explored in the literature [8, 10, 11, 12, 13] . However, the majority of approaches leverage digit-by-digit, scientific notation or base 10 formats to encode and decode numbers [14] . This discretization of numbers leads to inherent prediction artefacts, especially for unstructured and non-synthetic numerical input data. For this reason, XVAL has proposed to introduce continuity as an inductive bias [15] , to reduce the impact of discontinuous tokenization. Despite the improved performance in regression tasks and the reduction of artefacts, XVAL limits the magnitude by normalising the numerical inputs. In this work, we built upon XVAL incorporating continuity as an inductive bias, and pose the question: Does increasing the expressivity of the numerical embeddings, by allowing each number to attend to its surrounding text, improve the performance on number prediction in mathematical questions? For instance, it is intuitive that the number "2" should attend to words like "eggs" or "apple" and the number "0.0001" to words like "infinitesimal" or "fraction". The second question we aim to address: Can a routing layer (as in Mixture-of-Experts) differentiating between text and numbers induce structure in the prediction of the model to improve the generation of numbers?Regarding number tokenization, a large number of LLMs use variants of Byte Pair Encoding (BPE) [16] to tokenize numbers in the same way as text. As shown by [8] and [12] this introduces a limit to the number understanding of these models. Charton 10 demonstrates that transformer models can perform complex matrix operations and proposes a scheme tokenizing numbers as a series of tokens for sign, mantissa and exponent with varying levels of precision. As noted by [15] , this approach does not take into consideration the continuous nature of numbers which is important in scientific domains, and instead, the authors introduce an architecture called XVAL which bypasses the need for number tokenization but not for number preprocessing and assigns a to each number a single numerical embedding modulated by the magnitude of the number itself. In the arithmetics applications [11] that concern this paper [14] , the work by [17] shows that by adding positional encodings to digits it is possible to generalize out of the maximum training length of digits. 
our approach: 
 In this work, we introduce Multimodal Decoding (MMD) which builds on the XVAL [15] architecture. XVAL embeds numbers by introducing a number token, <num>. The embedding of this number token is multiplied by the actual value of the number, effectively scaling the embedding of the number token. To decode numbers XVAL introduces two heads, one for tokens and one for numbers. During decoding if the <num>-token is predicted the model embedding is routed to the number head predicting a number.Our architecture is similar. We also embed numbers separately but instead of scaling an embedding, we use an MLP allowing each number a distinct embedding direction learned by attending the surrounding text and as such providing more expressivity (see Figure 1 (a) ). Another major change consists of the introduction of a routing layer inspired by Mixture-of-Experts (MoE) models [18] . Instead of relying on the text head to predict the <num>-token and then passing the model embedding to the number head, the routing layer predicts whether a given model embedding is routed to the text or number head (see Figure 1 (b) ). This approach is motivated by the following insights: i) An embedding for a given number will be attending to the text that surrounds it, and different numbers will attend to different words; ii) The routing layer will force text and number embeddings to have a different distribution so that arithmetic calculations can be performed with numerical embeddings inside the transformer blocks; and iii) Instead of hard-coding an inductive bias for numerical embeddings with this method allows for finding optimal embedding distributions that might otherwise be missed. 
experiments: 
 To evaluate the efficacy of our architecture, we conduct experiments on two progressively more complex tasks. We begin with a relatively straightforward problem: predicting numerical answers to arithmetic problems. In this task, we constrain the arithmetic problems such that all answers are numbers (Numbers Only).We then extend our evaluation to more challenging math problems. In this case, we introduce problems that require either numerical or textual responses, or a combination of both (Text and Numbers).Since we propose fundamental changes on numerical embeddings and different losses, we selected a simple dataset, the Mathematics dataset [19] , see Appendix C. We use two versions of the dataset. One containing only numerical answers and the second containing prompts with both numbers as text as possible answers. We train a total of three variants of our model. The first with a MLP to encode numbers and similarly a one-layer MLP as the number head. In the second variant, we predict the log-transformed numbers and in the third variant, the number head consists of a three-layered MLP. We compare these models to three baselines: One transformer model trained with a BPE tokenizer, i.e. numbers are tokenised digit by digit, a word level tokenizer and lastly XVAL. All three of the variants of the model outperform the baseline on MAE, RMSE and R 2 (see Table 1 ). Especially encouraging is analysing the log-log plot of the predictions vs the targets (see Figure 2 ). Here we observe specifically for BPE and word-level artefacts with the model predicting multiple magnitudes off the target. On the other hand, we also observe failure modes in our architecture. While the predictions and targets track well starting from around 10 2 , the model fails to predict small values. This problem is addressed by log-transforming the numbers yielding a log-log curve closely following the identity.We then extend the evaluation of our method to mathematical problems beyond pure arithmetic, encompassing tasks which comprise interleaved text and numbers in various problems (e.g., algebra). We leverage the same methods and baselines reported in Section 3 and report the results in Table 2 .In this case, our methods obtain performance metrics that are several orders of magnitude better than the baselines, demonstrating the challenge of text-only models to process interleaved sequences. In particular, log-transformed numerical representations result in improved accuracy compared to the other two variants, with the only expection of the R 2 score for the base MMD version. Moreover, the results show that logarithmic pre-processing is pivotal for differentiating text and numbers, with an F1-score on the modality classification that is 0.11 higher than the other methods. As expected, the results are generally worse than in Section 3, due to the higher complexity of the interleaved multi-task dataset. 
conclusion: 
 In this paper we investigated two strategies to interleave text and numbers as input to a standard encoder-decoder transformer architecture. The first strategy allows more expressive numerical embeddings by assigning a distinct vector space to different numbers. The second strategy introduces a routing layer that differentiates between text and number tokens. This serves to create a desirable structure in the embedding space differentiating between text and numbers without explicitly introducing any inductive bias. From our experiments on predicting the numerical outcomes of arithmetic computations, we observed that our approach exhibited the most balanced set of metrics and when comparing predicted vs true values, our method was the closest to the identity showing consistent performance along magnitudes ranging from 10 -3 to 10 8 . We also note that even if a baseline performed well for a particular metric, the true vs predicted plots reveal artifacts and undesirable biases. The main limitation of this study is the simplicity of the arithmetic's dataset, which we intend to address in future work by incorporating more complex datasets and perform experiments against a wider set of numerical encodings and architectures in the literature.See Figure 3 for a graphical description of the different baselines and our methods. All experiments start with the same question and answer pair but text and numbers processed differently. For the baselines BPE and Word-level in yellow, the routing layer is inactive because numbers are processed as text only.  
 
 Appendix: Model: The backbone structure is a standard encoder-decoder architecture with 4 layers in the encoder transformer block and 4 layers for the text decoder block and a 2-layer fully-connected MLP for the numeric head. The activation functions are Gaussian Error Linear Units (GELU). All experiments were performed using a single NVIDIA A100 40GB GPU. Models were trained for 50 epochs by setting the learning rate to 1e -4 for all the experiments.Metrics: The metrics we use to evaluate number predictions for the result of arithmetic calculations are Mean Absolute Error (MAE), Mean Squared Error (MSE), Mean Relative Error (MRE), Median Relative Error (MedRE) and the R 2 coefficient. Instead of reporting one of them, together they present a better picture of the distribution and biases of errors with respect to the range of the true answer.Examples of question and answer pairs of the Mathematics dataset [19] that were used for training of Numbers Only and Text and Numbers experiments presented in the paper. 



Reviews: review: title: title: Solve Math Problems by interleaving text and number embeddings
main_review: main_review: Summary: 
The paper addresses the challenge of integrating text and numerical data in large language models (LLMs). The paper builds on recent advancements that propose continuous numerical encoding as a more effective way to incorporate numbers into LLMs. The authors employ MLP to assign distinct directions in the embedding space to different numbers and propose a routing layer that differentiates between numerical and text embeddings.

Quality: 
The paper is technically sound.

Clarity: 
The paper is well-written and easy to follow. 

Originality: 
This paper introduces Multimodal Decoding (MMD) which builds on the XVAL. I think the authors make good claims in the method section. The illustration of the routing layer at the end of the paper is good.

Pros:
- Proposed two contributions that resulting in an improved performance in the evaluation phase regarding F1-SCORE, MAE, RMSE, R^2 compared with other methods.

Cons:
- it would be helpful to include significant tests when comparing with other methods.
title: title: New multimodal model structure achieving better results on simple arithmetic dataset
main_review: main_review: The paper proposed a multimodal model, MMD, based on the XVal model, to solve the arithmetic problems. The two new strategies are: (1) individual embedding for each number and (2) an additional routing layer. The MMD model achieved better result for the Mathematic Dataset. The paper is well written with sound logic. The authors noted that the dataset used is simple and planned to test on more complex dataset in the future. Besides testing different dataset, I think the paper can benefit from:
(1) Experimenting on only the embedding change and only the routing layer - is the improvement coming more from one of them or is it a combined effect?
(2) Explaining why Multi-MLP MMD is better for numbers only problems and MMD-Log is better for text-and-numbers problems
(3) Listing details on model structure and training process such as embedding size to increase reproducibility
(4) Discussing computation trade-off - does adding a routing layer increase training time? does giving each number its own embedding increased overall model size?
(5) Providing sampled example where the XVal result is not as accurate as MMD
title: title: Approach aims to address limitations of XVal but the value of technique isn't evident from paper.
main_review: main_review: **Summary:** The paper presents a technique called Multimodal Decoding (MMD) which extends the numerical encoding technique from [XVal](https://arxiv.org/abs/2310.02989). MMD uses an MLP to encode the number and at decoding time adds a lightweight routing layer that acts as a soft classifier on the transformer representation and routes the embedding between text / numeric head. The numeric head directly predicts a number given an internal representation. Unlike the previous approaches that used a decoder-only architecture, the paper explores their technique on encoder-decoder architectures.

**Strengths:**
* **More Expressive Variant of XVal:** XVal represents numbers using a dedicated <num> token with input representation being <num> token scaled by number value and output representation being inferred via a numerical head for any <num> tokens generated in output. The current paper’s MMD approach makes this approach more flexible in two ways. First, it encodes the number using an MLP (with optional log application) rather than a single scaled embedding vector. Second, it gets rid of the intermediate token generation stage in decoding and instead uses a routing layer to optionally but directly activate the numerical head on intermediate representation.

**Weaknesses:**
* **Encoder-Decoder Architecture choice doesn’t seem well motivated:** Standard LLMs including XVal that this paper extends use decoder-only architecture. This paper uses Encoder-Decoder architecture (motivation isn’t specified) which prevents a cleaner comparison of the numerical embedding technique and limits the applicability of this paper.
* **Lacks Details:** Multiple relevant details seem missing from the paper e.g. a) the specifics of routing layer b) the specifics of the numeric and text-numeric dataset like size, train/eval splits etc c) the loss functions for which numerical head and the routing layer optimizes.
* **Technical Contribution seems Incremental:** The paper’s contributions w.r.t XVal while well-motivated are still quite incremental. The paper doesn’t address the key limitations of XVal. There is no comparison against open-source LLMs even at small scale (this was a limitation of XVal) which is crucial because despite the perceived issues in current digit based tokenization, it is undeniable that models using these tokenization techniques do fairly well on basic and even competitive math tasks. Which highlights that tokenization schemes cannot be evaluated in a vacuum and instead need to consider both the size of the transformer and the scale / diversity of the training input. XVal however highlighted their focus on continuous / large scale computations with focus on scientific datasets (something that LLMs do struggle with) but this paper doesn’t evaluate on those datasets. *So the value of MMD is not clear.*
* **Experiments do not seem well motivated:** The MMD technique is evaluated on a subset of data from [Saxton et al](https://arxiv.org/pdf/1904.01557) (which contains basic school level math problems). As also highlighted in the previous point, the value of these experiments is not really clear as today’s LLMs even at small scales like [Llama 1B and 3B](https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices) already excel at harder version of these problems (e.g. [Hendrycks MATH](https://arxiv.org/abs/2103.03874)) despite the classic digit level tokenization.

**Questions:**
* **Is Numerical head trained using MSE ?** It wasn’t clear from the paper how the numerical head is trained. Does it use MSE similar to XVal ? Also do you have any explicit losses for training your routing layer ?
* **Routing Layer Specifics:** Do you have any explicit losses for training your routing layer ? What are the architectural specifics of the routing layer (is it like a one layer binary classifier ?)
* **Input Normalization:** Besides applying log in the log-MMD variant, do you have any other normalizations on your numerical inputs (within the network or in the training datasets). How do you handle numbers of large magnitude ?
* **Why Encoder-Decoder Architecture:** XVal used GPT-2 and other standard LLM use the decoder-only architecture. Is there some specific reason, your technique uses encoder-decoder ?
* **MAE / RMSE for Text-Numeric Dataset:** What do these metrics mean for predictions with both text and numbers ? Are you only computing it on numbers ?
* **Variance across reruns:** Given that MAE / RMSE are quite sensitive to magnitude of answer, did you observe high levels of variance across reruns of your approach (and if operating in log space reduced it) ?

**Suggested Corrections:**
* **Corrections to Figure 2:** Correct WorldLevel -> WordLevel. Baselines are in the top row and MMD is in the bottom row, caption says the opposite.
* **Rendering Issues:** Not sure if there are specific image formats used in this pdf but for some reason this paper wasn’t rendering properly for me with the reader lagging a lot in scrolling (tried multiple browsers, Preview, Adobe and to no avail). The issue was specific to this pdf. Wanted to highlight in case this issue was observed by / other reviewers. Apologies in advance if I may have missed any key details in my review above.