Paper: xVal: A Continuous Number Encoding for Large Language Models (ICLR-2024)

Reviewer comment (Novelty):
Not well motivated. The paper says it's LMs have historically struggled to solve simple arithmetic problems but somehow many Chain-of-Thoughts paper contradicts the claim. There is no discussion and not literature review enough for this part.

Author response:
Indeed, a number of recent work have explored methodologies that force LLMs to use Chain-of-Thought reasoning in a “scratchpad” setting and have proven effective for improving numerical predictions of LLMs tasks like arithmetic or mathematical word problems. We have now added this information and relevant citations to the introduction.

However, a pre-requisite for these methods to be applicable is that Chain-of-Thought reasoning for solving the given task (or a similar task) must be available during train time. In our work, we are motivated by the need to improve numerical analysis of datasets such as those seen in many scientific contexts. For these tasks, such chain-of-thought reasoning is not known -- indeed often the goal is to discover the relationships in the first place -- and therefore a step-by-step scratchpad would not improve the results.

Reviewer comment (Novelty):
The method itself is not very interesting.

Author response:
We believe that our findings and results would be of interest to the community as evidenced by the comments from other reviewers. In particular, we show state-of-the-art out-of-distribution generalization performance in a number of tasks.

Reviewer comment (Evaluation):
Evaluation is rough, what are FP15, P10 and so on? No clear elaboration on this.

Author response:
These tokenization schemes (FP15, P10 etc) are described in detail in Table 1 (including an example for clairty) and at the beginning of Section 3 titled **Comparison with other number encodings**.

Reviewer comment (Presentation):
Page 2 is in low-resolution. A not ready draft IMO.

Author response:
Thank you for pointing out this accidental glitch in the rendering process. We have corrected this in our latest version.

Reviewer comment (Experiments):
I believe at least the paper need to show empirically why the task it difficult.

Author response:
In our experiments, we compare xVal with existing text-based number encoding schemes that are currently SOTA on the type of numerical analysis tasks that we are interested in. We demonstrate that these other schemes have several notable weaknesses, especially in terms of robustness to outliers and performance in out-of-distribution generalization.

Reviewer comment (Experiments):
In experiment setting, the proposed method is only evaluated in the supervised training setting. It is unclear on the impact of the proposed method on pre-training tasks.

Author response:
Thank you for bringing this up. While our training regimen is self-supervised masked language modeling, we primarily evaluate the performance of the model on values that are directly present in the training corpus. We suspect that this is what the reviewer is referring to as "evaluated in the supervised training setting". 

To more directly evaluate the model in a setting where the output is not explicitly present in the input, we performed two new tests. The first a fine-tuning task and the second a roll-out task.

* We fine-tuned the temperature forecasting model on a binary downstream task, specifically we use the fine-tune the pretrained model to predict if the first reporting station is located on water or on a landmass. In principle, this problem is solvable by looking at the latitude & longitude of the station, but because we fine-tune on only 500 samples, there is not enough information to accurately learn the mapping between latitude & longitude and water coverage over the earth. Therefore, the model needs to use other information (e.g. temperature variation etc.) to make this inference. We find that xVal outperforms the other models, achieving a ROC AUC of 0.620 ( vs .600 and 0.580 for P10 and FP15 respectively). We expect that the longer context length of P10 is adversely affecting its fine-tuning behavior. We report these findings in Sec. B.1.3.

* In the temperature forecast problem, we rolled-out the temperature forecast on 5 extra time-steps to predict the future development of the temperature across all sites. This is similar to text generation in language models that are only trained to predict the next time-step. We find again that xVal outperforms the other encoding schemes. This is not surprising since xVal has the best one-step prediction performance.

Reviewer comment (Experiments):
How the training hyper-parameters are configured and why different encodings have different configurations.

Author response:
Thank you for asking this clarifying question. To determine the hyper-parameters of our runs, we performed a coarse grid search for the learning rate (with a linear grid in log space). We chose the best performance on a validation set and reported the results on an unseen test set. In many cases, the validation performance of different learning rates were very close and/or within variance of each other, but we still chose the configuration with the best performance. This is one cause of differing learning rates between different experiments. Another reason for differences between the optimal hyperparameters for xVal and the text encodings is the very different statistics of the input for these different models. Whereas the input to the text models is always a text embedding, the input to xVal cares about the numbers that appear in the text.

We have added this description to the supplementary materials Sec. B.

Reviewer comment (Presentation):
An issue with the rendering quality on page 2, affecting some figures, diminishes the readability of both the text and data points. The compromised legibility of axis labels and data points could potentially hinder the reader's comprehension and impact the overall impression of the paper.

Author response:
Thank you for pointing out this accidental glitch in the rendering process. We have corrected this in our latest version.

Reviewer comment (Experiments):
The experiments conducted in the paper exclusively utilize structured data in JSON format, focusing on scenarios like multi-digit multiplication and multi-operand calculations. While these experiments demonstrate the effectiveness of the proposed approach in these specific contexts, they do not adequately showcase the method's capability to understand real numbers in the broader context of natural language. This limitation may raise questions about the universal applicability and effectiveness of the proposed approach. Expanding the scope of experiments to encompass real-world language contexts would provide a more comprehensive evaluation of its capabilities.

Author response:
We agree that it would be instructive to demonstrate the use of our encoding scheme in more general settings. The goal in this current work was to demonstrate better performance when applying LLMs to the analysis of numerically dense scientific data from heterogeneous sources.

