Better autoregressive regression with LLMs via regression-aware fine-tuning
Authors: Michal Lukasik
Zhao Meng
Harikrishna Narasimhan
Yin-Wen Chang
Aditya Krishna Menon
Felix Yu
Sanjiv Kumar
Venue: ICLR 2025 Spotlight (2025)
Keywords: regression
LLMs
Primary area: foundation or frontier models, including LLMs

Abstract: Decoder-based large language models (LLMs) have proven highly versatile, with remarkable successes even on problems ostensibly removed from traditional language generation.  One such example is solving regression problems, where the targets are real numbers rather than textual tokens.  A common approach to use LLMs on such problems is to perform fine-tuning based on the cross-entropy loss, and use autoregressive sampling at inference time. Another approach relies on fine-tuning a separate predictive head with a suitable loss such as squared error. While each approach has had success, there has been limited study on principled ways of using decoder LLMs for regression. In this work, we compare different prior works under a unified view, and introduce regression-aware fine-tuning(RAFT), a novel approach based on the Bayes-optimal decision rule. We demonstrate how RAFT improves over established baselines on several benchmarks and model families.

Published as a conference paper at ICLR 2025
BETTER AUTOREGRESSIVE REGRESSION WITH LLMs VIA REGRESSION- AWARE
Michal Lukasik, Zhao Meng, Harikrishna Narasimhan, Yin-Wen Chang, Aditya Krishna Menon, Felix X. Yu, Sanjiy Kumar (mlukasik, mengzhao, hnarasimhan, yinwen, adi tyakmenon, felixy,,snniikkkeggoeeeooo Google Research
ABSTRACT
secoder-basedd large language models (LLMs) have proven highly versatile, with remarkable successes even on problems ostensibly removed from traditional lan- guage generation. One such example is solving regression problems, where the targets are real numbers rather than textual tokens. common approach to use LLMs such problems is to perform fine-tuning based on the cross-entropy loss, and use autoregressive sampling at inference time. Another approach relies on fine- tuning a separate predictive head with suitable loss such squared error. While each approach has had success, there has been limited study on principled ways of using decoder LLMs for regression. In this work, compare different prior works under a unified view, and introduce regression-aware fine-tuning (RAFT), a novel approach based on the Bayes-optimal decision rule. We demonstrate how RAFT improves over established baselines on several benchmarks and model families.
INTRODUCTION
Decoder-based large language models (LLMs) (Brown et al., 2020; OpenAl et al., 2023; Anil et al., 2023; Touvron et al., 2023; Gemini Team et al., 2024; Grattafiori et al., 2024; DeepSeck-Al et al., 2024) have set new benchmarks in challenging generative tasks (e.g., summarization, translation, open-ended dialogue). Such models' versatility has further prompted their exploration for classic predictive tasks (e.g., classification, regression, ranking) (Liu & Low, 2023; Fernandes et al., 2023; Qin al., 2023; Vacareanu et al., 2024b; Yang et al., 2023; Dukić Snajder, 2024; Lukasik et al., 2024; Vacareanu et al., 2024a), once the purview of encoder-only models such as BERT (Devlin et al., 2019). Such exploration is expected to increase given the sustained efforts towards building ever-larger decoder-based LLMs, with limited parallels in scaling encoder-based models.
Our interest is the predictive task natural language regression, where the goal to predict real-valued target given textual input. This covers important practical applications such as semantic similarity prediction (Cer et al., 2017), automatic quality assessment of translation (Koemi & Federmann, 2023) written text (Chiang & Lee, 2023), and sentiment analysis (Zhang al., 2024). At first glance, it is not apparent how to perform numeric prediction via model operating on textual tokens. Existing works have successfully followed two broad approaches. Autoregressive regression approaches directly predict as text the numeric targets (c.g., predict 12.34 by iteratively predicting tokens: '1' '2', '3', '4') or corresponding discretized categories (e.g., predict one of "bad". "ok", "good" 1) (Fernandes et al., 2023), via either standard autoregressive decoding (Liu & Low, 2023; Yang et al., 2023) or suitable modifications (Gruver al., 2023; Lukasik al., 2024; Requeima et al., 2024). Such approaches have been explored for both in-context learning, well as standard cross-entropy loss fine-tuning (c.f. Figure (a)). Predictive head approaches side-step autoregressive decoding entirely, and instead learn separate head on representations derived from the inputs. Common representations include pooold output embeddings (Zhuang et al., 2023), and the final-positee logit for special token (e.g., <extra_id 0> in T5) (Fernandes et al., 2023).
Both autoregressive and predictive head approaches have proven successful for natural language regression tasks. However, there has been (to our knowledge) no systematic comparison between these methods; further, each of them has a conceptual shortcoming. The autoregressive regression approach does not exploit the numerical nature of the regression targets, and thus does not consider

---

Published as a conference paper at ICLR 2025
Epetem (M-Epenea ((str(s) ((st(()))))))
log p(g" [a)
score
score
is
(a) Standard fine- mone-tunng ($2.4)
(b) RAFT ($4.1)
Figure 1: Illustrative example contrasting approaches to autoregressive regression fine-tuning for decoder-only Transformer. We consider an input string x ending with tokens { 'score', 'is' }. and numerical target y", which for simplicity we assume to comprise a single string token. With abuse of notation, we use y" as both its numerical value and the corresponding string token. (a) Standard (CE) fine-tuning optimizes the likelihood log ply" x) of individual tokens in the numerical string representation, essentially treating numeric generation as classification task. In general, y" can have a multi-digit representation (i.e., y/ ... Vū). which case the CE loss applied over each token. (b) RAFT leverages the LLM's predictive distribution over the next numerical token (from some grid (gria) to compute an expected numerical value, and minimizes the squared difference the target.
the fact that for a target of 1, predicting 11 is worse than predicting 1.1. On the other hand, the predictive head approach deviates from the pre-training objective typically used in secoder-based LLMs, viz. next-token prediction (Radford et al., 2018), and thus may not use the model in an optimal manner. This prompts numerical nature of targets for natural language regression tasks?
ask: how
respect both the LLM pre-training objective and the
In this work, introduce regression-aware fine-tuning (RAFT), a novel approach to autoregressive regression which makes use of the numerical nature of the targets. We prove theoretical limitations of established alternative approaches to autoregressive regression, and prove that RAFT mitigates them. We systematically compare RAFT against autoregressive and predictive head baselines, and consider several ablations for understanding the crucial design decisions for making decoder-bant based LLM work under different settings. See Table for an overview of both the previous works and the approach introduced in this work. Overall, our contributions are as follows:
We identify theoretical limitations of standard approaches to LLM-based regression ($3).
We propose regression-awart fine-tuning (RAFT), novel approach autoregressive regression, and prove that it mitigates the theoretical limitations of prior works ($4.1). We further present unified view of decoder-based LLM regression approaches, capturing both the autoregressive and the prediction head approaches ($4.2).
(iii
We systematically compare autoregressive regression, predictive head and RAFT approaches across multiple datasets and LLMs. and consistently find RAFT to be the most performant. We further attempt to pinpoint the sources of differences in the performance between different approaches, explicating the design choices behind the effectiveness of RAFT ($5).
BACKGROUND
We first introduce notation and review previous works on applying decoder-basee based LLMs to regression.
2.1 NOTATION AND PROBLEM SETTING
For finite vocabulary V of tokens (e.g., words in English), let X C V** be set of inputs comprising strings of tokens, and y C R be a set of real-valued targets. We restrict our consideration to targets with finite 10) bit representations, thus excluding irrational numbers. We further assume that each E y has a unique string representation str(y) E V*; for example, the integer 1 has the string encoding "1" ". Let P denote ground-truth distribution over X x y, with the decomposition
2

---

Published as a conference paper at ICLR 2025
Approach
Autoregressive
Fine-tuning
Inference
References
Zero-sho decoding
Standard decoding
Koemi Federmann (2023)
Fine tuning and decoding
Cross-etmiyy
Standard decoding
Fernandes et all (2023)
Zero-shot RAII.
V
None
Regression aware decoding
Lukasik et al. (2024)
Fine-tuning and RAIL.
V
Cross enntsyy
Regression
aawere aware decoding
thix work
Predictive head
Target metric
Point estimate
Fernandes et al. (2023)
RAFT
V
Target metric
Regression-awe decoding
this work
Table I: Summary of the approaches to applying decoder-basse LLMs to natural language regression tasks. Previous works rely either on using the model autoregressively (i.e., analogously how it was pre- pre-trained) or as an encoder (i.e., an output is constructed based on embeddings or logits obtained for the inputs). Different training and inference approaches have been considered for both approaches.
The natural language regression problem involves learning a predictor ÿ: X - IR that minimises the mean squared error over (input, target) pairs drawn from P:
F(x. 3) = -(z)) (py 2).
L(ÿ) >~ [yy* )
The mean squared error is a canonical choice in regression problems (Fernandes et al., 2023). The Bayes-optines predictor minimizing the above is ((z) -r(lll/)))) Classical natural language regression problems include sentiment analysis and semantic similarity prediction. More broadly, many problems can be framed as natural language regression tasks, e.g.. by crafting suitable prompt that summarizes a raw feature vector into a natural language description (Requeima et al., 2025).
We seek to employ large language models (LLMs) for such regression tasks. An LLM specifies a distribution p over strings V Given an input x E X, let p(. x) denote the corresponding conditional distribution over possible continuations. Note that it may be possible for p(= x) 0 where Z E V does not have a numerical representation; we discuss this issue more in Section 2.2.
LLMs are typically pre-trained large corpora via self-supervised objectives (Radford et al., 2018 10, few-shot and can fow or in-context learning given suitably crafted prompts (Brown et al., 2020). For example, if the goal is to predict the probability that user will enjoy movie titled "Cure", may construct an input x = "Hereditary: 0.7 Ringu: 0.9 | Cure: ", and probe the LLM's estimate of plausible continuations via p(- æz). We next discuss inference (or decoding) procedures for deriving a predictor ý given a pre-trained LLM.
2.2 STANDARD LLM INFERENCE FOR REGRESSION
Standard LLM inference involves predicting numerical targets in a generative manner, by performing autoregressive decoding to draw sample from the distribution p(- Ir::
(*R(#) = float(=)
(1)
1~p(-112 p()12)
Here, z E V* is generated autoregressively on a token-by-token basis. Further, float(z) denotes an operator that converts given string (e.g., 312.34") a corresponding numeric value (e.g., 12.34); if 2 does not have a numeric representation (e.g.. "banana"), then we assume that a suitable default value is returned. Unless otherwise stated, we assume float(=) = 0.0 for z g y.
Different algorithms may be used for the autoregressive generation of 2, e.g., greedy decoding and temperature sampling (Naseh et al., 2023). Many such algorithms seek approximate the mode:
imode (x) == arg max p(y (r). yey
Note that one may also forcibly restrict the decoding output to comprise numerical targets, e.g., by employing a form of constrained decoding (Geng et al., 2023). However, in practice, the targets from high-quality LLMs tend to be numerical even under shot settings (Lukasik et al., 2024).
2.3 RAIL: REGRESSION-AWARE LLM INFERENCE
Recently, Lukasik et al. (2024) pointed out a limitation of decoding the most likely target when employing autoregressive models for regression. Decoding of the most likely targets can be shown
3

---

Published as a conference paper at ICLR 2025
to minimize the 0-1 loss E(V. 1y ū). which may not be well aligned with the square loss of interest in regression. As a remedy, instead autoregressive decoding per Equation 1, Lukasik al. (2024) proposed the regression-awes awwre inference (RAIL) method. which given a loss & and model prediction p(: x)) estimates the Bayes-optimal minimizer of the expected loss: i.e.,
ÚRAIL (x) = arg min By-pp-lw)
[ [((float(y), w]],
(3)
VER
where float(-) is as per the previous section. For the squared loss ( 0) (y - the optimal decision rule can be shown take the following closed-form solution:
ÚRAI (x) Ey--p( s) [float(y)]-
Since p( 2)) a distribution over all possible strings, it is typically intractable to compute the above expectation exactly; this remains true even if we restrict attention to those strings corresponding to valid numerical value (of which there are infinitely many). In practice, Equation can be estimated either via sampling finite number of values, or via scoring of targets. In the latter, suppose we have some restricted target grid Agrid C y. Then, the RAIL predictor is averaged over ygrid, yielding:
ÚRAIL (a;; Vgnd)
Σ p(str(y) x) yy WEMaria
(5)
Note that E p(str(y) x)) # is possible, so the above is technically not an expectation; however, in practice, high-quality LLMs tend to concentrate most mass on numerical targets.
There are several choices of Ygrid available to the practitioner. For discrete targets y of moderate size, one may simply choose Ygria = For bounded y, one choice is equally spaced targets covering the range of y, e.g. integers or fixed-precision numbers (e.g. 2 decimal points) (Lukasik et al., 2024).
2.4 STANDARD FINE-TUNING FOR REGRESSION
The above approaches operate on a pre-trained L.M via few-shot prompting. However, it has been consistently observed that direct fine-tuning of LLMs the task of interest can be beneficial (Liu et al., 2022). Fine-tuning seeks to adapt an LLM to the target distribution I' by minimizing
L(p) = E(p [K(y" p(. x))]
(6)
for suitable loss function f: Dy. + R. where As denotes the of distributions over a S. Given sample S E (XC x y) of N (input, target) pairs drawn from P, the empirical loss is
M(p)
Σ (K.N")ES
Elv", .pp x)).
A standard choice of f is the log-loss (also referred to cross-entropy):
"(v", p(- x)) = log (ogr(ssr(y") [z),
(8)
recalling that str(y") denotes the string representation a numeric target y² E lR. More generally, one may use categorical descriptions of the target after discretising some finite grid ygrid y;
e.g., (vvery bad", "bad", "ok", "good", "very good" } (Fernandes et al., 2023).
model obtained from standard fine-tuning can rely on either of the above decoding procedures (standard decoding $2.2 or RAIL $2.3) at inference time.
2.5 PREDICTIVE HEAD APPROACHES TO REGRESSION
Predictive head approaches are an alternative to standard fine-tuning. Here, one constructs a predictor Ú(x) by utilizing activations embeddings from forward pass of the LLM. Abstractly, tract an input representation & (x)) E R", which then fed into a regressor s: IR > IR. Canonically, the regressor is linear model, yielding (xx) b w ((x) for learnable
first
R..b R.
Various choices for ((x) have been considered in previous works. To describe these, we need some additional notation. Let V = [V]. Given string x. V of length T., Transfermer-bused language
4

---

Published as a conference paper at ICLR 2025
model (Vaswani al., 2017) first constructs input embedding (2) Con E Dx where Ein E RPVV is matrix of D-dimensional token embeddings, and Col (z) E RVxL, is the one-hot embedding of each token in Next, this input embedding is passed through stack of attention and MIP layers, to produce the output emhedding fout (x) E IR DxL, One further projects this to the vocabulary space to produce output logits fout(a) (1) = F cout (a) E RVXL where Fout RDV Finally, one transforms these to distribution Pour x)) = softmax(fout(z) E [0. 1 over possible tokens via the softmax operator. For certain models (e.g., Gemma), one ties En Eour.
Given the above, one may extract an input representation through multiple means, most commonly pooling or selection of the output token embeddings, outpur logits, or output probabilities. For example, we may consider the final-token logit activation for special token v, E V (Fernandes et al., 2023; Zhuang et al., 2023), or mean-pooling the output token embeddings (Zhuang et al., 2023).
Given suitable predictor, one may directly optimize the mean squared error during fine-tuning, Note here that no autoregressive decoding conducted inference.
LIMITATIONS OF STANDARD FINE-TUNING FOR REGRESSION
We analyze limitations of standard fine-tuning for autoregressive regression (proofs in Appendix A).
3.1 LIMITATIONS OF STANDARD FINE FINE-TUNING AND STANDARD DECODING
A natural baseline is to employ cross-entropy based fine-tuning by minimizing Equation 7, and to then apply standard decoding (see Equation 2). Since the log-loss is strictly proper, minimizing Equation 7 recovers the Bayes distribution P(. x²) in the population limit (Gneiting & Raftery, 2007). In practice, however, the fine-tuned model distribution p(- 2) may deviate from IP( 2). The following Lemma shows that even when the fine-tuned model distribution p(- x) perfectly fits IP, the standard decoding predictor can incur a high squared error compared to the Bayes-optimal predictor.
Lemma 1. Assume [8] 2 2 and 0 E yd with N = max(y). For any c CE |0,, 0.5). there exist IP, p such that: Eg (P(-I=) I c, and Eg (E. (Eg (=) Imode (z)) z (2)²
Thus, using cross-entropy fine-tuning with standard decoding is not well-aligned with the eventual goal of approximating Ey-~F((la)) (x) [y"). Simply stated, although the distance between a predicted probability distribution p. and the true distribution IP may be small (or even 0), the squared error between the mean of the distribution and the mode of the distribution p can disproportionately large. This is because the mode of a distribution can be far from its mean, leading to high error.
3.2 LIMITATIONS OF STANDARD FINE-TUNING AND RAIL DECODING
Given that the standard decoding can lead to arbitrary large errors due to the predictor being ill suited for regression tasks, one may expect a better approach would be to employ cross-entropy based fine-tuning by minimizing Equation 7, and to then apply the RAIL decoding (see Equation 5). This approach indeed mitigates the issue of standard decoding being misaligned with the Bayes-optimal predictor. However, P by even a small error, the predictor can lead to significant squared error.
show that when the fine-tuned model distribution p(: [x) deviates from
Lemma 2. Assume |8] 2 2 and O E y. with N = max()). For any c E 0.5), there exists P.p such that: E. [DKL (P(-I=),p(-x>) < e, and Eg [[,))=(())) (x)) 2 (2)²
Thus again, using cross-entropy fine-tuning with RAIL might not align with the goal of approximating Ey--<(t1>>) (x)--t(t->>l"" [g"]. Intuitively, cross-entropy tine-tuning treats all "wrong" predictions the same, as is unaware of the difference in the magnitude of the numerical values represented by the tokens; e.g., if 100 and 1, 000 are two incorrect predictions each represented by single tokens, then placing large mass on the token representing 100 is penalized similarly to placing a large mass the token 000.
One solution to the above issue is to directly employ the RAIL predictor in the fine-tuning process. This requires going beyond the log-loss in Equation 8, detail next.
5

---

Published as a conference paper at ICLR 2025
4 REGRESSION-AWARE AUTOREGRESSIVE LLM TRAINING
To overcome the drawbacks of using RAIL with traditional fine-tuning, we propose a novel regression- aware objective that directly minimizes the squared loss on the RAIL predictor.
4.1 REGRESSION-AWARE LLM FINE-TUNING
As an alternative to the standard log-loss, we propose the following. Definition 1. Define the regression-awarcs fine-unning (RAFT) loss as follows:
GRAFT(Y" "eerryy,,tu ,p(- x)) = (y" (e
(9)
Equally, this uses the RAIL predictor (z) to construct numeric value from the LLM, and measures the square loss against the target y". Given finite grid Agrid C y and fine-tuning set S, the empirical loss corresponding to Equation 9 is:
LRAFT(P: Ygrid)
Σ (x,9))E
y
Σ p(str(v) [x) DERaria
(10)
N
Note that computing this loss only requires scoring each E ygrid under the model; we do not need perform any explicit sampling or decoding during training.
Compared to standard fine-tuning with RAIL decoding, we attempt to avoid the issue in Lemma 2 by seeking to directly minimize E. (Fvy=~((11)))"" - ŪRAN. (x)))² Surprisingly, despite computing
@RAIL (1) over the restricted target space Ygrid, under mild conditions the minimizer of Equation exactly mimics the Bayes-optimee predictor over the full space y.
Lemma 3. Suppose y C IR, and arg min E Vgrid and arg max E &ygria- Let (- ()) be the minimizer the RAFT loss from Definition over all distributions p(: x). Then the RAIL predictor ÚRAIL (2; Mgria) EveBaca (s) x))
constructed from p" (: x)) satisfies:
ÚRAIL (1; Vgria) [=)-=((()))
The intuition behind this result is that any numerical target in y can be expressed by convex combination of the smallest and largest numbers y, and can thus be realized by the RAIL predictor.
In Appendix C we contrast RAFT loss against a variant optimizing regression metrics via sampled model predictions, following the Minimum Bayes Risk prediction literature (Kaiser et al., 2000; Shannon, 2017; Prabhavalkar et al., 2018).
4.2 CONTRASTING AND UNIFYING RAFT AND PREDICTIVE HEAD APPROACHES
Our discussion of RAFT highlighted its close relation to autoregressive RAIL decoding, which appears rather different to predictive head approaches. However, the case of single-digit grid Ygrid (wherein each element corresponds to single token in V), the predictor function for RAFT bears similarities to the predictive head approaches. Note that y E Ygrid corresponds to a single token, by definition p(str(y) .:) Pout (- s)str(u).L- Then, the RAII, predictor becomes
ÚRAIL (2; Agria)
Σ y p(str(y) (z) yEM
Σ geMaria
y- V- x)str((),
We may now contrast RAFT against various predictors considered in prior work (see Table 2 for summary), particular, is instructive to compare this with the final-token logit activation method from Table 2. Both take the following form for an activation Ws and weight vector w E R
(xx) b+ TO y (fout (z)..).
Contrasting the RAFT and the final-token logit method, observe the following differences:

---

Published as a conference paper at ICLR 2025
Category
Approach
Predictor function j/xx)
Fine-tuning loss
Standard decoding shot (Koemi & Federmann, 2023)
arg max P(Y) a)
N/A
Anterogressive
(prior works)
zero-shot ukkaiik et al., 2024)
Evey V pp(y(ar)
N/A
Standard fine-tuning and standard decoding (Fernandes et al., 2023)
ang max P(yla)
logpty" xt)
Autoregressive
RAIL standard fine-tuning
Evey V/ p(y/(r)
logpív" 2t)
(this work)
RAFT (this work)
General RAFT
Egeyy p(yy(u)
(@(x) - w²"
Single-digit RAFT
Eiee vf (Pnn
(0(x) - =*)
Predictive head (prior works)
Final-token logit (Femandes et al,, 2023)
6- fau(())or, b+ w pool(eous(ar)
(9(x) =/
Pooled oulpul embeddings (Zhuang et al., 2023)
(Ø(r)) =²)"
MLP on the final-token logits
6+ MLP(E Faan
(8(x) w-
Predictive head (this work)
Probability vector projection
6 10 "on
(B(r)) -0"
Leamable regression-awwe aware training
Evey Wey Pou Poml (2))12
(0(x) 3/'
Table 2: Approaches for applying decoder-based LLMS regression. Here, p(- | 2) denotes distribution over possible outputs given an input string x, and ((x) E R predictor given by predictive head approach. and w are learnable parameters, v* operator, L is the length of input I, and ddiiits denotes all digits covering the range of targets (unless otherwise stated, '1 '5' The first 4 rows show the autoregressive baselines: standard decoding (Section 2.2), RAIL zero-shot (Section 2.3), standard fine-tuning and decoding (Section 2.4), RAIL with standard fine-tuning (Section 3.2). The next 2 rows show RAFT: the general autoregressive form (y ygid for general output spaces), and the single digit version (e.g. {1, 2, 4, 5)). The following 2 rows present the prior works from Fernandes et al. (2023); Zhuang et al. (2023). The last 3 rows present new predictive head approaches that attempt to mimic the behavior of RAFT.
a fixed token, pool is a pooling
Activation: for single-digit RAFT, the softmax activation that converts (x)) to the proba- bility vector Pout (- 2). For the final-token logit, V is the identity activation.
Weight vector: for the final-token logit, is one-hot vector with corresponding the special token position. For single-digit RAFT, w. = float(u) for each E V; note that, as a result, positions corresponding to non-digits have weight
Initialization: another important factor is that the RAFT predictor initialization exactly co- incides with RAIL decoding, and thus, forms a strong predictor for zero-shot inference with LLMs (Lukasik et al., 2024). By contrast, most predictive head approaches will incur a high error initialization due to deviating from the next token prediction task. Therefore, RAFT can be seen as a predictive head approach with strong performance at initialization, potentially making optimization easier; indeed, we empirically observe RAFT to converge faster compared to the baselines (see Figure 4 and Figure 5 in Appendix E.7).
In light of the close similarities between RAFT and the final-token logit approach, it is prudent to carefully analyze these differences and identify whether any of these choices play an important role in RAFT's performance. Therefore, introduce the following
predictive head variants:
MLP on final-token logits: this is a variant of the final-token logit method, wherein a 2-layer MLP with non-linear activation (sigmoid) is employed on the entire final-token logit vector, rather than selecting the logit for single special token:
(x) b MLP (.)....)
(11)
Learnable-RAFT this a variant of RAFT, wherein the weights over the output model probabili- ties are learned, rather than being fixed the vector w. = float(v):
((x) = Σ Wy Pout Pour (x)str(())))
(12)
N'EY
The learnable-RAFT variant adds more flexibility to the predictor function ý over the vanilla RAFT method. However, as with other predictive head methods, it deviates from the next-token prediction pre-training task. Which of these two factors predictor flexibility, and alignment to the pre-training task - the most important? To answer this question, we compare Carmable-RAFF against RAFT, and also experiment with fine-tuning from randomly initialized (as opposed a pre-trained) model.
7

---

Published as a conference paper at ICLR 2025
Datuset
Model size
Zero-shot standard decoding
Zernsshot
Stundard line-tniing decoding
Stundard fine line-tnning RAII.
Predictive head
RAFT
Wireless
28
0.88
0.70 + 0.01
0.67 = 0.01
0.51 = 0.01
0.47 0.01
(Amazin)
98
0,86 +
0.00
Personal cure
28
1.00
0.77 0.01
0.74 #: 0.01
0.52 + 0.02
0.49 0.00
(Amazon)
Music
0.11
0.00
(Amazon)
STSB
28
0.01
0.57
0.54 0.00
0.54
9B
1.37
0.58 0.00
0.53 d: 0.00
0.52 0.00
0.51 0.01
Table 3: RMSE across datasets, methods, and Gemma-2 models of varying sizes. Fine-tuning methods report mean = std dev from model retraining. See Table 10 (Appendix) for Gemma-2 27B.
EXPERIMENTAL RESULTS
We now present experiments and ablations comparing the autoregressive regression, predictive head and RAFT approaches on natural language regression datasets. We make the following main empirical findings: RAFT outperforms all autoregressive regression and predictive head baselines across datasets and models; (ii) ablations indicate the importance of aligning fine-tuning to the pre-training loss; (iii) RAFT tends work well when the grid Ygrid corresponds to digit tokens.
5.1 EXPERIMENT SETTINGS
Datasets. We use the following natural language regression datasets, with RMSE as the main metric:
(i) US Amazon reviews, where aim to predict the 5-star rating for product review (Ni et al., 2019). We consider a few categories from the Amazon reviews datasets, each forming a separate dataset: Wireless, Music, Personal products. We use 1,500 examples for the test set (after Lukasik et al. (2024)), 1,500 for validation and 10,000 examples for training.
(ii) Semantic Textual Similarity Benchmark (STSB) (Cer et al., 2017). comprising of sentence pairs haman-annotated with a similarity score from 5. To measure the impact of varying the dataset size, we also consider 1,000 examples for training (STSB IK; see Table 8 in Appendix).
(ii) MovieLens-IM, where construct movie rating prediction task following Luo al. (2024).
We summarize the dataset statistics and the prompts in Table 6 and Table 7 (Appendix).
Models. We experiment with Gemma-2 (Team et al., 2024) and PaI.M2 (Anil et al., 2023) instruction- tuned model families of different sizes (see Appendix D for training details). Where standard deviations are reported, fine-tuning performed times.
Methods. We compare the following methods: (1) autoregressive baselines (Section 2.2), RAIL zero- shot (Section 2.3), RAIL with cross-entropy fine-tuning (Section 3.2); (2) predictive head approaches from Fernandes et al. (2023); Zhuang et al. (2023); (3) the new RAFT method (Section 4.1); (4) new predictive head approaches that attempt to mimic the behavior of RAFT (Section 4.2). In zero-shot standard decoding. we use greedy decoding: in zero-shot RAII,, the predictor is obtained by sampling 32 targets with temperature T (Lukasik et al., 2024). We also run ablations with replacing causal attention masking with bi-directional attention masking, following previous works on classification with decoder-based LLMs (Dukić & Snajder, 2024; Qorib et al., 2024).
Implementation of the RAFT objective. An important practical consideration is the grid ygrid- Unless otherwise stated, we choose Ygrid = 1,2, 3,4,5 as targets from all considered datasets belong to [0, 5). For Amazon reviews datasets, Igrid ly, while for STSB, Yguid C y (as the targets take floating point values). Recall that RAFT can represent real valued targets even under such a choice for y gri (see Lemma 3). Nonetheless, we analyze the impact of Ygrid the results.
8

---

Published as a conference paper at ICLR 2025
1.0
2.536
Stan derd fine-tuniing Preciction head
Standard fine-turing Presictive head RAFT
1.706
0.227
AMSE
0.015
0.005
Fro trained
Pre-Irwineed
Random
(a) MovieLens- -M.
(b) Synthetic data ("Original #1"),
Figure 2: Comparing standard fine-tuning, predictive head and RAFT approaches on Moviel (left) and the synthetic "Original #1" (right) datasets with Gemma-2 2B. RAFT outperforms the baselines on MovieLens. RAFT also improves over the baselines on the synthetic dataset "Original when using pre-trained but a randomly initialized checkpoint, (Best viewed in color.)
Approach
Gemma-2 2B
Gemma-2 9B
Standard fine-tuming
0.83
0.82
Special-token logit (Fernandes et al., 2023)
0.51
0.46
Special-token logit 2-layer MI.P
0).88
0.46
Pooled output embeddings (mean) (Zhuang al., 2023)
0.50
0.47
Pooled outpul embeddings (min)
1.38
1.18
Pooled output embeddings (max)
1.32
1.13
RAFT
0.47
0.45
Learnable-RAee over all tokens
0.48
0.45
Learnale-RAFTT over digits '1'-15'
0.47
0.45
Table 4: RMSE Gemma-2 models Amazon Wireless across predictive head and RAFT variants. RAFT outperforms the different predictive head variants. Additionally, we find that RAFT with learnable decision rule (learnableRRAFT) does not improve over RAFT with the fixed decision rule.
5.2 RAFT LEADS TO BETTER AUTOREGRESSIV REGRESSION
We compare different autoregressive and prediction head approaches across across Gemma-2 models of varying sizes in Table 3. We report additional results from PaLM-2 models on STSB in Table 11 (Appendix) to verify the findings across an additional model family. We make several observations.
First, we verify the value of both (1) use of an appropriate decoding strategy (greedy versus regression- aware inference). and (2) fine-tuning over zero-shot inference. Indeed, we find that zero-shot greedy decoding, RAIL (see Section 2.3), standard fine-tuning with greedy decoding (see Section 3.2) and standard fine-tuning with RAIL inference (see Section 4.1) work increasingly better.
Second, we find that the predictive head approach outperforms the autoregressiie baselines, including those that perform standard fine-tuning. This corroborates Lemma and Lemma 2, which pointed at the limitations of standard fine-tuning due to it being misaligned with the squared error.
Finally, we find RAFT to outperform the predictive head and the autoregressive approaches across almost all settings. RAFT outperforming the autoregressive approaches corroborates the posited importance aligning the fine-tuning loss in regression tasks to a regression loss. RAFT outper- forming the predictive head approach corroborates the posited importance of not deviating from the autoregressive setting, which aligns with the next-token prediction pre- pre-training task.
To further evaluate RAFT, next contrast the key approaches (RAFT, predictive head and standard fine-tuning) on a large scale dataset for the movie recommendation problem (MovieLens- 1M) with Gemma-2 2B model. We again find RAFT improve over the baselines (see Figure 2(a)).

---

Published as a conference paper at ICLR 2025
1.85
1.89
Standard fine-tuning Predictive head
Standard fine tun-tunnng Predictive head
1.5
RMSE
1.0
RASE
1.0
0.94
0.92
0.5
0.47
0.49
0.5
0.44
siverges
diverggu
0.0
Pre-trained checkpaint
Random
Pre-trained biddirctional
0.0
Pre-trained checspoint
Random checkpoirt
Pre-trameed byddrectiona
checkpaint
(a) Gemma-2
(b) Gemma-2 9B.
Figure 3: Comparing standard fine-tuning, predictive head and RAFT variants on Amazon Wireless. We find that changing causal to bi-directional attention masking does not significantly affect the results of the predictive head and RAFT, while expectedly making standard fine-tuning untrainable (due to the model being able to attend to the predicted target). We also analyze the effect random initialization, and find all methods to perform worse. RAFT improves initializing from a pretrained checkpoint, but not when initializing from a random checkpoint, empirically corroborating the argument for better alignment of RAFT to the pre-training task.
predictive head when
5.3 ABLATING DESIGN CHOICES IN RAFT AND PREDICTIVE HEAD APPROACHES
Predictive head non-linear variant. We experiment with variants of learnable-RAFT (see Equation 12): one where the weight vector is learnt for all vocabulary entries, and the second, where only entries corresponding digits '1'-'5' are learnt, while other entries are fixed to O. For both variants we found it necessary to initialize from the solution corresponding to RAFT, as random initialization did not lead good training dynamics. Overall, as reported Table 4, we find both approaches to not improve over RAFT. We also consider adding a non- lonelinnar function over the special-token logit in the form two-layer MLP (see Equation 11). This again does not lead to improvements over RAFT. Both results demonstrate that it may be more important to align the fine-tuning to the pre-training loss, as opposed to only try make the predictor more expressive.
Predictive head design choices. We experiment with additional predictive head variants, including pooling over the full sequence of token embeddings from the LLM, instead of taking the final token activation. As shown in Table 4, find it to not lead significantly better results than the special- token logit method, Additionally, as seen on Figure 3, we find bi-directional masking of attention does significantly affect the results of the predictive head and RAFT, while expectedly making standard fine-tuning untrainable (due to the model being able to attend to the predicted target during training). Overall, none of the predictive head variants improves over the RAFT approaches, again suggesting the importance of not deviating from the pre-training loss in fine-tuning.
shed light on the importance of aligning the method with the pre-training task, we experiment with fine-tuning from randomly initialized checkpoint (see Figure 3). We find that standard fine-tuning does not converge to a reasonable result, RAFT converges to a poor predictor, and predictive head fares the best. We next experiments dataset from Vacareanu et al. (2024a) (the Original #I dataset) and report results in Figure 2(b). We corroborate the finding that RAFT improves checkpoint, and not when the model weights are initialized randomly. This supports the hypothesis that RAFT outperforms predictive head due alignment with the pre-training task.
Role of pre-training.
synthetic regression
predictive head when initialized from a pre-trained
DISCUSSION AND FUTURE WORK
We introduced regression-aware fine-tuning (RAFT), new method for fine-tuning decoder-based LLMs to predict numeric targets. We demonstrated empirically that RAFT can consistently out- perform existing methods that perform standard cross-entropy fine-tuning, well methods that construct separate predictive heads. An interesting direction for study would applications such techniques to problems like time-series forecasting, well problems of ordinal regression.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent EI Shafey, Yanping Huang, Kathy Meier-Hellstern, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry. Christopher A. Proquette-Choo. Aakanksha Chowdhery. Clément Crepy. Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vlad Feinberg, Fangxiaoyu Feng, Vlad Fienber, Markus Freitag, Xavier Garcia, Sehastian Gehrmann, Lucas Gonzalez, Guy Gur-Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui, Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagielski, Wenhao Jia, Kathleen Kenealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontaek Lim, Hanzhao Lin, Zhongtao Liu, Frederick Liu, Marcello Maggioni, Aroma Mahendru, Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Bric Ni, Andrew Nystrom, Alicia Parrish, Marie Pellat, Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao, Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Sacta, Rajkumar Samuel, Renee Shelby, Ambrose Slone, Daniel Smilkov. David R. So, Daniel Sohn, Simon Tokumine, Dasha Valter. Vijay Vasudevan, Kiran Vodrahalli, Xuezhi Wang, Pidong Wang, Zirui Wang. Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng Yin, Jiahui Yu, Qiao Zhang. Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu. Palm 2 technical report, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Berbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, Curran Associates, Inc., 2020.
volume 33, pp. 1877-1901.
Daniel Cer, Mona Diab, Eneko Agirre, Inigo Lopez-Gazpic and Lucia Specia. SemEval-2017 task I: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the I 11th International Workshop Semantic Evaluation "SemEval-2011 Association for Computational Linguistics, 2017. 18653/v1 /5172200111
doi: 10.18653/v1ss11720001 URL //dx:://xx..e doi. org/10.
Cheng-Han Chiang and Hung-yi Lee. Can large language models be an alternative to human evaluations? In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the s1 Annual Meeting the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15607-15631, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: URL htt.s: htttp:://cclanthologg. org/ 2023 acl-long. 870/.
DeepSeck-AI, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao. Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang. Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Hanwei Xu, Hao Yang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jin Chen, Jingyang Yuan, Junjie Qiu, Junxiao Song, Kai Dong, Kaige Gao, Kang Guan, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Liyue Zhang, Meng Li, Miaojun Wang. Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li. Ning Tian, Panpan Huang, Peiyi Wang. Peng Zhang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L.. Jin, Ruiqi Ge, Ruizhe Pan, Runxin Xu, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Size Zheng, T. Wang, Tian Pei, Tian Yuan, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang. Xiaojin Shen, Xiaokang Chen, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Liu, Xin Xie, Xingkai Yu, Xinnan Song, Xinyi Zhou, Xinyu Yang. Xuan Lu, Xuecheng Su, Y. Wu, Y. K. Li, Y. X. Wei,
Y. x. Zhu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui
11

---

Published as a conference paper at ICLR 2025
Wang, Yi Zheng, Yichao Zhang, Yiliang Xiong, Yilong Zhao, Ying He, Ying Tang, Yishi Piao, Yixin Dong, Yixuan Tan, Yiyuan Liu, Yongji Wang, Yongqiang Guo, Yuchen Zhu, Yuduan Wang, Yuheng Zou, Yukun Zha, Yunxian Ma, Yuting Yan, Yuxiang You, Yuxuan Liu, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhewen Hao, Zhihonng Shao, Zhiniu Wen, Zhipeng Xu, Zhongyu Zhang, Zhuoshu Li, Zihan Wang, Zihui Gu, Zilin Li, and Ziwei Xie. Deepseck-v2: A strong, economical, and efficient mixture-of-experts language model, 2024. URL httph:httaaaete wwwpvv.ravvvvvvv/o/vv org/abs/220 /4005 04434.
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
David Dukić and Jan Snajder. Looking right is sometimes right: Investigating the capabilities of decoder-only LLMs for sequence labeling. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics ACL 2024, pp. 14168-14181, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics. URL https: / www.ahhooogty....oommooooo org/ 2024. findings- acw.ggssacll... 843.
Patrick Fernandes, Daniel Deutsch, Mara Finkelstein, Parker Riley, André Martins, Graham Neubig, Ankush Garg, Jonathan Clark, Markus Freitag, and Orhan Firat. The devil is in the errors: Leverag- ing large language models for fine-graide treeggraied machine translation evaluation. In Philipp Koehn, Barry Haddow, Tom Koemi, and Christof Monz (eds.), Proceedings of the Eighth Conference on Machine Translation, pp. 1066-1083, Singapore, December 2023, Association for Computational Linguis- tics. doi: 10.1 8653/v 1/2023. wmt-11.100. URL https: / ////a/ccaantholooyyoooooo www.orgg//22 wmt-1. 100.
Gemini Team et al. Gemini: A family of highly capable multimodal models, 2024.
Saibo Geng, Martin Josifoski, Maxime Peyrard, and Robert West. for structured NLP tasks without finetuning. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023. URL https: http://pppereeeeeeeeeee /wt/ wawoorrmmmmikkkyyo 1i==KHYYWNDDII
decoding
Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and esti- mation. Journal of the American Statistical Association, 102(477):359-378 2007. doi: 10.1 10.119800162155000000011
URL httpt:tttts:toorrrt 10. 1991 0 /0//0//121155000001100
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong. Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song. Danielle Pintz, Danny Livshits, Danny Wyatr, David Esiobu, Dhruv Choudhary. Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lec, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, mmanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yeaa Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic. Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishman, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic,
Plawiak,
12

---

Published as a conference paper at ICLR 2025
Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vitor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitey Meers, Xavier Martinet, Xiaodong Wang. Xiaofang Wang. Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, 7oe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape. Bing Liu, Bo Wu, Boyu Ni. Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Shhing-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang. Duc Le. Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni. Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpem, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj. Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean- Saan-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang. Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang. Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A. Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manay Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari. Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar. Ozlem Kalinli. Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji. Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj. Qian Liang. Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng. Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyighe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kochler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru,
13

---

Published as a conference paper at ICLR 2025
Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang. Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang. Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen. Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang. Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary Vito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang. Zhiwei Zhao, and Zhiyu Ma. The Llama herd of models, 2024. URL https: /taps://arxivv http:///aaaiiv.ooaaa//////e abgga 21783.
Nate Gruver, Mare Finzi, Shikai Qiu, and Andrew Gordon Wilson. Large language models are zero-shot time series forecasters, 2023.
Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regu- larized optimization. arXiv preprint arXiv: 10191111.3777 2019.
Janez Kaiser, Bogomir Horvat, and Zdravko Kacic. A novel loss function for the overall risk criterion based discriminative training of hmm models. In 6th International Conference on Spoken Language Processing (ICSLP 2000), pp. vol. 2, 887-890, 2000. doi: 10.21437/CCSLPP2000-412
Tom Kocmi and Christian Federmann. Large language models are state-of-the-art evaluators of translation quality. In Mary Nurminen, Judith Brenner, Maarit Koponen, Sirkku Latomaa, Mikhail Mikhailov, Frederike Schierl, Tharindu Ranasinghe, Eva Vanmassenhove, Sergi Alvarez Vidal, Nora Aranberri, Mara Nunziatini, Carla Parra Escartín, Mikel Forcada, Maja Popovic. Carolina Scarton, and Helena Moniz (eds.), Proceedings of the 24th Annual Conference of the European Association for Machine Translation, pp. 193-203, Tampere, Finland, June 2023. European Association for Machine Translation. URL https: / http:///acaaathlloo.oooo
or g2233 eamt-1.
Zongxi Li, Xianming Li, Yuzhang Liu, Haoran Xie, Jing Li, Fu lee Wang, Qing Li, and Xiaoqin Zhong. Label supervised Ilama finetuning, 2023. URL https: //atxi://///xv....///// ab5/2300. 01208.
Haokun Liu, Derek Tam, Mohammed Mugeeth, Jay Mohta, Tenghao Huang. Mohit Bansal, and Colin A Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. In S. Koyejo, Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 1950-1965. Curran Asso- ciates, Inc., 2022. URL @cde695b83bd 86c1 (e1eee33dd166111111
neurips. ctppiisssppppppppceeeoeoee patp/ppperrrff/ fit.ss//appre//22//iiee. pdf.
Tiedong Liu and Bryan Kian Hsiang Low. Goat: Fine-tuned LLaMA outperforms GPT-4 on arithmetic tasks, 2023.
Michal Lukasik, Harikrishna Narasimhan, Aditya Krishna Menon, Felix Yu, and Sanjiv Kumar. Regression aware inference with LLMs. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 13667-13678, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10.18653) 1122044findings-emmlpp
URL https: / hatp//aaaaanntoggooyooooo orrp:/llggooooooo//ooooooo 2024. http///2222..nnsseeeeeeee findings-emnlp 799/.
Sichun Luo, Yuxuan Yao, Bowei He, Yinya Huang, Aojun Zhou, Xinyi Zhang, Yuanzhang Xiao, Mingjie Zhan, and Linqi Song. Integrating large language models into recommendation via mutual augmentation and adaptive aggregation, 2024. URL https: //apxvaaaiiv.///aa org. orggabss/240 13870.
Ali Nasch, Kalpesh Krishma, Mohit Iyyer, and Amir Houmansadr. Stealing the decoding algorithms of language models. In Proceedings of the 2023 ACM SIGSAC Conference on Computer and Communications Security, CCS '23, pp. 1835-1849, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700507.
doi: 10.1 10.1145/3576115336166
URL
https://dol org/10 /00 11000..145544 /141/1576915..... 36662..
Jianmo Ni. Jiacheng Li, and Julian McAuley. Justifying recommendations using distantly-labeled reviews and fine-graines aspects. In Kentaro Inui, Jing Jiang, Vincent Ng. and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),, pp. 188-197, 197, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10. 20.18653/v1/DD 1018. URL
14

---

Published as a conference paper at ICLR 2025
Made Nindyatama Nityasya, Haryo Wibowo, Alham Fikri Aji, Genta Winata, Radityo Eko Prasojo, Phil Blunsom, and Adhiguna Kuncoro. On "scientific debt" in NLP: A case for more rigour in language model pre- pre-training research. In Anna Rogers, Jordan Boyd-Graber. and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 8554-8572 Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/1120233.ee acl acl-long..477 URL https////caaathloogg.. org/2023 acl- 477.
OpenAl et al. GPT-4 technical report, 2023.
Rohit Prabhavalkar, Tara N. Sainath, Yonghui Wu, Patrick Nguyen, Zhifeng Chen, Chung-Cheng Chiu, and Anjuli Kannan. Minimum word enror rate training for attention- based sequence-to- sequence models. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4839 4843, 2018. doi: 10. 10.1109/CASSP22018.... 8...5..8008..4611899
Zhen Qin, Rolf Jagerman, Kai Hui, Honglei Zhuang, Junru Wu, Jiaming Shen, Tianqi Liu, Jialu Liu, Donald Metzler, Xuanhui Wang, and Michael Bendersky. Large language models are effective text rankers with pairwise ranking prompting, 2023.
Muhammad Qorib, Geonsik Moon, and Hwee Tou Ng. Are decoder-only language models better than encoder-only language models in understanding word meaning? In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics ACL. 2024, pp. 16339-16347, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics. URL https:/ / aalantholoogyy.....o org/2024, findings-adee .67.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
James Requeima, John Bronskill, Dami Choi, Richard E. Turner, and David Duvenaud. Llm processes: Numerical predictive distributions conditioned on natural language. CoRR, 0bss/2405.12856 2024. www.uniitte tniitrrer... http:teeeeeeeerd/eeetetetterette www.dd///ooooaaleee /ournalss/cor//orr222//
URL http: .ni.ppuu.iiiim
htbs##bs 2405-12855
James Requeima, John Bronskill, Dami Choi, Richard Turner, and David Duvenaud. Llm
processes: Numerical predictive distributions conditioned on natural language. Advances in Neural Information Processing Systems, 17:609619-1 2025.
Igal Sason. On reverse Pinsker inequalities. arXiv preprint
2015.
Matt Shannon. Optimizing expected word error rate via sampling for speech recognition, 2017. URL https://arxiv. wwtp:///////////o/oooo 02776.
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, Johan Ferret, Peter Liu, Pouya Tafti, Abe Friesen, Michelle Casbon, Sabela Ramos, Ravin Kumar, Charline Le Lan, Sammy Jerome, Anton Tsitsulin, Nino Vieillard, Piotr Stanczyk, Sertan Girgin, Nikola Momchev, Matt Hoffman, Shantanu Thakoor. Jean-Bastien Grill, Behnam Neyshabur, Olivier Bachem, Alanna Walton, Aliaksei Severyn, Alicia Parrish, Aliya Ahmad, Allen Hutchison, Alvin Abdagic, Amanda Carl, Amy Shen, Andy Brock, Andy Coenen, Anthony Laforge, Antonia Paterson, Ben Bastian, Bilal Piot, Bo Wu, Brandon Royal, Charlie Chen, Chintu Kumar, Chris Perry, Chris Welty, Christopher A. Choquette-Choo, Danila Sinopalnikov, David Weinberger, Dimple Vijaykumar, Dominika Rogozińska, Dustin Herbison, Elisa Bandy, Emma Wang, Eric Noland, Erica Moreira, Evan Senter, Evgenii Eltyshev, Francesco Visin, Gabriel Rasskin, Gary Wei. Glenn Cameron, Gus Martins. Hadi Hashemi, Hanna Hiimezak-Plucinskke Harleen Batra, Harsh Dhand, Ivan Nardini, Jacinda Mein, Jack Zhou, James Svensson, Jeff Stanway, Jetha Chan, Jin Peng Zhou, Joana Carrasqueira, Joana Iljazi, Jocelyn Becker, Joe Fernandez, Joost van Amersfoort, Josh Gordon, Josh Lipschultz, Josh Newlan, Ju yeong Ji, Kareem Mohamed, Kartikeya Badola, Kat Black, Katie Millican, Keelin McDonell, Kelvin Nguyen, Kiranbir Sodhia, Kish Greene, Lars Lowe Sjoesund, Lauren Usui, Laurent Sifre, Lena Heuermann, Leticia Lago, Lilly McNealus, Livio Baldini Soares, Logan Kilpatrick, Lucas Dixon, Luciano Martins, Machel Reid, Manvinder Singh, Mark Iverson, Martin Görner, Mat Velloso, Mateo Wirth, Matt Davidow, Matt Miller, Matthew Rahtz, Matthew Watson, Meg Risdal, Mehran Kazemi, Michael Moynihan, Ming Zhang, Minsuk Kahng, Minwoo Park, Mofi Rahman, Mohit Khatwani, Natalie Dao, Nenshad
15

---

Published as a conference paper at ICLR 2025
Bardoliwalla, Nesh Devanathan, Neta Dumai, Nilay Chauhan, Oscar Wahltinez, Pankil Botarda, Parker Barnes, Paul Barham, Paul Michel, Pengchong Jin, Petko Georgiev, Phil Culliton, Pradeep Kuppala, Ramona Comanescu, Ramona Merhej. Reena Jana, Reza Ardeshir Rokni, Rishabh Agarwal, Ryan Mullins, Samaneh Saadat, Sara Mc Carthy, Sarah Perrin, Sébastien M. R. Arnold, Sebastian Krause, Shengyang Dai, Shruti Garg, Shruti Sheth, Sue Ronstrom, Susan Chan, Timothy Jordan, Ting Yu, Tom Eccles, Tom Hennigan, Tomas Kocisky, Tulsee Doshi, Vihan Jain, Vikas Yadav, Vilobh Meshram, Vishal Dharmadhikari, Warren Barkley, Wei Wei, Wenming Ye, Woohyun Han, Woosuk Kwon, Xiang Xu, Zhe Shen, Zhitao Gong, Zichuan Wei, Victor Cotruta, Phoebe Kirk, Anand Rao, Minh Giang, Ludovic Peran, Tris Warkentin, Eli Collins, Joelle Barral, Zoubin Ghahramani, Raia Hadsell, D. Sculley, Jeanine Banks, Anca Dragan, Slav Petrov, Oriol Vinyals, Jeff Dean, Demis Hassabis, Koray Kavukcuoglu, Clement Farabet, Elena Buchatskaya, Sebastian Borgeaud, Noah Fiedel, Armand Joulin, Kathleen Kenealy, Robert Dadashi, and Alek Andreev. Gemma 2: Improving open language models at a practical size, 2024, URL https:/ https:://axxvvorrpp/ abs/2408. 001 18.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
Robert Vacareanu, Vlad Andrei Negru, Vasile Suciu, and Mihai Surdeanu. From words to numbers: Your large language model is secretly a capable regressor when given in-context examples. In First Conference on Language Modeling, 2024a. URL LzpaUxcNFK.
Robert Vacareanu, Vlad-Andrei Negru, Vasile Suciu, and Mihai Surdeanu. From words to numbers: Your large language model is secretly capable regressor when given in-context examples. arXiv preprint 44X1v:2404.07544, 2024b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Proceedings of NeurlPS 2017, 2017.
Zhen Yang, Ming Ding, Qingsong Lv, Zhihuan Jiang, Zehai He, Yuyi Guo, Jinfeng Bai, and Jie Tang. Gpt solve mathematical problems without calculator, 2023.
Wenxuan Zhang, Yue Deng, Bing Liu, Sinno Pan, and Lidong Bing. Sentiment analysis in the of large language models: A reality check. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Findings of the Association for Computational Linguistics: NAACL 2024, pp. 3881-3906, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 18.18v53/v1/ 2024. 1o2.findings-nacll.22 URL https: httpp//:///aaaaahhoooooooo org/ 2024. wt2...f..nnii.. findings-naacl 246.
Honglei Zhuang, Zhen Qin, Rolf Jagerman, Kai Hui, Ma, Jing Lu, Jianmo Ni, Xuanhui Wang, and Michael Bendersky. Rankt5: Fine-tuning 15 for text ranking with ranking losses. In Proceedings of the 46th International ACM SIGIR Conference on Reseurch and Development in Information New York, NY, USA, 2023. Association for Computing doi: 10. 18444//553918835224444 URL https: / //doi. org/10 100.
Retrieval, SIGIR '23, pp. Machinery, ISBN 978 9781450394086, 0445/3536188.. 35928777
16

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
A Proofs of theoretical results
18
A.1 Proof for Lemma
18
A.2 Proof for Lemma 2
18
A.3 Proof for Lemma
18
B Related works
C Comparison to MBR-based fine-tuning
21
D Additional experimental details
23
E Additional experimental results
24
STSB-1K
24
Synthetic data
24
E.3
Comparison encoder- based baselines ................................
Gemma-2 27B and PALM-2 models
24
E,4
24
E.5
Choices for tokens in RAFT
25
E.6 Choices for loss, normalization and token initialization in RAFT
26
E.7 Convergence speed RAFT
29
E,8 Distribution over tokens in the RAFT Predictor
30
17

---

Published as a conference paper at ICLR 2025
A PROOFS OF THEORETICAL RESULTS
A.1 PROOF FOR LEMMA 1
Lemma (Restated). Assume [9] 2 2 and 0 y, with N = max(U). For any [0, (0,0.5), there exist P. such that: Eg [KKL (P(()))))(((I))
and E. (B. (By)~>((1a) imode (2))
N
Proof. Pick any E [0. 1]. Recall that N = max (). Consider distribution IP where, for any -0.50 and all other targets attain probability
example a E X, IP( x) 2 P(N x) 2
1+0.5a
Next, consider the model distribution: p(0 z) = 1-0.5a p(N 2) = 1 and all other targets
2
attain probability 0. Then, for any x E X, we have: (0mode (2) - By-~p(y"]) = (*) (1+ 9)
()² and [|P(+|z>= - p(-1x))l1 = a. Using the reverse Pinsker's inequality (Sason, 2015), we further have
- DRi. (P(())))))) S log 1+ - p((1=)11
2a
log
1+ 1 0.5c
S log (1 + 4a²).
minzeyp(yla)
where we use the fact that a I 1. Setting
Live I, for any E [0,])]())], we have:
NN 2 2
Eg [Tx (P(())))))) S;
Ey
(Wmoste - -xy=~=(""1)
A.2 PROOF FOR LEMMA 2
Lemma (Restated). Assume [3] 2 and 0 E y. with N max(y). For any e (0,0. there exists P.P such that: E, [Ki. (Pp(2z)()))))) ((--\=)) e, and E., (B) (x)) )
Proof. Pick any (0, 1]. Recall that N = ma) Consider distribution IP such that, for any example a E x, (P x)
1+0.5a IPN (x) 1-0.50 and all other targets attain probability
2
2
Next, consider the model distribution: p(0 x) = 1-0.5a
p(N x) = 1+0.3a and all other 2
2
targets attain probability 0. Then, for any E X: (Ev-~~x [ - ÚRAN. (x)) = (a²) 2 and
[IP(-|z) p(11=)]l a. Using the reverse Pinsker's inequality (Sason, 2015), we further have
IDKL (P(-(())) & log
[+ [[[((2>) - p(-1=)11))
= log (1+ 1 1- 0.5c
(1+
2a
< log (1 + 4a?).
mingey
where we use the fact that < 1.
Setting a = -eve I, for any c E (og(()))),, we have E, [TK... (P(: (P(-(x),, (0(()x))(((a))) S c, and:
N (e 1) 2 4 (1+e - 1) = N
Ex (LE wDt un) u>lv"| DRAIL (x))²
4
where we use the fact that et 2 c.
A.3 PROOF FOR LEMMA 3
Proof. For simplicity, we avoid explicitly stating conversions from float to string, and vice versa. For any x, we wish to minimize:
27
Σ p(ylx) y 3²
Ey
18

---

Published as a conference paper at ICLR 2025
Equating the derivative w.n.t. p(ylæc) to O, we derive the first-order condition for optimality:
p(u'(z) (2)x)).11 of - 1/ 5r) ]]
Σ
P((l2)
y = Vy E Agria-
D'EMgria
So the optimal solution is achieved when:
E p(y'zz) v' = -o-
From the conditions in the lemma, we have that the smallest and largest numbers in y are present convex combination of the smallest and largest
in Egrid- Since (e [y"] can be expressed
numbers in y. it can also be expressed as convex combination of numbers in Agrid- Hence, there
exists probability distribution p" (3)z²) such that ÚRAIL (z; Vgrid) Ey'Eygrua P (1/ (x) 3
Ey-xP(-)e) [y"]. thus satisfying the condition for optimality.
L
19

---

Published as a conference paper at ICLR 2025
B RELATED WORKS
Encoder-based models (e.g., BERT) relying on the masked language modeling pretraining tasks have been primarily employed discriminative tasks (including classification and regression) (Devlin et al., 2019). Decoder-based large language models (LLMs) (c.g.. GPT, LLaMa) on the other hand, mostly relying on the next-token prediction pretraining task, showed state of the art results across range of generative tasks. (OpenAl et al., 2023; Anil et al., 2023; Touvron et al., 2023; Gemini Team et al., 2024) While there an on-going research regarding whether the encoder or decoder architecture is better tailored to predictive tasks (Nityasya et al., 2023; Li et al., 2023; Dukić & Snajder, 2024; Qorib et al., 2024), this work we focus on the question of how do we best apply decoder models to predictive tasks?
Generative models have been successfuly applied to numeric prediction, where number is generated autoregressively token-by-token. For example, Gruver et al. (2023); Requeima et al. (2024) considered zero-shot learning setup for time series prediction, Vacareanu et al. (2024a) experimented with zero-shot regression problems, and Liu & Low (2023); Yang al. (2023) considered the autoregressive finetuning over numerical targets applied to arithmetic tasks. Multiple prior works used predictive distribution from L.MM towards improving the prediction, including the median rule (Gruver et al., 2023; Requeima et al., 2024) and more broadly, Bayes- Bayes-optimal rule optimizing regression metric of choice (Lukasik al., 2024). The importance of tokenizing the numerical targets into individual digits has been raised by previous works (Liu & Low, 2023; Yang et al., 2023), and importance of prompt selection was analyzed by Requeima et (2024).
20

---

Published as a conference paper at ICLR 2025
C COMPARISON TO MBR-BASED FINE-TUNING
The RAFT loss may be contrast against ideas in the literature on Minimum Bayes Risk (MBR) prediction literature (Kaiser et al., 2000; Shannon, 2017; Prabhavalkar et al., 2018), which optimize non-regression metrics via approximation using sampled model predictions. For the squared loss, this may be formulated follows:
Definition 2. Define the sampled regression-avare loss as follows:
(MBR p(- x)) Byy-p(-1)) (y* float(y))
(13)
Compared to the loss Equation 9, the key difference that the expectation over the model outputs appears outside the square loss. A naïve empirical implementation of this objective requires explicitly sample responses from the model p(- z); this can be expensive and incur high variance. As with ÚRAIL Agrid), may instead consider practical variant that approximates the expectation using a restricted grid of targets dggii" C y/:
(PBR(P; ygria)
1
Σ
Σ p(str(v) | |x) (v" - y)²
(14)
N
(a. yEBaria
Even this variant has a notable disadvantage: the minimizer of Equation 14 is one-hot distribution that places all its probability mass on one of the targets in ygria C y:
Lemma 4. Let y (z) P(-a)( [y*] denote the Bayes-optimal prediction for input We assume (P( x) is supported numerical targets only. The minimizer of the approximate sampled regression-aware loss Equation 14 over all model distributions p(- 2x) of the form:
{1
if arg min y 11 y (2)ll2
plu 1w1x)) (x)
else
Therefore, the quality of the minimizer p(- ()) entirely depends on how well ygrid approximates the original target space y. For example, if Ygrid is a set of integers, the minimizer of Equation 14 will also be limited to predicting integers, even when numbers of arbitrary precision. As shown in Lemma 3, RAFT does not suffer from the loss of precision resulting from the use an approximate target space, and also avoids the high variance associated with sampling. In Table 5, experimentally verify better performance of RAFT over the MBR-based fine-tuning, and unless otherwise stated, we focus our attention to RAFT. Also, see $E.8 for analysis of the RAFT predictor in terms of the learnt distributions over tokens.
original target space y contains foating-poine
We next prove Lemma
Proof. Notice that:
Ey->P(1>) [Ey-m( [(y y²] = ~p((a)))
Σ E(lla)) (y - y²
Σ p(y(2)) By [(y y² )²
Since this convex combination, the optimal value achieved for p to be a one-hot vector with on the index arg min,, Ey -P((a)) (n) [VV y² ?² We thus want y that minimizes:
By By--p(-1w)) [v y )² By =F(1)) [D] + (v*)² 2 v²]
Equivalently, we want y that minimizes:
y² 2 y Byy-P((())))) [y*].
or equivalently:
3² - 2 y. Ey. -Py-~((112) [v-] + (By wB-1o) [w']]² (y By-~r(-le) lw'l)²
21

---

Published as a conference paper at ICLR 2025
Method/ablation
RMSE
Sampled regression aware (Definition (2))
0.98
Regression aware (Definition (1)
0.40
Table Root mean squared error (RMSE) on STSB across regression aware approaches and their variants on Gemma-2 9B.
We next provide a general version of Lemma 4. Lemma 5. The minimizer of the following objective:
Eg up" (- popttl [4y"", ([)",))],
for a loss function e y y -> R4, is one-hot distribution over targets such that all probability mass is on a target 9 which minimizes By (+1x)) [e(y",
Proof. The proof is elementary. Expanding the above objective:
By (1z) (ux() [K(y", v)] By up (1)) [e("",))
By- (1x) [E(y" "0"""""0 v)l p(ulæ)) dy
wEY By--p (1)) [8(y" 0) p(ulz) dy
vEN
(.x)) [E
(+xz) [8 y)] p(ylz) dy, = By--pr((
vEY
where the third step follows from the fact that minimizes Byy' (C(y"; )] the final step, p(-12)) is probability distribution that has point mass on
22

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EXPERIMENTAL DETAILS
We update all parameters during the fine-tuning. We summarize specific settings below. For Gemma-2, we use the following settings:
We use dropout rate 0.1 and batch size 16.
We train for 200A steps and select the best step using the held out validation set (see Table for details on the train/test/validation splits).
We use constant learning rate schedule. We select the learning rate value over the validation set from the values: 10-4 10-5, 10-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: 10 30 <2 10-3 decay rate
= 0.8.
For PaLM-2, we use the above settings, except we use batch size 64 and no dropout, and train for 5K steps and report the results from the last checkpoint.
For Moviel we use AdamW W optimizer and sweep learning rates from the range (10-1, 10-5, 10-6). We use cosine decay schedule for the learning rate, with 10K steps of warmup from learning rate 10
In Table 6, we report the prompts used in experiments, and in Table 7 we report the dataset statistics.
Dataset
Input prompt
Target range
STSB
What is the sentence similarity between the following two sentences measured on a scale to Sentence #1). (Sentence #2]. The similarity mensured scale of with O being unrelated and 5 being related is equal to
[0, 5]
Amazon reviews
product_eategory category: (product category) product_ title: (product title) review date: [review date) review headline: freview headline review hody: (review body) Question: In star rating 1, 4,5 4, 5, the higher the better, what would be the star rating of the above review? Please only give me the final rating and not need any explanations.
product_
2, 3.
MovieLens- IM
Instruction:
Predict the rating of target movie based the user's historical movie ratings. Rating History: [Rating history| | Candidate Item: [Candidate Item]. Output:
2. 3. 4,5
Synthetic (Original #1 from (Vacureanu et al., 2024a))
The to provide your
"outpul score' based "input
provide
[0. 9]
and only that, without any additional text. Input score: (Input score). Output score:
Table Prompts used for different datasets and the corresponding target ranges. Curly braces denote inputs specific to an input example. For Synthetic (Original #1 from (Vacareanu et al., 2024a)) we normalize the targets to correspond to [0, 9].
Dataset
Train size
Validation size
test size
Wireless
10,000
1,500
1,500
Personal care
10,000
1,500
1,500
Music
10,000
1,500
1,500
STSB
4,887
863
1,500
STSB 1k
1,000
863
1,500
MovieLens-IM
797,758
10,145
10,145
Synthetic (Original #1 from (Vacareamu et al., 2024a))
10,000
1,000
1,000
Table 7: Summary of dataset statisties.
23

---

Published as a conference paper at ICLR 2025
E ADDITIONAL EXPERIMENTAL RESULTS
We provide additional experimental results corroborating the findings in the main paper.
E.1
STSB-1K
In Table 8, we report results from sub-sampling the training set for STSB down to 1K examples. We find similar trends as observed in the case of other datasets in Table 3.
Model size
Zero-shot standard decoding
Standard fine-tuning
Predictive head
RAFT
standard decoding
2B
1.10
0.61110.01
0.01
0.58
9B
1.31
0.57=00.01
0.57=0.01
0.56++0.01
Table 8: RMSE on STSB-1 K across methods, and Gemma-2 models of varying sizes. Fine-tuning methods report mean std dev from model retraining.
E.2 SYNTHETIC DATA
For a simple setting, we consider synthetic regression dataset from Vacareanu et al. (2024a) referred to the Original #1 dataset by the authors. We report results in Table 9 and corroborate our observation from the language regression task experiment that RAFT improves over the predictive head approach when initialized from pre-trained checkpoint, and not when the model weights are initialized randomly. This provides additional support for our hypothesis that the alignment of RAFT to the next-token prediction pre-training task the underlying reason better performance over the predictive head.
initialiation
model size
autoregressive
predictive head
RAFT
Pre-trained checkpoint
2B
0.018
0.013
0.005
9B
0.017
0.017
0.006
2B
2.536
0.327
1.704
Random
9B
2.536
0.147
1.092
Table 9: The role of initialization to the pre-trained checkpoint on synthetic regression dataset from Vacareanu et al. (2024a) (the Original #1 dataset). We compare RMSE across different Gemma model sizes, and across different fine-tuning methods: autoregressive, predictive head and autoregressive regression aware. We corroborate our observation from language regression tasks that RAFT improves over the predictive head approach when initialized from a pre-trained checkpoint, and not when model weights are initialized randomly.
E.3 GEMMA-2 27B AND PALM-2 MODELS
We report results on Gemma-2 27B across all dataset in Table 10 and on PALM-2 models on STSB in Table 11, corroborating the findings of RAFT improving in most settings.
E.4 COMPARISON TO ENCODER BASED BASELINES
In Table 12, we report results with additional baselines that use prediction head over RoBERTa representation for the input sequence. We include: mean-pooling and CLS token variants, and both frozen RoBERTa weights and unfrozen weights in the fine-tuning. We also include the SMART method (Jiang al., 2019). In keeping with previous work, in addition to RMSE, also report performance on the Pearson and Spearman metrics for STSB (Jiang et al., 2019) (where available). We find RAFT surpass all the included baselines.
24

---

Published as a conference paper at ICLR 2025
dataset
zero-shot
autoregressive
predictive head
RAFT
Wireless
0.87
0.83
0.50
0.44
Personal care
0.94
0.89
0.50
0.47
Music
1.26
1.23
0.47
0.49
STSB 1k
1.29
0.60
0.56
0.55
STSB
1.29
0.62
0.51
0.48
Table 10: Comparison RMSE across datasets for Gemma-2 27B, In most cases, RAFT outperforms all other methods.
model family
model
autoregressive
predictive head
RAFT
PALM-2
1B
0.79
0.02
0.61 + 0.01
0.62 + 0.01
PALM-2
24B
0.63 0.03
0.56 1 0.00
0.53 1: 0.00
Table 11: Comparison of root mean squared error (RMSE) STSB across different PALM-2 model sizes, and across different fine-tuning methods: autoregressive, predictive head and autoregressive regression aware. Each model is for 3 times to obtain standard deviations. For PALM-2 IB, we find no difference in the performance of predictive head and RAFT, and both outperform the autoregressive approach. For PALM-2 24B, we see significant improvement from RAFT over both the autoregressive and the predictive head approaches.
E.5 CHOICES FOR TOKENS RAFT
Sensitivity to the grid size Following Lemma 3, ÚRAn. (2r) express any numerical target in the population limit, even with a coarse grid. Empirically, however, one might expect different grids to affect the results. We now assess this point on the Amazon reviews Wireless dataset, comparing the following choices for Agrid:
(1) { '1' '2', '3', '4', '5' (the default choice for RAFT),
(2) { '5' } (viz max(!)),
(3) { ('1',99' ) (the smallest and largest digit in V),
(4) { ('4', '5' (the two largest digits in for all datasets),
(5) { '7', '8", '9' (the two largest digits in V).
Method/ablation
Parameter count
RMSE
Pearson corr.
Spearman corr.
RoBERTa Base CLS
110M
0.64
90.84
90.59
RoBERTa Large CLS
356M
0.59
91.99
92.02
RoBERTa Large mean-pooling
356M
0.63
91.65
91.56
RoBERTa Large mean-pooling freeze
356M
1.08
72.72
74.16
RoBERTa Large CLS freeze
356M
1.30
56.48
54.76
SMART BERT (Jiang et al., 2019)
356M
90.00
89.40
SMART RoBERTa (Jiang et al., 2019)
356M
92.80
92.60
Gemma-2 2B RAFT
2B
0.54
93.55
93.22
Gemma-2 9B RAFT
9B
0.51
94.30
94.18
Table 12: Parameter counts and performance metrics on STSB across basclines. Results from SMART (Jiang et al., 2019) taken reported in the paper (RMSE was reported).
25

---

tends to converge Published as a conference paper at ICLR 2025
Dataset
Model size
'1'-15'
'5'
'1','9'
'4','5'
'7',8', "7'"""
Wireless
2B
0.47:10.01
0.48:00.1
0.48:10.00
0.48:00.00
0.48:10.01
9B
0.45::0.01
0.45=00.01
0.47##0.02
0.46:
0.47:00.01
Personal care
2B
0.49:10.00
0.49:+0.00
0.49+00.2
0.49+0.00
0.49+00.00
9B
0.47:10.01
0.48:10.01
0.48:10.02
0.47:10.00
0.47
Music
0.50:10.00
0.50
0.50110001
0.50
0.50110.01
0.46:10.01
0.47=00.00
0.48
0.47
0.48
Table 13: Comparison of RMSE (mean + std dev) across variants of RAFT with varying sizes of the grid ggrid The choice of '1'-'5' outperforms the alternatives.
In Table 13, report the results for different choices for Agrid- We find that, aligned with Lemma limiting the grid does not significantly impact the results. However, when comparing the number of steps convergence (see Table 21), we find that the default choice (1) in most faster to the best solution than other choices. One explanation for this finding is the following. As shown by Lukasik et (2024), the choice of '1' '2', '3', '4', '5' yields a reasonable result for the zero-shot RAIL approach, and thus it provides a good starting point for fine-tuning, Recall that the zero-shot RAIL corresponds to the RAFT approach step 0 of training, since the predictors for each are equivalent.
Sensitivity to the grid token indices The next question we pose is about the importance of strictly sticking to numeric tokens: what would happen if the RAFT predictor ýR (R)FT(Z)) used non- numeric tokens? To analyze this question, let us consider a more general form of the predictor:
ORAFT-NN(I)
Σ
p(token(y) zz) yy
(15)
y ygrid
where (y)en V denotes token of choice corresponding the numerical target y-
We keep Agria as composed of the digits '1', '2', '3', '4', '5' ), and use the predictor in Equation 15 with the following choices for tokens:
(1) token for each digit becomes alphabet token starting with 'a' and ending with 'a',
(2) each token is digit (i.e., '1' becomes '5', '2' becomes '4"),
(3) we only consider digit '5'.
As shown in Table 14, in most settings, the choice of digits { '1', '2', '3', '4', '5' is as performant any other choice. In certain settings, we even find large drop in performance (e.g.. the choice of characters or months). However, by enlarge, we find the results do not worsen significantly when different tokens are used.
E.6 CHOICES FOR LOSS, NORMALIZATION AND TOKEN INITIALIZATION RAFT
In Table 15 we compare the MSE loss to distillation style le log loss STSB. The target is scaled to be
between and for this set experiments. The log loss is defined as -y" log/i - (1 (1-y") (1-p)) where pi is the probability of digit 'I'. We find that both the MSE loss and the log loss yield similar results, and that scaling the range the targets does not negatively affect the performance.
We next analyze whether enforcing the probabilities over the grid to sum to I (by normalizing them by the sum the probabilities of all numbers in ygria) can improve the performance of RAFT. Table 16 shows that applying normalization to the probabilities for numbers in Vgrid does not significantly affect the results.
In Table 17, compare three initialization methods for the embedding of numbers in autoregressive and RAFT methods. The tokens used are of granularity of 0.1 for both autoregressive and RAFT methods, except for RAFT "/'-'5', where digits from to are used for the grid (granularity 1.0).
26

---

Published as a conference paper at ICLR 2025
dataset
model size
digits '''5'
characters 'a'='e'
reversed digits
'January' - 'May'
Wireless
2B
0.47 = 0.01
0.48 = 0.01
0.48 0.00
0.83 = 0.01
9B
0.45 + 0.00
0.46 0.00
0.46 0.01
0.47 = 0.01
Personal care
2B
0.49 + 0.00
0.48 + 0.00
0.48 = 0.01
0.85 + 0.00
9B
0.47 + 0.01
0.47 + 0.00
0.48 = 0.01
0.48 + 0.01
Music
2B
0.50 0.00
0.50 0.01
0.50 0.01
0.50 0.00
0.46 0.00
0.48 0.01
0.47 = 0.01
0.61 0.25
Table 14: Comparison of RMSE (mean + std dev) across variants of RAFT where different tokens are used for the prediction formula of RAFT. Each experiment repeated for times. In most settings. the choice of digits '1'-5'' is at least as performant as any other choice. In certain settings, we find a large drop in performance compared to the choice of digits (i.e., months "Jauaar'' "MMay'.
unscaled targets
targets scaled [0, 1]
model size
predictive head
RAFT
RAFT '0" and '1'
log loss
0.61
0.62 0.01
0.63 + 0.01
0.63 0.01
24B
0.56 0.00
0.53 0.00
0.54 1 0.01
0.53 0.00
Comparision of RMSE (mean + std dev) on the STSB dataset for MSE loss and log loss when the target is scaled to be between 0 and 1. The models are PALM-2 1B and 24B models. The MSE loss uses digit 'O' and '1' to compute the predicted value. The log loss is in the form -y" log 1 - y²) log(1 - P1) where pi is the probability of digit
Table 15:
Overall, we find that RAFT is less sensitive to the initialization method than the autoregressive approach.
In Table 18, we experiment with random initialization for the tokens in the RAFT grid. We find random initialization lead to worse results compared to using pre-trained token embeddings.
To this end, we vary the granularity of Agrid by constructing a list of equally spaced numbers covering the range of y. For example, choosing granularity to be 0.1 when y = [0. 5| yields Ygrid [ 1'0.0', '0 1', '0.2' '4.9', '5.0' 0' }. In our implementation, all numbers in Ygrid are represented by single tokens that we add to the vocabulary. We initialize the token embedding with either the First the Average method. In the First method, we initialize the embedding with the embedding of the first digit of the number (e-g. token '0' embedding for the number '")."). In the Average method, we initialize the embedding with the average of the embedding from the constituent tokens (c.g. use the average embeding of token '0', and I' for the number '0.1".) We report the results on the STSB datasets with PALM-2 IB model in Table 19 and find no significant difference in the results across different choices for the granularity of Vgrid-
We next analyze the impact of the choice of y Egrid in computing the RAIL predictor.
Additionally, in Table 19 we also include the autoregressive method utilizing additional tokens from Ygrid as constructed with the methodology outlined above (i.e., with varying granularity). Here, contrary to RAFT, we find the initialization method to affect the results, with First performing better than Average. Note that the autoregressive method is equivalent to the generative classification from (Fernandes et al., 2023), where the classes correspond to the numbers from the grid.
Lastly, we would like to note that in the case of generative classification, there is trade-off between how fine-grained the grid is and how many examples per token are observed during training. In particular, if the classes are too coarse, we observe a loss in performance. On the other hand, the classes are too fine-grad finengrained, there may be insufficient training examples per label to learn the embeddings for new tokens. For example, with granularity 0.05, 17 out of the 101 numbers in the grid do not appear in the training data. This also lead to loss in performance.
27

---

Published as a conference paper at ICLR 2025
model size
not normalized
normalized
LB
0.63 + 0.00
0.63 =: 0.02
2/B
0.53 + 0.00
0.53 + 0.01
Table 16: The effect normalization of grid token probabilities in RAFT on the STSB dataset. The models are PALM-2 1B and 24B models.
initialization
autoregressive
RAFT
RAFT '1'-1''
Zero
1.20 + 0.01
0.62 + 0.00
0.63 + 0.00
First
0.80 + 0.02
0.63 + 0.00
0.62 = 0.02
Average
0.98 + 0.01
0.64 + 0.01
0.61 + 0.01
Table 17: Comparison of different initialization methods for embeddings for tokens corresponding to numbers autoregressive and RAFT methods on STSB dataset with PALM-2 1B model. We consider granularity 0.1 for constructing the tokens for autoregressive and RAFT reported in the first 2 columns, whereas the final column reports RAFT with tokens to 5. Zero denotes initialization with O values, First denotes the initialization with the embedding corresponding to the first token of the number (e.g. use token 'O' embedding for the number '0.1"), and Average denotes the initialization with the embedding corresponding to the different tokens in the number (e.g. average embeddings for tokens '0', and '1' when initializing the embedding for the number '0.1"). RAFT is less sensitive to the initialization method. We find the performance does significantly worsen with different initialization methods compared to using the pre-trained token embeddings at initialization (see Table 15).
RAFT with '1'-55'
RAFT with '5'
baseline
random
baseline
random
0.63 + 0.00 0.66 0.02 0.62 + 0.01 0.64 + 0.01
Table 18: RMSE across different initialization (baseline and random) of the tokens in RAFT. We find random initialization of tokens in the RAFT grid to hurt the performance. The model is the PALM-2 1B model.
autoregressive
RAFT '1'-5'
granularity
First
Average
First
Average
0.05
0.85 + 0.01
1.11 + 0.01
0.61 + 0.01
0.63 + 0.00
0.1
0.80 = 0.02
0.98 0.01
0.63 0.00
0.64
0.2
0.83 = 0.02
0.89
0.64
0.66
0.5
0.84 + 0.01
0.87
0.63
0.63
1.0
0.81 = 0.01
0.85
0.62 0.02
0.61 0.01
Table 19: Comparison of RMSE on autoregressive and RAFT methods with different granularity of tokens. We append tokens that represent numbers with various granularity. For example, with granularity 0.05, the tokens are '0.00, '0.05, ..., '5.00'. For autoregressive, this is equivalent generative classification method from (Fernandes et al., 2023). First and Average denotes different ways to initialize the embeddings of the tokens, with details described in the caption of Table 17, The model is the PALM-2 IB model.
28

---

Published as a conference paper at ICLR 2025
model size
'1'='5'
'5'
'1'-99'
'b'-'p'
'5'-11'
"Jnn. '- "May"
IB
0.63
0.00
0.62 = 0.01
0.63 = 0.01
0.61 = 0.01
0.62 = 0.01
0.64 + 0.01
24B
0.53 0.00
0.53 = 0.01
0.53 = 0.00
0.54 = 0.02
0.54
0.00
0.55 + 0.01
Table 20: Comparison of RMSE (mean = std dev) on the STSB dataset across variants of RAFT where different tokens are used to construct the grid for computing the RAFT predictor. The models are the PALM-2 24B and IB models.
dataset
model size
'1'-'5'
'5'
'1','9'
'4','5"
'7", 17', ''','9'
Wireless
2B
1,000
3,000
2,000
1,600
2,200
9B
1,600
2,800
3,400
3,200
4,200
Personal care
2B
2,200
2,200
2,400
1,600
2,200
9B
4,200
2,000
6,200
400
2,400
Music
2B
800
1,200
1,800
1,800
1,800
9B
1,000
2,200
3,000
1,600
3,000
Table 21: Comparison of the number of steps to convergence of RAFT where different number of numerical targets are used for the grid RAFT. Results on the Amazon Wireless dataset.
Steps Convergence regeeee
Predictive Head
4,000
RAFT
3,000
2,000
1,000
10%
20%
30%
40%
50%
60%
70%
80%
90%
100%
Percentage of Training Data
Figure 4: Comparison of the number steps to convergence RAFT and predictive head, with different percentage of the training set.
E.7 CONVERGENCE SPEED OF RAFT
In Table 21 we report the number of training steps to convergence to the best result the held out validation set of different methods. We can see that RAFT with digits '1"-'5' majority of cases converges faster than other choices for the grid.
In Figure 4 we compare the number of steps to convergence RAFT and predictive head on STSB. The figure shows that RAFT converges with fewer number of steps across different percentages of training data used for training.
In Figure we report RMSE the test set function of the training step. We find that compared to predictive head, RAFT starts from a much lower RMSE score (as it corresponds to the RAIL method before training), and then converges in fewer steps.
29

---

Published as a conference paper at ICLR 2025
4.0
Predictive head
3.0
RMSE
2.0
1.0
0.0
1000
2000
3000
4000
Training step
Figure 5: RMSE on the test set function of the training step on Amazon Wireless for Gemma-2 2B. Compared to predictive head, RAFT starts from a much lower RMSE at step 0, and then converges to its best RMSE in fewer steps.
300
Frequency
200
-1.0
-0.5
0.5
Per-example entropy change
Figure 6: The histogram of per-example change (after vs. before fine-tuning) entropy over the probabilities for the digits in the RAFT predictor. The entropy on average increases after the RAFT fine-tuning, meaning that the probabilities are overall more spread over the digits than prior to RAFT fine-tuning.
E.8 DISTRIBUTION OVER TOKENS IN THE RAFT PREDICTOR
We next investigate the distribution over tokens in the RAFT predictor. We find that, while the error decreases with training, the entropy increases after RAFT training, as we show in Figure 6. This corresponds to the model on average spreading the probabilities over tokens more than prior to fine- tuning, shown for specific examples in Figure 7. 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 4).
30

---

Published as a conference paper at ICLR 2025
Probability Untrained
Probability Trained
target prediction
target
prediction
Probability
Probabiltty
0.6
0.4
0.2
0.0
0.0
Digit
Digit
(a) Example
Probability Untrained
Probability Trained
1.0
1.0
target
target
0.8
prediction
0.8
prediction
Probability
0.6
Probability
0.6
0.4
0.4
0.2
0.2
0.0
0.0
1
2
4
1
2
4
Digit
Digit
(b) Example #2
Figure 7: Comparison the distribution over digit tokens RAFT predictor (Gemma-2 2B on STSB) before and after fine-tuning, We find the entropy over the digit token probabilities overall increases after fine-tuning with the RAFT objective.
31


Reviews:
[{"review_id": "wOWWJdzY22", "reviewer": "Reviewer_7cP3", "summary": "The paper addresses regression tasks for language models, specifically using decoder transformer models to predict a numerical value $y$ based on an input sequence $x$. In this setup, two standard methods for regression are examined: (i) quantizing $y$ into discrete bins, where values are mapped to tokens (e.g., sentiment values represented by tokens 1, 2, 3, 4, 5) or (ii) training a multilayer perceptron (MLP) head on top of a decoder language model with an $L_2$ loss.\n\nEach of these approaches has known limitations. Quantization lacks sensitivity to numerical proximity, treating all incorrect tokens as equally wrong (e.g., predicting 4 is treated the same as predicting 5, regardless of proximity to a true value of 3). Meanwhile, an MLP head with $L_2$ loss does not fully leverage the language model’s pretrained token prediction capabilities. To address these issues, the paper studies the MALI estimator, which uses a grid of tokens and calculates $\\sum_{y \\in \\mathbf{y}_{\\text{grid}}} p(\\text{str}(y) | x ) . y $ \n\nAlthough MALI typically serves as a post-training estimator, the authors explore its application in fine-tuning language models for regression tasks, providing extensive comparisons and ablation studies to validate its performance against existing methods.", "questions": "1. **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.\n2. **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.\n3. **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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-organized, with a clearly defined objective and comprehensive experimental comparisons, making it accessible and straightforward to follow. I enjoyed reading the paper.", "weaknesses": "1. **Performance Drop on Alternative Grids (Lines 902–903 and L531)**: You note a performance drop when using a grid of letters a–e, yet Table 11 shows similar MSE to the numerical grid. Could you clarify this inconsistency?\n\n2. **Lemma 1 (Line 187)**: Based on the appendix, does the norm in Lemma 1 refer to an $L_1$ norm rather than $L_2$? The crafted example in the appendix (where a probability mass of $\\epsilon / 2$ is positioned at 0 and $y_\\text{max}$ implies that the $L_1$ loss should satisfy $E_x [ \\| P(. | x)  - p(.|x)  \\| ] <= \\epsilon $ (the norm being L1), given that line 747 defines $ | \\mathbb{P}(. | x) - p(.|x) |_{1} < \\epsilon $.\n\n3. **Token Granularity and Grid Complexity**:\n   In the RAFT experiments, was the grid composed solely of single tokens? If so, does this also apply to the months-named grid?\n   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.\n4. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5ePvz4sOnW", "reviewer": "Reviewer_Tf6L", "summary": "This paper provides a comparison among several methods for autoregressive regression with LLMs, where the goal is to predict numerical value given the text input (e.g., text similarity). The authors propose a new fine-tuning approach, namely RAFT, which computes squared loss directly on the grid of possible outcomes without explicit sampling from the model. Empirical results show the effectiveness of the proposed approach across various datasets.", "questions": "- 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). --\n- In Table 3, the predictive head closely follows the results of RAFT. Could you please elaborate more on that? (more than described in 429-430)\n- 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", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Paper (in general) is well-written and easy to follow (even though there are still small things that require further clarification)\n- Simple and effective idea \n- Shown results indicate improvements across different tasks\n- Additional analysis on several aspects (grid, pre-training, masking)", "weaknesses": "The predictive head approach closely follows the RAFT; however, there is no extended discussion of what sets RAFT apart.\n\nNote that the below weaknesses are mostly minor.\n\nThe 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\n- Since some things are still unclear to me in the experiment setup (please see question), I am hesitant to assign a score", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KQzTTxwn2t", "reviewer": "Reviewer_CzSf", "summary": "This work introduces a new fine-tuning method for autoregressive language models to improve their performance on regression tasks, such as sentiment analysis, semantic similarity, and quality estimation for translation. \n\nThis work argues that the cross-entropy loss is not suited to regression tasks, since small deviation of the model from the ground truth can lead to high regression error. Indeed, the cross-entropy objective does not take into account the numerical value associated with text tokens. Instead, authors propose a new fine-tuning method (RAFT) based on the \"Metric-aware LLM inference for regression\" (MALI) approach.", "questions": "- 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). There is one small experiment comparing the sampled regression aware approach (table 8), but it is only in the appendix. 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.\n\n- 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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The research question is relevant and interesting.\n - The presentation of MALI (previous work) is clear.\n- The counter-example of Lemma 1 is good. Since it's short, I would include it in the main body if possible.\n- The experiment on using semantically unrelated tokens for the regression task (Table 11) is interesting.", "weaknesses": "- On a first read, I misunderstood the implications of Lemma 1. I think it would be worth slightly rewriting the paragraph after Lemma 1. In particular, I would replace the sentence \n> Intuitively, log-perplexity fine-tuning treats all “wrong” predictions the same, *and does not account for the magnitude of their differences from the ground-truth target*.\n\nby\n\n> Intuitively, log-perplexity fine-tuning treats all “wrong” predictions the same, *as it is unaware of difference in magnitude of the numerical values represented by the tokens. For example, assuming that \"100\" and \"1000\" are represented with a single token, placing $\\epsilon$ too much mass on tokens representing \"100\" is penalized similarly as placing  $\\epsilon$ too much mass on the token \"1000\".\n\nIf you have a better rephrasing, you can of course use something else, but I think the aforementioned sentence can be improved (e.g. using my example).\n\n- 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?\n\n- **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):\n    - Prediction head regression over a mean-pooling of RoBERTa representation (over sequence length)\n    - Prediction head regression over the output for the CLS token of RoBERTa model\n    - Both variants with either frozen RoBERTa weights, as well as unfrozen weights\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IFRJFRLagZ", "reviewer": "Reviewer_AkBL", "summary": "This paper explores two approaches: (1) fine-tuning an LLM using the log-ppl loss with autoregressive sampling during inference, and (2) adding a predictive head and fine-tuning it with an appropriate loss function. Following an in-depth analysis, this paper introduces RAFT, a regression-aware, Bayes-optimal method for optimal LLM fine-tuning. Extensive experiments demonstrate that RAFT consistently outperforms baseline methods across various datasets and models.", "questions": "Could you clarify the scenarios in which a predictive head performs better versus those where autoregressive sampling is more effective? Also, is RAFT the optimal choice in all situations, or are there cases where other methods might be preferable?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "1. This paper provides a well-rounded formulation for LLM-based regression approaches.\n2. It presents thorough evaluations across diverse tasks (e.g., Wireless, Music, Personal Care) and models (e.g., Gemma and PaLM).\n3. The proposed RAFT method demonstrates effectiveness across these benchmarks.", "weaknesses": "1. While this paper presents a strong approach, it addresses a relatively niche problem. In my view, 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)