Better autoregressive regression with LLMs via regression-aware fine-tuning

While this paper presents a strong approach, it addresses a relatively niche problem.

As we discuss in the introduction, we would like to point out that regression problems encompass a wide range of problems, including semantic similarity prediction (Cer et al., 2017), quality estimation (Kocmi & Federmann, 2023; Jain et al., 2023; Fernandes et al., 2023), and sentiment analysis (Zhang et al., 2024). Moreover, there has been a large interest in utilizing LLMs to regression problems, e.g. (Gruver et al., 2023; Liu & Low, 2023; Yang et al., 2023; Fernandes et al., 2023; Zhuang et al., 2023; Lukasik et al., 2024). Thus, we believe that it is a pertinent question as to what is the most effective approach to regression with LLMs.

The choice between autoregressive sampling and a predictive head largely depends on the specific task, making it less crucial to claim that your method is universally superior to both across all tasks.

Note that our approach is more grounded theoretically than predictive head when utilizing autoregressively pre-trained LLMs. In other words, it is well motivated for decoder based LLMs across regression problems.

The proposed method does not show a significant improvement over the predictive head. For instance, in Table 9, RAFT performs worse than the predictive head on the Music dataset.

Please see a common response to all reviewers “RAFT improvements over predictive head”.

Could you clarify the scenarios in which a predictive head performs better versus those where autoregressive sampling is more effective?

Can we ask for a clarification of what the reviewer refers to as “autoregressive sampling”: RAFT or regular autoregressive fine-tuning + decoding? We will address either interpretation below by contrasting the following three scenarios: predictive head, autoregressive fine-tuning + decoding and RAFT.

Autoregressive fine-tuning + decoding In our work, we prove autoregressive based approaches can be arbitrarily poor for any regression problem due to misalignment to the evaluation metric: see the newly added Lemma 1 expanding this theoretical understanding to autoregressive fine-tuning + vanilla decoding, and also the Lemma 2 treating autoregressive fine-tuning + MALI decoding.

Predictive head The predictive head deviates from the pre-training task, making it unclear whether it is well suited for finetuning decoder-based models. We add additional synthetic experiments to expand the understanding of the role of pretraining for decoder-based model finetuning, and further corroborate how when using pretraining checkpoint from the next-token prediction task, even on a simple synthetic regression dataset, predictive head underperforms RAFT; please see the common response to all reviewers “RAFT improvements over predictive head”.

RAFT We argue for several empirical and conceptual benefits of RAFT over predictive head. Please see “RAFT improvements over predictive head”.

Is RAFT the optimal choice in all situations, or are there cases where other methods might be preferable?

Our work argues that RAFT is the most theoretically grounded approach. Moreover, at step zero it corresponds to MALI, and thus, provides a reasonable performance even at zero shot. It also converges much faster than predictive head approaches even where there is no significant difference in the final performance. (e.g. see Table 12 and Figure 1).

it would be worth slightly rewriting the paragraph after Lemma 1

Thank you for the suggestion, we now incorporated it in the updated manuscript.

On the line 371, you are explaining that you are experimenting with a reduced training set for the STSB benchmark (using 1k examples instead of the whole training set). I understand that it is interesting to vary the amount of training data, but why are you only reducing the training set size on STSB and not on the other datasets?

We selected STSB as one example regression dataset. To additionally analyze the impact of the data size, we add additional datasets: MovieLens and a synthetic dataset, thus expanding the range of dataset sizes (see Table 7 in the updated draft for the summary of all datasets and their corresponding sizes). Overall, we confirm RAFT improves over baselines across dataset sizes. Please see a common response to all reviewers “RAFT improvements over predictive head” for details.