**On structure:** We note that even though our samples were formatted as a JSON, this formatting is not necessary for achieving the results in the draft. Violating this strict structure, for example removing structural components of the JSON (e.g. commas and brackets) or providing information out of place (e.g. putting a key in a different part of the JSON hierarchy) does not degrade the performance, so long as these were also seen during training. (In comparison, methods that rely on the strict structure of JSON would fail when the structure is violated.)

**Broadening the scope:** Demonstrating better number understanding in real-world language contexts would be the logical next step. However, training a language model with broad language understanding requires dramatically more computation time and power than the experiments run in the current paper.

Reviewer comment (Evaluation):
Could you please demonstrate how well the proposed method can handle situations where it needs to refer to previously mentioned real numbers in the context, ensuring these numbers remain unaltered? How does this embedding method impact a Language Model's capability to preserve real numbers in the given input?

Author response:
We thank the reviewer for this suggestion. We address this as follows.

**New Experiment:** In response, we have carried a new experiment demonstrating how a language model trained with xVal can indeed recall numbers. The experiment and explanations can now be found in Sec. B.6 of the Supplementary Materials, and we summarize them below:

We implemented a simple number lookup experiment where each sample is a combination of a dictionary and an example of a dictionary lookup. An example is “{d:1.53, e:-1.33, a:2.53, i:0.0232} e=-1.33”. We train this via self-supervised masked language modeling and then evaluate the performance of the lookup task on unseen samples. We find that with xVal the model can perform the number lookup with very good MSE on the target number, but often fails to recover the exact number to 3 digit precision. (We also train this task using FP15 and P10 and provide the comparison in Sec. B.6).

To see whether or not this performance generalizes to number ranges unseen in the training set, we evaluate out of distribution performance of the number lookup task. We again find that xVal generalizes better compared to P10 and  FP15. 


**Architectural Perspective:** From an architectural perspective, the model can access the numbers present in the context because of the residual connections which connect the input embedding directly to the final output of the transformer blocks. Therefore, in principle the transformer can make sure that this information remains available by leaving that subspace of the latent space unchanged.

Reviewer comment (Experiments):
How is the capability of the proposed numerical encoding scheme affected by extremely small or extremely large numbers? Is it able to maintain representation accuracy and robustness in the presence of such numerical extremes?

Author response:
We discuss this in Section 2.1 (“xVal: A Continuous Number Encoding”). As shown in Figure 2, the input dynamic range of xVal is indeed limited, particularly outside of the range [-5, 5]. Before training, we therefore preprocess our data to scale the values to fall within this range. 

We discuss the shortcoming of the limited dynamic range of our proposed method in the Conclusion section. We expect that that this shortcoming can be alleviated through more elaborate continuous encodings.  (Also see our response to Weakness #2 above.) However, we believe that by focusing on the simplest continuous encoding (i.e. our multiplicative implementation) we can highlight the core advantages - as well as disadvantages - of this approach more clearly.

Reviewer comment (Experiments):
While the continuous nature of the xVal embedding can obviously be an advantage in some domains, I'm not sure how well it would work in general. For example, tasks like summarization, or question answering, where numerical values such as years/dates/account numbers are not meant to be worked with in the sense of arithmetic or other mathematical operations but simply carried about may lose performance. That said, for domain specific applications (like science), this may not be an issue. A hybrid approach may also work (e.g., representing $7.4$ as "+ , 740, e-2, 7.4 * [NUM]").

Author response:
We agree that for a general language model a hybrid approach would be useful. As the reviewer suggests, an approach combining text-based encoding with a continuous encoding to leverage advantages of each method can lead to improvements on both. One can also envision an encoding scheme that treats different number types (integers, floats, etc) differently. 

In this work, we wanted to demonstrate the advantages of a continuous number encoding over text-based encodings in the setting of heterogeneous scientific data analysis. We leave the exploration of combining different tokenization schemes to future work.

Reviewer comment (Evaluation):
I'm not entirely sure why the runtime is so dramatically reduced. Is this due to the reduced length of each input/target and the vocab size? If the latter is a big factor, I'm surprised that what seems like a fairly small additional overhead of the softmax size for everything but FP15 setting would make that big of a difference.

Author response:
The reduced runtime of xVal is a combination of both fewer tokens and small vocabulary size. The bigger factor is the length of the tokenized input. Text-based encodings (other than FP15) have dramatically increased number of tokens compared to xVal and FP15  (in Sec. 3 around 8000 and 5000 tokens per sample respectively for P1000 and P10 compared to 1800 tokens for FP15 and xVal (see Table 1 for a summary of different tokenization lengths). This leads to a longer runtime for the encodings with higher number of tokens per number. In comparison, the overhead of FP15 compared to xVal is coming from the size of the vocabulary leading to large embedding and logit-head layers. Note that the relative increase in computational cost from the embedding and logit-head is larger for small transformer models. In our case, the embedding matrix of FP15, for example, is a large fraction of the total number of parameters of the model.

Note that in the runs with equal number of total tokens, for efficient parallelization purposes, the number of tokens seen per batch is still different for the different methods. This also leads to a difference in run-time, even when the total number of tokens seen during training is the same.

Reviewer comment (Experiments):
I'm a bit put off by the dependence on _parsing_ numerical quantities accurately, especially in messier settings. Curious if that posed any difficulties.

Author response:
In our code, the preprocessing step is implemented using a short regular expression to parse the input text and replace numerical values with the “[NUM]” token. This is done efficiently using the Re package and the processing adds little overhead. This simple preprocessing step would not be sufficient for a general purpose language model, especially if we want to extract the number types as suggested above. However, for such a use case, lightweight parts of speech tagging models can be used to efficiently and accurately tag numbers and number types (e.g. Jurafsky Speech and Language Processing Chapter 8).