Main weakness: While your paper studies regression tasks using autoregressive models, you observe on lines 465-467 that when training from scratch, the predictive head methods work best. Since encoder-only models such as BERT are widely used for tasks mentioned in the introduction (e.g. sentiment analysis, regression, ranking), it would be valuable to include a baseline of a large fine-tuned RoBERTa model on the same tasks that you present in the paper. Your findings would remain interesting and relevant in case the RoBERTa baseline was beating RAFT, but I think knowing how RAFT compares to fine-tuning RoBERTa is very relevant to practitioners. Therefore, I would suggest adding the following experiments (with the same hyperparameter tuning as RAFT): *Prediction head regression over a mean-pooling of RoBERTa representation (over sequence length) *Prediction head regression over the output for the CLS token of RoBERTa model *Both variants with either frozen RoBERTa weights, as well as unfrozen weights

Please see our answer to this question in the common response to all reviewers: “Additional baselines”.

There could be more practical details on the fine-tuning process. Are you updating all the parameters of the model or did you freeze some of them (eg. update only the last linear layer)? Which optimizer are you using? If you are using Adam, what values of the betas, weight decay, and epsilon are you using? What batch size are you using? How long does one training run takes on average? How many machines did you need? The answers to these questions are important to help other researchers reproduce your results.

We added a clarification regarding the settings in Appendix C. We update all parameters during the fine-tuning. We summarize specific settings below:

• DROPOUT_RATE=0.1.
• BATCH_SIZE=16.
• We train for 200k steps and select the best step using the held out validation set (see Table 7 for details on the train/test/validation splits).
• We use a constant learning rate schedule. We select the learning rate value over the validation set from the set of values: 1e-4, 1e-5, 1e-6.
• We use the Adafactor optimizer to save memory during the fine-tuning (we find Adam to not perform better). The parameters for Adafactor are: epsilon1 = 1e-30, epsilon2 = 1e-3, decay_rate =0.8.
• Each run takes approximately several hours on a TPU chip. We note that RAFT brings no additional overhead over standard fine-tuning, as it doesn’t introduce any new parameters, and it doesn’t slow down the training.

I don't understand why the authors are comparing MALI (regression on continuous values) against the Minimum Bayes Risk (MBR) prediction literature, since it optimizes non-regression metrics on the discrete grid (section 3.4). (...) It would be helpful to either strengthen the connection between MBR and MALI/RAFT, or consider moving this section to the appendix if it's not central to the paper's main argument. If you choose to keep it in the main text, consider explaining more clearly why this comparison is important for understanding RAFT's contributions.

Our aim with this comparison was to additionally strengthen the motivation for the RAFT formulation by contrasting it against an MBR-style objective. We agree that the comparison between RAFT and MBR is non-central to the paper, and following the suggestion from the reviewer, we now delegated this discussion to the Appendix D in the updated manuscript.

Similar to the experiment presented in Table 11, did the authors consider training new token embeddings for the regression task? For example, instead of using digits 1-5, they could start from 5 untrained tokens. I would expect this to work worse, but I am curious about the results.

Thank you for this suggestion. We include this ablation in Table 19 in the updated draft. We confirm the reviewer's intuition that it underperforms compared to when using digits 1-5.

The predictive head approach closely follows the RAFT; however, there is no extended discussion of what sets RAFT apart. (e.g. Table 3 and 429-430).

Please see a common response to all reviewers “RAFT improvements over predictive head”.

The authors highlight the issues of the decoder-based LLMs. However, there is no baseline comparison (i.e., it is not clear what the potential upper bound is with current technologies). In the literature, performance for datasets like SSTB is often reported as Pierson/Spearman correlation

We answer this question in the joint response to all reviewers: “Additional baselines”.

In Table 3, you report wireless predictive head results for the 2B model as 0.51, while in Table 4, the best result is 0.48 (special logits token). –

We would like to clarify that 0.48 in Table 4 corresponds to the “+ 2-layer MLP” variant. The predictive head variant instead scores 0.51 in both Table 4 and Table 3.

While comparing the effect of the grid, why have you excluded the STSB benchmark? It's a most interesting case since the grid 1-5 is a subset of the output space. I would be curious to see results and see if adding 0 makes any difference, for example

Per reviewer request, we include additional ablations of the grid for STSB in Table 19. For example, we find that only using token '5' in RAFT performs as well as using tokens '1'-'5'.

We would like to clarify that adding 0 doesn’t make a difference to the loss, because in the RAFT rule, we multiply the probability by the digit, and so p(0|x) would get multiplied by 0.

Since some things are still unclear to me in the experiment setup (please see question), I am hesitant to assign a score

Please do not hesitate to let us know if you have further questions we could clarify.

We thank the reviewer for the positive review and thoughtful comments.

“grid of letters a–e, yet Table 11 shows similar MSE to the numerical grid” and “Lemma 1 (Line 187)”

Thank you for catching both misspellings. We now corrected the text to reflect the Table. We also fixed the norm in the Lemma statement to refer to the l1 norm.

Dataset Examples: Including a table of dataset examples and inputs/outputs of the RAFT model would further help with clarity. It would help in verifying dataset characteristics too.

We report additional details on prompts and target ranges in Table 6.

Post-Training Distribution of the MALI Predictor: After RAFT training, what does the MALI predictor’s distribution look like over grid elements? Since Lemma 2 emphasizes the grid’s max and min values, it would be informative to know how the grid’s distribution evolves post-training, including entropy relative to a uniform distribution over tokens.

This is a great suggestion. To investigate this, we inspect the distribution of softmax probabilities over the grid of tokens from RAFT for the Gemma-2 2B model on STSB over test examples.

We find that, while the error decreases with training, the entropy increases after RAFT training, as we show in Figure 2 in the updated draft. This corresponds to the model on average spreading the probabilities over tokens more than prior to fine-tuning, as shown for specific examples in Figure 3 in the updated draft. We posit this to be beneficial, and indeed, RAFT fine-tuning does not restrict uncertainty of the model, contrary to the MBR fine-tuning (compare Lemma 3 and Lemma 5).

Baseline Comparisons: While the paper compares MALI-based methods with other similar approaches, are these baselines considered state-of-the-art for your chosen tasks like sentiment analysis or the review estimations? Could simpler models, such as random forests, achieve lower MSE on these tasks? In other words, the experiments provide a good comparison among the methods based on finetuning a pre-trained transformer, but I wanted to clarify if these MSEs are actually state of the art compared to other methods too.

We answer this question and provide more baselines in the joint response to all reviewers: “Additional baselines”.

Token Granularity and Grid Complexity: In the RAFT experiments, was the grid composed solely of single tokens? If so, does this also apply to the months-named grid? If the pretraining hypothesis holds that a simple grid (like 1–5) leverages pretrained information effectively, would a more complex grid (e.g., values such as 0.97, 2.3, 3.001, etc.) yield similar performance? I think this would be a very important question that can substantially strengthen the paper.

MSE Variations on Different Grids: The MSE increase when using a months-named grid (versus alphabetic grids) is curious. Are the months tokenized as single or multiple tokens? Could this affect pretraining-based semantics, or could it be due to the number of tokens in use? Understanding why alphabet tokens retain performance while month names don’t could provide insight.

Indeed, all grids we analyzed (including the months) in our work are single tokens in the vocabulary.

We do not have a clear answer as to why months perform worse than character tokens. Our main hypothesis is that the pretraining alignment has the largest impact, which can explain why digits perform the best overall, but does not explain why months perform the worst (the length of the targets does not explain this observation).

We agree more complex grid values are a great direction for exploration to understand the pretraining hypothesis better. We would like to clarify that more complex grid values that the reviewer gave as examples (0.97, 2.3, 3.001, etc.) are not in the pre-training vocabulary. Evaluating such more complex grids is certainly an interesting suggestion we will prioritize in explorations. However, it requires engineering effort to implement properly; e.g. when running multiple autoregressive steps at training, we need to carefully reuse key-value caches to avoid repeated passes through all transformer layers when computing activations for multi-token grid elements. A naive implementation significantly slows down the training.