Bridging the Training-Inference Gap in LLMs by Leveraging Self-Generated Tokens
Authors: Zhepeng Cen
Yao Liu
Siliang Zeng
Pratik Chaudhari
Huzefa Rangwala
George Karypis
Rasool Fakoor
Venue: Accepted by TMLR ()

Abstract: Language models are often trained to maximize the likelihood of the next token given past tokens in the training dataset. However, during inference time, they are utilized differently, generating text sequentially and auto-regressively by using previously \emph{generated} tokens as input to predict the next one. Marginal differences in predictions at each step can cascade over successive steps, resulting in different distributions from what the models were trained for and potentially leading to unpredictable behavior. This paper proposes two simple approaches based on model own generation to address this discrepancy between the training and inference time. Our first approach is Batch-Scheduled Sampling, where, during training, we stochastically choose between the ground-truth token from the dataset and the model's own generated token as input to predict the next token. This is done in an offline manner, modifying the context window by interleaving ground-truth tokens with those generated by the model. Our second approach is Reference-Answer-based Correction, where we explicitly incorporate a self-correction capability into the model during training. This enables the model to effectively self-correct the gaps between the generated sequences and the ground truth data without relying on an external oracle model. By incorporating our proposed strategies during training, we have observed an overall improvement in performance compared to baseline methods, as demonstrated by our extensive experiments using summarization, general question-answering, and math question-answering tasks.

Published in Transactions on Machine Learning Research (01/2025)
Bridging the Training-Inference Gap in LLMs by Leveraging Self-Generated Tokens
Zhepeng Cen
zoenSandrew.cmu. cmw. edu
Cornegie Mellon University
Yao Liu
ycoltuaiSoaazzo..com
Amazon Web Services
Siliang Zeng University of Minnesota Twin Cities
zeng017763umn.edle
Pratik Chaudhari
prticBamazon.com
Amazon Web Services
Huzefa Rangwala Amazon Web Services
rhuzefa@amazon.com
George Karypis Amazon Web Services
gkargpis@emazon.com
Rasool Fakoor Amazon Web Services
faboor@amazon.com
Reviewed OpenReview: https: openreeveee net/ forum? id wd=pWSrm3oP88
Abstract
Language models are often trained to maximize the likelihood of the next token given past tokens in the training dataset. However, during inference time, they are utilized differently, generating text sequentially and ruto-regressively by using previously generated tokens as input predict the next one. Marginal differences in predictions at each step can cascade over successive steps, resulting different distributions from what the models were trained for and potentially leading to umpredictable behavior. This paper proposes two simple approaches based on model own generation to address this discrepancy between the training and inference time. Our first approach is Batch-Schoduled Sampling, where, during training, we stochastically choose between the ground-truth token from the dataset and the model's own generated token as input to predict the next token. This done in an offline manner, modifying the context window by interleaving ground-truth tokens with those generated by Correction, where we explicitly
the model. Our second approach is wwwerencoo-Anee---baeee incorporate self-correction capability into the model during training. This enables the model to effectively self-correct the gaps between the generated sequences and the ground truth data without relying on external oracle model. By incorporating our proposed strategies during training, we have observed an overall improvement in performance compared to baseline methods, as demonstrated by our extensive experiments using summarization, tasks.
general question-answering, and math
1 Introduction
The common approach to training auto-regressive models known as teacher forcing (Williams Zipser, 1989). In this method, the ground truth token from the previous time step is utilized the input for the

---

Published in Transactions on Machine Learning Research (01/2025)
model at the current time step. This technique allows the model to learn relationships between tokens more effectively, facilitating laster convergence during the training process. While this training technique has been widely adopted in previous works (Cho et al., 2014; Gregor et al., 2014; Bahdanau et al., 2015; Vinyals et al., 2015; Parmar et al,, 2018; Fakoor et al., 2017: 2018; 2020; Esser et al., 2020; Chang et al., 2022; Li et al., 2024; Liu et al., 2024), it can also lead to overfitting and undesirable behavior (Bengio et al., 2015; Bachmann & Nagarajan, 2024). In particular, when a model is solely trained on the provided ground-truth tokens, it may fail to behave reliably when encountering its own generations later, which can include unseen tokens. This occurs because, during inference, the model must rely solely its own previous generations/ generations/preedction prediccions rather than the actual ground truth; hence, any small error can propagate through subsequent time steps, resulting in compounding errors and, therefore, unpredictable behavior. This issue is commonly known as exposure bias (Bengio et al., 2015; Ranzalo et al., 2015; Schmiddt, 2019; He et al., 2019).
Current transformer-based (Vaswani et al., 2017) methods for aligning large language models (LLMs) with human preferences, such as Reinforcement Learning from Human Foedback (RLHF), are also auto-regressive models (Ziegler et. al., 2020; Ouyang el. al., 2022; AT, 2023; Gemini- 2024). Supervised Fine-Tuning (SFT) is employed to fine-tune the pre-trained model using human demonstrations (utilizzng teacher forcing) as the initial step of this alignment method. This fine-tuned model is then further refined using reinforcement learning through a reward model that serves as a proxy for human preferences (Ouyang et al., 2022). Since SFT is used as a standalone alignment method and also serves as the initialization for other steps in RLHF, such as the RL step that exclusively uses its own generations, having an SFT model that can utilize its own generations during training to develop tolerance for the shift between ground-truth data and its own generations seems both necessary and increasingly important. Therefore, it is important develop an SFT model that is not solely trained on ground-truth data, which is more prone to overfitting and, consequently, to training-inferent discrepancies, but also incorporates its own self-generated data during training to closely align with what it encounters during inference.
To bridge the discrepancics between how the model is trained and how it is used during inference, we propose two approaches that both leverage the model's own generations to address this issue. First, we train the model in manner akin to how samples are generated during inference. Specifically, instead of relying solely on ground-truth tokens during training (teacher forcing), we expose the model to its own generations, adopting the scheduled sampling (Bengio et al., 2015) for LLMs but in an offiine and batch manner, particularly during SFT training. While scheduled sampling and its variants have been popular with smaller models, especially recurrent- based ones (Lamb et al., 2016; Goyal et al., 2017; Li et al., 2020), they have not been widely adopted for LLMs due to their complexity and the practical challenges of incorporating them during training. We propose an offline and batch version of scheduled sampling called Batch-scheduled Sampling (BASH), which is more practical and can be easily adopted during SFT training. Despite its effectiveness, one of the main problems with BASH is that the model's own generated tokens at each time step can deviate from the ground-truth tokens. When interleaved with ground-truth tokens, the resulting sequence can differ significantly from the ground truth, complicating training and ultimately leading to slower training time, which negatively impacts results. To address this, we introduce Reference- Answerr-ae Correction (RAC), where we explicitly incorporate a self-correction capability into the model. This approach resembles Dataset Aggregation (Ross et al., 2011) in imitation learning but employs a self-supervised objective without relying on an external oracle model.
To evaluate the effectiveness of our proposed approaches, we provide a comprehensive empirical comparison and ablation study of our method across a range of standard benchmark tasks, such summarization (Stiennon et al., 2020) and general (Ding et al., 2023) and math question-answerigg tasks (Cobbe et al., 2021; Hendryeks et al., 2021). This evaluation is conducted in settings where we have access only to the human demonstrations data, which is primarily applicable to the SFT stage. Our results, based on win rates against the reference for the summarization task and length-controlled win rates (Dubois et al., 2024a) on the AlpacaEval 2.0 benchmark (Dubois et al., 2024b) for QA tasks, clearly demonstrate that our proposed approaches are effective in improving performance. Additionally, we demonstrate that initializing model trained using our approach, followed by fino-tuning with preference data through direct preference alignment method (Rafailov et al., 2024), leads to better results compared to initializing with a standard SFT model.
2

---

Published in Transactions on Machine Learning Research (01/2025)
2 Background
Given pre-trained auto-regressive language model parameterized by w, objective is fine-tune this model using human demonstration (a.k.a. expert) data² to ensure that it generales text aligned with the demonstration data. Consider a dataset D = {(4,, (:),, where and yi represent query/prompp and its corresponding continuation, respectively. Each example is a sequence of tokens x; ..;,...,, x7) and yi = .y.....,,"" ) where and indicate the lengths of the prompt and continuation, respectively".
To fine-tune this auto-regressive model with D, we employ a maximum-likelihon, objective, defined as follows:
Σ log Po (د,
D
(=
where 0 is the model parameters initialized from w, y² .'..y. 1) 30 shows the beginning-of-sentenn token, L denotes length of the contimuation sequence and x = (z'...., x²) indicates a prompt of length T. To maintain consistency with current. literature, we refer to this method as SET (Ouyang et al., 2022).
2.1 Auto-regressive Generation
After the model is trained, we use to perform conditional auto-regressive generation by specifying prefix sequence I::T (i.e. prompt/query) and sampling the remaining sequence (i.e. continuation) one token at the time, using the prefix and the previously generated tokens context:
po(12x 2 xk), k 0, ... K
(2)
where the current context looks as follows:
T-k+1 context
prompt
generated takens
This process is repeated until the stopping condition is met (i..., the maximum sequence length is reached the end-of-sentence token is generated). Throughout the remainder of this paper, whenever generation is mentioned, we will utilize the auto-regressive approach explained above.
3 Methods
The objective function in Eq. (1) is known as teacher-forcing (Williams & Zipser, 1989) learning method for rego-segressvve models, where it utilizes the ground -outh token from the previous time step as input to the model at the current time step. This can help the model learn more quickly, but it also lead to overfitting, In particular, models trained exclusively and undesirable behavior when faced with their own generated tokens during inference, especially if the generated token was not seen during training. Motivated by these challenges and to bridge the gap between training and inforence, we propose two approaches in the following sections to address the shortcomings of the teacher-forcing method, making training resemble inference as closely as possible.
ground-truth data might exhibit inconsistent
3.1 Batch-schedules Sampling
Scheduled sampling. To align the model's behavior during training with how it functions during inference, and to account for its auto-regressive nature, update Eq. (1) that it consumes its own generated tokens in addition to the ground truth tokens during training, but in controlled manner. To achieve this, the
also be syuthetically generated from 2To simplify notation, drop the subscript from equations whenever it clears from the context.
Note that demonstration data does not always need from humans,
another model.
3

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm Batch-scheduled Sampling (BASH)
Input: Pro-trained model w, training datatsct D.
Initialize x-u
2: for k == .,,,..., K1 do
3: Sample mini-batch B = ((o,y)}~1
4: VoJsrT(0)) -- Vo B
Σ (a.)) E, logpoly' x. (
+ a VoJrr
6: end for
7: for iteration 1......, , H do
8: Construct dataset D, in oflline manner as explained in Sec. 3.1
9: for h = 1,2,..... .0,...,,y K2 do
10: Sample mini-batch
{(x, v. 0) D, Σ (E E, log (5)
11:
VoJsFr(() V
|.|
12:
VoJBASH t V (.. E, log Doyp 1x, )0)
13: - 0 a vaJSFT(0) + VoJRASH(O))
14: end for
15: end for
Output: 0.
scheduled sampling (SCS) method from Bengio et al. (2015) can be utilized. The goal of the scheduled sampling method to stochastically include the model's generated tokens in an online manner during training:
Jscs (0)
Σ
Σ
log po (y² (
(3)
D
(=
where /0 mixture of ground truth tokens and the model's own generations. To construct g, we first sample 2 from the model output given the previous context ك input: 2 ~ go() Then, g is created by randomly selecting either 2 y² the token with probability B:
2) y otherwise
with prob. B
(4)
Here, I represents the mixing factor: when it equals 0, 9 is exactly the same y. However, when B equals 1, g becomes completely different from as every token is auto-regressively generated by the model, and the ground truth contimuation tokens are completely discarded. Notice that, the context g in Jscs(f) is a function of the current parameter 0, though we do not propagate the gradient through it in the standard SCS.
Remark 1 (Scheduled sampling is not scalable). While SCS (and its variants) has proven effective with small models (Bengio et al., 2015; Lamb et al., 2016; Goyal et al., 2017; Li et al., 2020), particularly recurrent-basee networks, its computational complexity outweighs its benefits for LLMs. This is why it has not been widely adopted for training large models. Specifically, because LLMs require distributed training across many GPUs, switching between training and inference modes per tokens would not only significantly slow down the training but also lead significant GPU under-utilization and memory related issues.
Mihaylova & Martins (2019) also highlighted the difficulty of applying scheduled sampling to transformer- based models and proposed structural changes to the transformer through a two-pass decoding strategy. In the first pass, model predictions are generated without accumulating gradients, and in the second pass, the generated data mixed with the ground truth update the model. Although this approach shows promise in some tasks, it has not been applied LLMs. This due to the required structural changes the model's architecture and the discussed scalability challenges.
4

---

Published in Transactions on Machine Learning Research (01/2025)
Question:
Nelly is very pleased with the painting she bought at the auction. tells daughter that outhid rival Joe by paying $2000 more than thrice his bid. If Nelly get
RAC Correction
Joe's $160,000, how painting?
Generated ... response label
S
S
SFT-moddl-eneeraeee Answer: Nelly paid $2000 more
which
meuns she paid 3 $160000 + $2000 = 484000
The answer is: 484000
Figure How does RAC correct mistakes in model-generated responses? In this example, SFT model makes mistake in calculating 3 $160, 000 + $2, 000, shown in yellow. However, RAC corrects the error by replacing the wrong token, 4, with the correct token, 2. This is achieved by forcing the model to fit ž that differ from the original generated response (highlighted in purple), enabling it to self-correct. This example is based the GSM8K dataset (Cobbe et al., 2021).
Batch-scheduled Sampling. To mitigate the limitations discussed in SCS for large models, propose simple yet effective offline approach. A new dataset, D, = {(z,, Ut. b.)) is created offline between training iterations by current model". While 2; remains identical to the original dataset D, Ds represents a "mixed" continuation constructed by stochastically combining ground truth tokens with generated ones from the model, as described in the previous section (see Eq. (4)). It important to emphasize that the offline and batch nature of BASH minimizes the cost of switching between training and inference times by creating D, only between each iteration of training. In contrast. SCS switches between training and inference modes at every gradient step, resulting in significant slowdowns in training time and practical challenges related to distributed training::
JBAHH(O)
Σ
log Po 2, (.)
(5)
[Da
(N.D))ED, j=0
In practice, to balance optimizing the objective function using generated data and ground-truth data, we combine Eq. (1) and Eq. (5) (i.e. J= (B)s)n)))))),, where the only difference is training data. Moreover, to minimize discrepancies between the current model and the one used for generation, this procedure can be repeated across different training itcrations using the updated model from the most recent itcration. It is important to note that since the model can generate random tokens at the beginning of training, which complicates the training process, we start by first optimizing Eq. (1) alone. After a few iterations, we then begin including BASH. These steps are detailed Algorithm 1.
It is worth noting that our method offers distinct advantage over existing approaches, it does not necessitate any modifications to the model structure, unlike methods such as Mihaylova & Martins (2019). Moreover, our method is scalable and can be applied to LLMs as shown in our experiments.
Remark 2 (Parameter 3 should be chosen to be small). Depending on the value of 3, there may be a distribution mismatch between the ground truth sequences y and the mixed ones g. This gap can become larger as the value of B approaches 1, as g increasingly differs from y, not just in terms of a few tokens, but at the sequence level. Also, since scheduled sampling can result in biased estimator (Huszár, 2015; Lamb et al., 2016), the training can get more harder, the model needs to learn to fit the ground truth data with the altered continuations which are far from each others. Therefore, it is important to keep the value of B small to avoid making the opt imizaatiioon problem harder to solve 5.
"If we drop i from D. ((1,) Ve (.)), retrieve the original D. Practically, this is useful during training, as D, also function the original (when &s is dropped), eliminating the need for separate data loaders train both SFT and BASH. "See pttpps//hugginge
ct//oc////cecccccccccc taaece//taaafforreeee wwweeerreee///oooereeeeeee perpperformm ormance to learn about the challenges distributed training. Tn our experiments this paper, we select, be equal to 0.2.
5

---

Published in Transactions on Machine Learning Research (01/2025)
3.2 Reference- Ansseer-aadd Correction
One of the main problems with scheduled sampling, which mixes ground truth tokens y and the model's gencrations 2², is that the resulting sequence g can diverge significantly from the ground truth sequence y- This means that the generated sequence g not only differs in individual tokens but also conveys a context and meaning that deviate substantially from Importantly, since the scheduled sampling approach (both online and offline) results in biased estimator (see Remark 2), such significant discrepancies between the ground truth and generated sequences can further complicate training and ultimately lead to slower training/comie reaiiing/convergence time. To give the model ability to recover in such scenarios, we propose Reference- Answer-based Correction (RAC), where we explicitly incorporate self-correction capability into the model.
To build RAC, we first construct dataset D, {(2;, Vi, 2,, 2,) in an offine manner. Here, denotes sequence of model's ' own generated lokens (see Eq. (2)) and E denotes new target/label sequence, also composed of the model's own generated tokens, which is constructed by greedy sampling from model 0 at each time step:
x = arg max mox
where f is prompt template The reason for constructing z in this way to enable the model leverage its self-correction capability by incorporating the grundd-ruth answer y into the input context. Hence, the maximum-lkkclhoodd objective for RAC can be written as follows:
JRAC.
Σ
Σ
log 1 21)
(6)
[Dr]
(D... J=0
In RAC, the model attempts to maximize the likelihood of correction label z conditioned on the original prompt æ and its own generated tokens 2. This contrasts with BASH, where 9 is used instead 2, and the target token is always yd, i.e,, Poly² x, gz)). It important to note that is generated using greedy sampling, conditioned on f(r, y) and 2. Conditioning on f y) rather than just allows for additional context and guidance during the generation process.
One issue that arises in Eq. (6) is when 2) becomes identical to the generated token zf, When this occurs, it can lead to model collapse (Shumailov et al., 2024), as the model is likely to only learn trivial solutions. To mitigate this issue, we mask out such tokens and rewrite the objective function as follows:
1
Σ E1(2 + +2)' >2) 11,
TRA
Drl (L.x,3)ED, f=0
Similar to BASH's algorithm, we combine Eq. (1) and Eq. (7) by first optimizing Eq. (1) alone for several iterations to avoid relying model generation the start of training. After this initial phase, then include RAC objective function and jointly optimize with SFT, i.c. J = Jat TR)c(O). These steps are detailed in Algorithm 2. Fig. 1 shows an example of how the self-correction capability of RAC helps in producing correct results.
4 Experiments
In this section, we present a comprehensive empirical comparison of proposed methods across range of standard benchmark tasks, including summarization and question answering (QA). These results show that the effectiveness and robustness of our approaches across different settings. See also Appendix A.2 for a description of our experiments and Appendix A.3 for more ablation studies.
6One example for prompt template f(, "I will give you question and reference respouse. You need to give response hased the reference response. Question: Reference response: y". See Appendix A.1 for the actual templates used in experiments

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm CnsereenceAAswer-beeeee Correction (RAC)
Input: Pro-trained model w, training datasct D.
Initialize x-u
2: for k= .,,,..., K1 do
3: Sample mini-batch B = ((x,,))~1
4: VoJsrT(0) -- Vo B
Σ (a.)) 2, logpoly' x., (5)
+ a VoJrr
6: end for
7: for iteration 1, ..... H do
8: D, -
9: for i= 1,2, ,N do
10: Generate model's response 2; by 2 ~ po(-11s, 1,2....
11: Generate RAC label & by & = arg max; (1)) ), 1,2, 1,
12: D, + D, ((Ts, yi, 24, 2)}
13: end for
14: for k = 1,2, ...., K2 do
15:
Sample mini-batch B
{(x, y, 2, 2)} D.
16:
VaJsFT(()) t- V V B8
(B E, log (o ود
VoJnnc(0) + Vo B (0, Ej 1(3) # 2) log (0)))))(3311
17:
18: + 0 a VoJsFr( VoJRAc(f))
19: end for
20: end for
Output: 8.
4.1 Setups
4.1.1
Benchmark Tasks and Evaluation Metrics
Summarization task. We use OpenAl TL;RR dataset (Stiennon et al., 2020) for this task, which includes posts from Reddit forum and their corresponding summaries from human labelers. We evaluate performance by calculating the win rate against the reference summary and reporting Rouge F1 scores (Lin, 2004) on its test set.
General QA task. For this task, use the Ultrachat-2005 dataset, high-quality 200K subset of the Ultrachat corpus (Ding et al., 2023), which contains approximately 1.4 million general QA dialogues generated by ChatGPT (3.5) Turbo API. We evaluate performance using the length-controlled (LC) win rate (Dubois et al., 2024a) on the Al AlpacaFvvll 2.0 benchmark (Dubois et al., 2024b).
Math QA task. We also compare the language model's ability in mathematical calculation and reasoning. To do this, we use two commonly used math QA datasets: GSM8K (Cobbe et al., 2021) and MATH (Hendrycks al., 2021), and evaluate the accuracy on their respective test sets.
4.1.2 Baselines Methods
Considering the setting of this paper, which applics in cases where access is limited only to human demonstra- tion data, we compare our methods against existing approaches that also rely exclusively on demonstration data: SFT (Ouyang et al., 2022), NEFTune (Jain et al., 2023), and SPIN (Chen et al., 2024). These methods represent strong set of supervised approaches basod on demonstration data. NEFTune introduces noise input token embeddings to improve the model's robustness, and SPIN utilizes self-replay to ensure that the generated outputs remain indistinguishable from the reference demonstrations.
7

---

Published in Transactions on Machine Learning Research (01/2025)
4.1.3 Training
We use Pythia-1B as the pre-trained model for the summarization experiments and Mistral-7B-v0.1 for the general QA and math QA experiments. Additionally, we use Mistral--7B-sll-beta as the SFT model for the general QA task, it is fine-tuned version of Mistral-7B-v0.11, Therefore, we report its performance SFT model and use it to fine-tune SPIN and our methods. For the summarization and math QA tasks, following SPIN (Chen et al., 2024), we first perform SFT the pre-trained models and then train SPIN, BASH, and RAC on top of them. See Appendix A.2 for more details on the baselines and our methods.
4.2 Main results
4.2.1
Summarization Task
For this task, we train SFT and NEFTuue for two epochs, starting from the Pythia-IB base model, and then continue fine-tuning SFT model with SPIN, BASH, and RAC for one more epoch, as no noticcable improvements were observed beyond that point. We compare their performance based on the win rate of the generated outputs against the reference responses on the test set, evaluated by the GPT-4 Turbo model. The evaluation prompt template is provided in Appendix A.1. To ensure more comprehensive results, repeat the training with three different random seeds and report the win rates along with the Rouge F1 scores (Lin, 2004) in Table 1.
Table 1: Comparison of the performance (higher is better) of methods against others the summarization task. We train each method using three different seeds and report the average win rate across them in addition Rouge scores. The gray denotes the standard deviations.
win rate (%)
Rouge
Rouge 2
Rouge
SFT
20.33
31.69
11.02
24.53
NEFTune
27.21 +0.40
31.97
11.16
24.73
SPIN
21.07
30.46
10.09
22.55
BASH (ours)
28.12::.43
32.00
11.27
24.89
RAC (ours)
28.02-10.39
32.01
11.20
24.80
As Table shows, although our methods outperform the others, the improvement is not particularly significant. This could be attributed to the nature of the summarization task, where the prompts and queries are quite long, but the generated responses are very brief. Since our method operates directly the response space, the relatively short length of the responses means that there is less to correct during the generation steps. Despite the inherent limitations of the task, our method still improves the base models, demonstrating its applicability across different lasks.
4.2.2 General QA task
Following SPIN (Ouyang et al., 2022), we randomly sample 50K prompts from the full training set to generate oflline datasets D, and D, for our methods and then use these datasets to train our BASH and RAC as explained Sec. 3.
To provide a comprehensive view of the results, we use the length-controlled win rate on AlpacaEval 2.0 our evaluation metric, which compares the generated out against those from GPT-4, following the standard AlpacaEval prompt template. For each method, evaluate performance using two types of generation: 1) Sampling-based generation with a temperature of 0.7, aligning with the default settings of the Mistral-7B model and its derivative, Zephyr-7BB Beta (Tunstall et al., 2023), on AlpacaEval', Given the inherent randomness in single generation evaluations, we repeat the generation three times and report the average evaluation for each generation. 2) Greedy generation, which selects the next token by choosing the one with the highest probability.
See httpp://hugginggeee httpe///ew...tgggaeccc///aaaaaaa
8See httpa://github. http//ttaataaabaaaaaaaaa http://www....eaaaaaaaa///aaaaaa hwwp://apppaaaaaa/ooaeeaaaeaa eval/nodala, http://wwa...aaaeee/aa/aa//////a htt.:/ppppp...bbbbtott//rrrggggo
8

---

Published in Transactions on Machine Learning Research (01/2025)
Table Comparison of the performance (higher is better) of our methods against others in the General QA task using the length-controlled win rate the AlpacaEval 2.0 benchmark. These results shows that our method consistently slands out, mainlaining ils effectiveness across different generation stralegies.
Temperature-0.1 0.7 (%)
Greedy (%)
Generation
Generation Generation 3
Average
SFT
7.90
8.45
7.74
8.03
7.06
NEFTune
7.60
6.99
7.56
7.38.10.28
6.64
iter-1
8.76
9.27
9.41
9.15++022
8.66
SPIN
iter-2
8.44
8.89
8.17
8.50+0.30
7.64
iter-1
8.77
8.56
8.48
8.60+0.12
7.31
BASH (ours)
iter-2
8.73
9.38
9.10
9.07110.27
8.42
iter-1
9.95
9.15
9.49
9.53 +0.33
9.04
RAC (ours)
iter-2
10.68
10.54
9.89
10.37 10.34
9.41
Table Comparison the test accuracy (higher better) our methods against baseline algorithms GSM8K and MATH tasks, averaged across three different generations used for evaluation.
GSM8K
MATH
SFT
56. 76.0699009
13.11#0.12
NEF
55.85+00.38
12.13+0.03
SPIN
10.91
6.85
BASH (ours)
60.22+0.31
14.59=0.09
RAC (ours)
59.41
13.75
As shown in Table 2, our RAC achieves the highest win rate in both sampling-based and greedy generation settings. Additionally, BASH also demonstrates consistent improvements compared to other methods, except for SPIN. It important to note that our methods show nearly monotonic improvements and consistent behavior across different iterations, unlike methods such as SPIN 9, This consistency highlights the applicability of our approach beyond single iteration, suggesting that it can be effectively utilized with larger models and in more iterations, resulting in consistent improvements over time.
4.2.3 Math QA task
In this setting, train SFT and NEFTune for two epochs, starting from the Mistral-7B-vO.1 base model, and then continue fine-tuning SFT model with SPIN, BASH, and RAC for one more epoch. We prompt the model to generate answers in the GSM8K and MATH test sets with query template attached in Appendix A.1. We then calculate the strict match accuracy of the generated answers with the ground truth. We use sampling-basel generation with a temperature 0.1, repeating the generation three times, and report the average accuracy in Table 3. As the results show, BASH and RAC outperform other methods on both GSMSK and MATH benchmarks. These results further demonstrate the effectiveness of our methods across diverse set of tasks.
4.2.4 Alignment with preference data
In this experiment, we further demonstrate that our model can serve as more effective initialization than standard SFT throughout the aligument pipeline. Specifically, we show that using our method to initialize a
"e observe decrease in the LC win rate of SPIN in the sccond iteration, primarily due to siguificant increase in generation length while the conteut remains similar. Since the length-controlled win accounts for output length, this results in performance drop. Similarly, the win for greedy generation is lower than for random sampling temperature 0.7, greedy generation tends be more verbose.

---

Published in Transactions on Machine Learning Research (01/202))
0.7
0.6
0.9
Embedding distance
0.5
0.8
0.5
0.3
0.4
0.7
0.4
0.6
0.3
0.2
0.5
0.2
0.4
Traiting exomple 1: buald
Triniing example 2: assoriated devices mitiigaioo.
Test exomale 1: impact of sestern Euroge economies.
Test example story about time
Caraal global
Figure 2: Visualization of the embedding distance between generated and reference responses. The left two ligures are based queries from the training set of the UltraChat-2008 dataset, while the right two figures are from the Lest set. Correspouding queries each ligure are summarized at bottom, with full queries available in Appendix A.3. We generate 256 responses from models and compute their embedding distances to the reference responses. Each violin plot includes an inner box plot that displays the maximum, third quartile, median (indicated by white line), first quartile, and minimum distances, while the shape of the violin represents the estimated probability density of the embedding distance.
model, followed by fine-tuning with preference data via DPO (Rafailov et al., 2024), leads to better results than initializing with a standard SFT model. For this, we closely follow the settings of the zephyr-7b-beta model (Tunstall et al., 2023) and use the preprocessod UltraFeedback dataset (Cui et al., 2024) as the preference data. The results this experiment, presented Table 4, illustrate the importance of our approach in enabling downstream alignment methods achieve improved performance when initialized with our method instead of standard methods like SFT.
Table Comparison the performance (higher is better) of DPO initialized with our methods versus others the AlpacaEval 2.0 benchmark. These results clearly demonstrate the effectiveness of method and provide further evidence of its applicability throughout the alignment pipeline.
0emperature=0.7 (%)
Greedy (%)
Ceneration
Generation 2
Ceneration 3
Average
SFT + DPO
14.17
14.74
13.28
14.06:00.0
14.22
BASH + DPO
14.66
13.34
13.61
13.87+0.57
14.39
RAC + DPO
16.40
15.21
15.95
15.85.10.49
15.35
4.3 Qualitative Analysis
Visualization. To illustrate how our methods address the training, we evaluate the discrepancy between model-generated responses and reference responses on the UltraChat dataset. We begin by selecting two queries from both the training and test and then use models fino-tuned with different methods to generate 256 responses by sampling with temperature of 0.7. To measure the deviation of the generated responses from the reference responses, first utilize Sentence Transformer 10 (Reimers & Gurevych, 2019) to extract the embeddings for each response. We then compute the distance between the embeddings of each generated response and its corresponding reference response, which we refer to as sentence distance. The distribution of these distances visualized in Fig. 2. As shown these plots, BASH and RAC reduce the discrepancy between the generated responses and the reference ones, except for the question with very open answers such as "write story" in test example indicating that our approach effectively narrows the gap between training and inference time.
gap in autoregressive model
10 htwp:////wwuggiieecc..cooeeeeeee
base-v2.
10

---

Published in Transactions on Machine Learning Research (01/2025)
5 Conclusion
In this paper, we propose principled methods to bridge the gap between training and inference time in LLMs by leveraging self-generated tokens. Specifically, this gap arises from the training strategy where the model uses the ground truth token from the previous time step as the input for the model at the current time step. While this strategy effective during training, the model must rely on its own predictions as input for subsequent steps during inference. This reliance on its own predictions can lead error accumulation and degraded performance, particularly when the generated sequences deviate from the conditions encountered during training. Scheduled sampling has emerged exposed to its own generated tokens during training, thus more closely simulating the conditions of inference. Despite its effectiveness, especially with smaller models and particularly in recurrent models, this approach has not been successfully adopted in training LLMs due to the computational and practical complexities of incorporating into the training process. However, our proposed methods are specifically tailored for LLMs without requiring structural changes to the model. Specifically, our approaches, BASH and RAC, incorporate the model's own generations in an offline manner, allowing training that closely mirrors the inference process. Morcover, RAC also builds in self-correcton capability during training, which becomes critical when the model's own generalion tineraaiios deviale siguificantly from the ground truth. These methods can be easily integrated into the training of LLMs without requiring changes to the training process or model architectures, which was not the case with previous methods. Through a series of comprehensive experiments, we demonstrate that our method outperforms existing approaches in both summarization and question answering (QA) tasks. The results indicate that by aligning the training process with the conditions of inference, we can enhance the model's performance and reliability, ultimately leading to more accurate and contextually appropriate results.
alternative approach, where the model is gradually
11

---

Published in Transactions on Machine Learning Research (01/2025)
References
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. PMLR, 21-27 Jul 2024.
Dzmitry Bahdanau, Kyung Hyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. International Conference on Learning Representations, 28, 2015.
Samy Bengio, Oriol Vinyals, Navdoop Jaitly, and Noam Shazeer. Scheduled sampling for sequence prodiction with recurrent neural networks. Advances in nearal information processing systems, 28, 2015.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William Freeman. Maskgit: Masked generative image transformer. In The IEEE Conference on Computer Vision and Pottern Recognition (CVPR), June 2022.
Zixiang Chen, Yihe Deng, Huizhuo Yuan, Kaixuan Ji, and Quanquan Gu. Self-play fine-tuning converts weak language models to strong language models. arXiv preprint 3rX1032401.01335, 2024.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rmn encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078, 2014.
Karl Cobbe, Vineet. Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Malthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. Xiv preprint arxiv:2110. 14168, 2021.
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Bingxiang Hc, Wei Zhu, Yuan Ni, Guotong Xie, Ruobing Xie, Yankai Lin, et al. Ultrafeedback: Boosting language models with scaled feedback. In Forty- first International Conference on Machine Learning, 2024.
Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations. aiv preprint arXie:2305.14230 14233, 2023.
Yann Dubois, Balázs Galambosi, Percy Liang, and Tatsunori B Hashimoto. Length- alpacaeval: A simple way to debias automatic evaluators. arXiv preprint orXiv:2404.04 044.24444.447 2024a.
Yann Dubois, Chen Xuechen Li, Rohan Taori, Tianyi Zhang, Ishaan Gulrajani, Jimmy Ba, Carlos Guestrin, Percy S Liang, and Tatsunori Hashimoto. Alpacafarm: A simulation framework for methods that learn from human feedback. Advances in Neural Information Processing Systems, 36, 2024b.
Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis, 2020.
Rasool Fakoor, Abdel rahman Mohamed, Margaret Mitchell, Sing Bing Kang, and Pushmeet Kohli. Memory- augmented attention modelling for videos. arXiv: 161 1..11.011 1.02261, 2017.
Rasool Fakoor, Xiaodong He, Ivan Tashev, and Shuayb Zarar. Constrained convolutionalrecurrett networks to improve speech quality with low impact on recognition accuracy. In 2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3011-3015, 2018
Rasool Fakoor, Pratik Chaudhari, Jonas Mueller, and Alexander J. Smola. Trade; Transformers for density cstimation. arXiv preprint arXim:2004.022 02244, 2020.
Rasool Fakoor, Jonas Mueller, Zachary C. Lipton, Pratik Chaudhari, and Alexander Smola. Timo-varying propensity score bridge the gap between the past and present. In ICLR, 2024.
Gemini-Teamn Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024.
12

---

Published in Transactions on Machine Learning Research (01/2025)
Kartik Goyal, Chris Dyer, and Taylor Berg- Kirkpatrick. Differentiable scheduled sampling for credit assignment. In Proceedings of the 55th Annual Meeting of the Associalion, for Compulational Linguisties (Volume 2: Short Papers), pp. 366 371, 2017.
Karol Gregor, Ivo Danihelka, Andriy Mnih, Charles Blundell, and Daan Wierstra. Deep autoregressive networks. Proceedings of the 31st International Conference on International Conference on Machine Learning Volume 32, 1CML'14, pp. 11-1242--1 II 1250, 2014.
Tianxing He, Jingzhao Zhang, Zhiming Zhou, and James Glass. Quantifying exposure bias for neural language generation. 2019.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.088 03874, 2021.
Ferene Huszár. How (not) to train your generative model: Scheduled sampling, likelihood, adversary?, 2015.
Neel Jain, Ping-yeh Chiang, Yuxin Wen, John Kirchenbauer, Mong-Min Chu, Gowthami Somepalli, Brian R Bartoldson, Bhavya Kailkhura, Avi Schwarzschild, Aniruddha Saha, et al. Neftunc: Noisy embeddings improve instruction finetuning, arXiv preprinl Xic:22230 05914, 2023.
Alex M Lamb, Anirudh Goyal ALIAS PARTH GOYAL, Ying Zhang, Saizheng Zhang, Aaron C Courville, and Yoshua Bengio. Professor forcing: A new algorithm for training recurrent networks. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, 2016.
Jianqiao Li, Chunyuan Li, Guoyin Wang, Hao Fu, Yuhchen Lin, Liqun Chen, Yizh Zhang, Chenyang Tao, Ruiyi Zhang, Wenlin Wang, Dinghan Shen, Qian Yang, and Lawrence Carin. Improving text generation with student-forcing optimal transport. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 9144 9156, November 2020.
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization, 2024.
Chin-Yer Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74-81, 2004.
Zuxin Liu, Josse Zhang, Kavosh Asadi, Yao Liu, Ding Zhao, Shoham Sabach, and Rasool Fakoor. Tail: Task-specific adapters for imitation learning with large pretrained models. Tn International Conference on Learning Representations (ICLR 2024), 2024.
Tsvetomila Mihaylova and André F. T. Martins. Scheduled sampling for transformers. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: Student Research Workshop, pp- 351-356, July 2019.
OpenAL. Gpt-4 technical report. arXiv preprint arXin:2303.0877, 2023.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer, 2018.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published in Transactions on Machine Learning Research (01/2025)
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1 67, 2020.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Nebworking, Slorage and Analysis, pp. 1- 16. IEEE, 2020.
Marc'Aurelio
Mare' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojeiech Zaremba. Sequence level training with recurrent neural networks. arXiv preprint 1rXic:1111.06732 2015.
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2019.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourleenth international conference on artificial intelligence and statistics, pp. 627-635. JMLR. Workshop and Conference Proceedings, 2011.
Florian Schmidt. Generalization in generation: A closer look at exposure bias. arXiv preprint or reXiv:1910.0920 2019.
Ili Shumailov, Zakhar Shumaylov, Yiren Zhao, Nicolas Papernot, Ross Anderson, and Yarin Gal. Ai models collapse when trained on recursively generated data. Nature, 631(8022))755-755) 2024.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-3021, 2020.
Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro von Werra, Clémentine Fourrier, Nathan Habib, et al. Zephyr: Direct distillation of Im alignment. arxiv preprint arXiv:2310. 16944, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Oriol Vinyals, Alexander Toshev, Samy Bengio, and Dumitru Erhan. Show and tell: A neural image caption generator, 2015.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Foi Xia, Ed Chi, Quoc Lc, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural informalion processing systems, 35:24824-24883 2022.
Ronald J. Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural Computation, (2):270-228) 1989.
Daniel M. Ziegler, Nisan Stiemnon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geolfrey Irving. Fine-tuning language models from human preferences, 2020.
14

---

Published in Transactions on Machine Learning Research (01/2025)
A Experiment Details
A.1 Prompt Templates
RAC prompt template f(x, in summarization task. In experiment of summarization task, is the original post and y is the reference summary from the TL;DR dataset. we will use the below template to generate RAC label.
POST]: REFERENCE SUMMARY]: y
Re-write summary of the post and cover the main content in the reference summary.
TL;DR:
RAC prompt template fzx,y) in general QA and math QA task. experiment of QA tasks, æ is query and y is reference answer from the datasets. We will use the below template to generate RAC label. (/s)is the special token indicating the end of sentence.
(Jsysteml)
You helpful assistant to answer user's question. Your will be given both question and reference response. You need to give a new response to the question and contain the main content in the reference response.(/s)
(Juser))
Answer this question based the following reference response: y/s) (Jassistant|)
Evaluation prompt template in summarization task. We use the below template to evaluate the win rate of generated summary against the reference summary with CoT. We also randomly shuffle the order between generated and reference summary to reduce the evaluation bias. In experiment, we will replace {post) by the post from test set and replace {summary A}, (summary B} by the shuffled generated and reference summaries.
Which of the following summaries does better job of summarizing the most important points in the given forum post, without including unimportant or irrelevant details? Judge based on accuracy, coverage, and coherence.
[post]
{post}
(summary A]
(summary A}
(summary B]
(summary B}
Instructions:
FIRST provide a one-sentence comparison of the two summaries, explaining which you prefer and why. line, state only A B to indicate your choice. Your respouse should the format: one-sentence comparison and explanation
SECOND,
Comparison:
Preferred: A B
Evaluation prompt template in math QA task. We use the template below to Lest the accuracy of the generated answer. Unlike other benchmarks, we use zero-shot (i.e., no QA examples are provided before query) and CoT (Wei et al., 2022) to prompt the language model to generate answers. In the experiment, we will replace {query) by the questions from the test set.
(Jsysteml) (/s) (Juser])
15

---

Published in Transactions on Machine Learning Research (01/2025)
{query) Let's think step by step.(/s) (Jassistant))
A.2 Hyperparameters and More Experiment Settings
More implementation details.
We implement baselines and our methods based on two codebases: sunmarize-fron-ooodbbabc (for sum- marization task) and Aligument-HHandbook (for general QA and math QA tasks), which use DeepSpeed ZeRO (Rajbhandari et al., 2020) for higher training efficiency and less computation overhead. In summariza- tion task, we generate the response from the fine-tuned model with temperature=0.01 for win rate evaluation following the codebase used.
We train SFT models for summarization and math QA tasks from the corresponding base pretrained models. In particular, we do not pack data during SFT training, which is introduced in the T5 model (Raffel et al., 2020) and is a default option in aligument-handboo in experiment. For SPIN, we use the official implementation and follow all training settings to generate data and fine-tune based on the SFT model in general QA and math QA tasks. The only difference is that adopt Mistral--7B-sf--eeta compatible with the summarization code repository. We also search the hyperparameter 3 of SPIN among (0.1, 1.0, 5.0] and report the highest win rate. Note that the enumeration starts from 0 in the SPIN paper, while ours starts from 1. Therefore, iteration 0 in the SPIN paper corresponds to itcration 1 in our paper.
repository. For NEFTune, we adopt the noise scale a
the base model. In summarization task, we implement SPIN ourselves to be
Hyperparameters.
We attach the hyperparamoters usod in the experiment in table
Table 5: Hyper-parameters used for experiments.
Summarization
General QA
Math QA
base pretrained model precision
pythia-11 1B
Mistral-7B-v0.1
Mistral-7BB-v.11
bfloat 16
bfloat 16
16
optimizer
AdamW
AdamW
AdamW
lcarning rate learning rate warmup steps learning scheduler
3 x 10-6
5 x 10-6
5 x 10-6
no warmup
10%
10%
cosine
cosine
cosine
global batch size SFT training epoch training iteration (BASH&RAC) training epochs in each iteration (BASH&RAC)
512
512
512
2
/
1
2
1
[1, 2]
mixture coefficient 3 in BASH generation
0.2
0.2
0.2
A.3 More Experiment Results
Queries used in Fig. for embedding distance computation.
Query of training example 1:
What factors influenced the decision to build the Panama Canal, and how did it transform global trade and transportation?
We implement algorithms based OpenAl summariae-ffom--cedbbbe and its clean-up version. 12See link of
13gee link of SPIN.
16

---

Published in Transactions on Machine Learning Research (01/2025)
Query of training example 2:
What the risks associated with IoT devices, and how they be mitigated?
Query of test example
How did the fall of the Soviet Union impact the economies of Eastern Europe?
Query of test example
Using vivid imagery and descriptive language, write a compelling short story about a time traveler who finds themselves transported the year 2077. Explore the world of the future, including advanced technol- ogy, cultural and societal changes, and environmental shifts. Consider adding a twist or unexpected turn to the plot to keep your andience engaged. Your story should be between 500-1000 words in length and should captivate the reader from beginning to end.
The performance comparison with different dataset sizes.
In this experiment, we compare our methods with baselines on training datasets with different sizes the summarization task. Specifically, we randomly choose subset of training data with proportion 10%, 20%, 50%, 100%. Then we keep all other scttings the same and train the model for the same steps as the full dataset training. For example, when training on the 20% data, will train for 10 epochs for SFT and 5 epochs for BASH or RAC (note that for the full data setting, train 2 epochs for SFT and 1 epoch for BASH or RAC). We still report the win rate against the reference summary on the whole test set using the same evaluation prompt.
29
28
win rate (%)
25
RAC
10%
20%
50%
100%
dataset size
Figure 3: The win rates of summarization task with different dataset sizes. Each results averaged three seeds. In each seed, the subset of trainig data different and we train model on the different subset for SFT, BASH and RAC. The win rate is evaluated the whole test set.
The final results are reported in Fig. We observe no notable performance drop when the datasot size exceeds 20%. When the dataset size becomes smaller, teacher-forcing-stye SFT training suffers from the data insufficiency, leading to more severe distribution shift between (Fakoor al., 2024) training and inference. In this case, our methods mitigate this issue and exhibit relatively smaller win rate decrease compared to SFT.
The performance of SFT training for more steps.
We train the SFT fine-tuned model by SFT for more steps as an ablation. We also compare its LC win rate on AlpacaEval 2.0 with our methods in Fig. 4. More SFT training steps improve the performance of model general QA task. However, there is performance drop at the third epoch. On the contrary, our
17

---

Published in Transactions on Machine Learning Research (01/2025)
methods ensure monotonic increase in LC win rate and can consistently outperform the SFT baseline in each iteration.
BASH
AlpacaEval 2.0 LC win rate (%)
10.0
9.5
8.0
model
epoch
epech
epoch
training epoch
Figure 4: The AlpacaEval 2.0 LC win rates comparison of our methods and SFT. Each results are averaged three generations. We leverage SFT the UltraChat dataset continue to train the existing SFT model and compare its performance with our methods. In the figure, the epoch 1 corresponds to the first iteration and epoch 2&:3 correspond lo second iteration for BASH and RAC. In the begimning of each iteration, will offline generate BASH sequences or RAC labels by current model.
The performance of training without combining the SFT loss.
Table 6: LC win rate comparison w.o. combining the SFT loss. We evaluate the average three sampled generation with 0.m.erature=0.7. The gray denotes the standard deviation of three evaluations.
with SFT loss
SFT loss
SFT
8.03
BASH
1.046
9.07
RAC
3.24+0.48
10.37
As shown in Table 6, there will be a large performance drop if do include the SFT loss in training. One potential reason is that the learning objective of BASH or RAC is biased estimator (Lamb et al., 2016) of the expert language model, which can be viewed as the underlying model of SFT dataset D. During training, the language model is enforced to fit the labels conditioned on student input distribution significantly shifts from the one in the SFT objective in Eq. (1). Therefore, we include SFT loss, an unbiased behavioral cloning objective, to train the LM to imitate the underlying expert model of the dataset.
The performance comparison between Scheduled sampling (SCS) and Batch-scheduled sampling (BASH).
Table Comparison of win rate on the summarization task. We compare the performance
win rate (%)
Rouge 1
Rouge
Rouge L
SFT
27.03+1038
31.69
11.02
24.53
SCS
28. 28.24410.46
32.03
11.22
24.95
BASH
28. 08
32.00
11.27
24.89
Table shows that scheduled sampling (SCS) achieves a win rate similar to BASHas expected. However, BASHis much more computationally efficient than SCS, which is crucial in the practical implementation of LLM training. To provide further perspective on the effectiveness of our approach, we conduct experiments comparing the computational overhead of SCS and BASH using 1B and 7B models.
18

---

Published in Transactions on Machine Learning Research (01/2025)
Table 8: Computation overhead comparison between SCS and BASHwith pythia-1B model on the summarization task. We compare the performance on 8xA6000 (18G) machine. For fair comparison, we set the local batch size as 16 and gradient accumulation slep as for all experiments. In this case, the global batch size is 16 4 4 512. The computation time of SCS includes both (online) dala generation and model training.
SCS
BASH (offline data generation)
BASH (training)
Computation time for global batch Memory usage each GPU
20.8s
0099
~4.4s
35GB
~ 17GB
28GB
Table 9: Computation overhead comparison between SCS and BASHwith Mistral-7B model on the general QA task. We compare the performance on an 8xA6000 (48G) machine. For fair comparison, we set the local batch size 8 and gradient accumulation step as 8 for all experiments. In this case, the global batch size is 8 512. The computation time of SCS includes both (online) data generation and model training.
SCS
BASH (oflliee data generation)
BASH (training)
Computation time for 1 global batch >3000s Memory usage of each GPU
3770s
755
44GB
22GB
36GB
As these results show, BASH requires significantly less GPU memory and computation time than SCS, making it more practical in LLM training.
A.4 More analysis of RAC
In this section, we first demonstrate how the RAC correction mechanism operates, followed by an example illustrating its failure case.
In Fig. 5, we present an example from GSM8K training set to illustrate how RAC leverages in-context learning corrects the error model-genee sener-generated response by appending the reference response in the query.
Reference Answer
his mother's ago, Peter 1/2"60 years
In four years, Peter will be 30-4 36 years old. Since age four years Harriet's age, Harrist will 34/2 years old in four years. Currently, Harriet The answer is: 13
RAC Correction
Question
Generated response
Peter
0
years
years, Harriot. If Peter's age currently holt of mother's s age, who is 60, how old is Harrtet now?
old.
RAC label
(the same as 2]
Teice
(same)
(samo)
..
c SFL-mode--emerrted Answer
R-C-unddd-gonerrttee Answer
Lat us first find Peter's age. Half of Patur's mother's ago 60/2 = 30, Sinee Peter's is currently half of his mother's age, Peter 30.
years. Peter 60/2 In four years, Harriet will be 30/2 - 15 years old. Currently, Harrict The answer isc 11
years This means that four years, Peter will be twice old as Harriet, Le. Harriet will be 34/2 17 old four Currently, Harriet is 17-4 13 years old. is:
four years, Peter will 30-4
years old.
Figure 5: How does RAC correct mistakes in model-generated responses? In this example, SFT model incorrectly calculates Peter's age in the initial slep, deriving it as 60/2 = 30. However, the model fine-tuned with RAC produces the correct result. Specifically, the SFT model's incorrect calculation leads to an age of 30, while the ground truth is 34. To address this error, RAC labels the next token after "Peter will be 60/2' as "4" instead of "=', guiding the model towards the correct computation. After training with RAC, the model successfully calculates Peter's age accurately, resulting in the correct answer.
We also present failure case of RAC correction. Despite being augmented by the reference answer, RAC is primarily trained on self self-generated reasoning trajectory and may struggle to identify complex reasoning errors in the model-generated response. In Fig. 6, the RAC model attempts to correct the final answer
19

---

Published in Transactions on Machine Learning Research (01/2025)
directly; however, the token-wise correction is inconsistent with the context of the preceding tokens and fails to address the underlying errors in the generated response.
Referenee Answer
slngll layer per
requires
RAC Correction
Question
As each wall has thickness bricks, he will 100
Generated response
200 bricks
Rence
the complete fence contalnns four walls, he will require a total of 200-4=800 bricks.
Gooper building brick fence around
RAC labol
(the same as :)
(samn)
aides property.
is:
long. bricks high, and bricks deep. Hour many bricks does Cooper need to complete his fence?
G SFT-modde-gemeratee Answer
RAC-modee-eemeraeedd Answer
Each wall
with
First, determine number beicks per
20"5-100 bricks per layer.
long.
Each wall is 5 beicks high, fence with four walls 514 20 bricks high.
Next, we need to datermine the number layers per wall. To do this, we perform 100/2=50 layers per wall.
walls 2* bricks deep.
bricks deep,
we perform 50"2=100
per
per
entire lence 80"20"g 3200 bricks.
Finally, need determine the total number bricks for the entline fence. do this, perform 100*4=100 bricks entire fence. The answer is: 400
The answer is: 3200
Figure 6: An example of RAC failure. Given the reference answer, RAC model labels the next token of 80 20 us 8 to match the ground-truth truth answer 800. However, such correction is not consistent with the previous context. Meanwhile, RAC correction fails find the deeper reasoning error that the entire fence is not long, 20 high or deep and does not correct it when the generated answer multiplies them together.
20


Reviews:
[{"review_id": "df0vrnpobr", "reviewer": "Reviewer_upQw", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nThe experiment results suggest the proposed methods perform better than a recent  baseline. The proposed batch-scheduled method should be faster than existing methods of the same purpose. Writing is generally clear.\n\nWeaknesses:\n\n1. While I see the motivation behind RAC, I found the token-by-token approach not super convincing. A factual error likely appears very early in the generation - for example, a response to \"Did Amy go to school on Monday?\" naturally starts with \"Yes/No\". If the model is forced to learn to complete something that starts with \"No\" when it *knows* that the reference answer is \"Yes\" (from the prompt), it looks like also a big distribution shift to me, arguably more than that caused by completing ground-truth vs self-generated prefixes. \n\ta. Also, while it empirically might have improved the performance, generating withe prompt during training time vs generating without the prompt seem like a bigger discrepancy to me.\n\n2. If I understand correctly, BASH is essentially SCS, but batched. The novelty thus looks low. Also, to show that BASH is actually useful, some experiments on the overhead vs batch-size (i.e. how often do we stop and regenerate the dataset) as well as the corresponding impact on performance is due. \n\n3. SCS was from pre GPT time when LMs were much less performing. Since modern LLMs are pretty strong, the mismatch between training and inference is less obvious, and the motivation of this paper might be weakened. Some analysis/experiment on how the improvement changes as models become bigger should be included.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed two methods (BASH and RAC) to fine-tune language models (LMs). BASH involves replacing ground-truth tokens with tokens generated by the LM itself, in order to mitigate the discrepancy between training and inference time. To align the self-generated target, RAC was also proposed where the self-generated target is generated with a prompt that includes the reference answer.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Overhead & performance vs batch size (Weakness 2)\n2. Experiments mentioned in weakness 3.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "H6GSF2jRBo", "reviewer": "Reviewer_PCM5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The paper proposes two simple and effective ideas to scale scheduled sampling. \n- The paper is generally well-written and easy to understand.\n- The experimentation is extensive -- showing the effectiveness of the approach across several different tasks and models. \n- There is some initial analysis into what makes these approaches better. \n\nWeaknesses \n\n- The analysis of why BASH and RAC work seems a bit preliminary and the paper might benefit from a deeper analysis. \n- The algorithms are variants of scheduled sampling and the paper does not compare BASH and RAC to older works on scheduled sampling (even on very small models as older scheduled sampling hard to scale)\n\nQuestions\n- Intuitively it seems that BASH should underperform RAC on the mathematical reasoning datasets as BASH allows the model to be trained on potentially incorrect reasoning chains. But the inverse is observed in the results. Why do you think this is the case? \n\nFormatting errors \n- Footnote on page 8 overflows to page 9.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes an efficient method to perform scheduled sampling at scale. Prior work has investigated that \"teacher-forcing\" can result in the LM learning an incorrect distribution of tokens. Scheduled sampling and its variants were proposed as solutions. \nThe authors claim that scheduled sampling is hard to scale, leading to it not being used while training large LMs. \nThe paper proposes BASH (Batched Scheduled Sampling), which creates a mixed dataset between teacher-forced and model-generated sequences. The authors also perform BASH between training instances which allows the approach to scale. \n\nThe authors note that BASH can result in creating \"incorrect\" sequences and to account for this, they propose an alternative method called RAC that self-corrects these model-generated sequences post-hoc. \nThe authors empirically show that BASH and RAC improve model performance on several downstream tasks like summarization, math and general QA, and instruction following. \nThe authors compare their approach against several baselines notably NEFTtune, SPIN, and vanilla SFT. They find that BASH and RAC outperform. \nThe authors also analyze the gap between training and inference using the distance of the embeddings of the training and inference sequences and they find that models trained using BASH and RAC exhibit lower distance indicating lesser discrepancy between the two.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Overall I feel positive about the paper. I believe that the paper would benefit strongly from a more detailed analysis of what aspects of BASH and RAC improve performance. However, I don't think that this is critical for acceptance.", "broader_impact_concerns": "No Broader Impact Statement is present. I don't see any ethical concerns surrounding this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JFZARjO1se", "reviewer": "Reviewer_bzhr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The methods are great. It is easy to understand.\n* The performance improvement is significant. \n* The author validates the method on different datasets, proving its effectiveness.\n\nWeakness:\n* Does should $z^j=p_{\\theta}(.|x,g^{j})$ be rewritten as  $z^j=argmax~p_{\\theta}(.|x,g^{<j})$?\n* For Equation 4, could you give an example to explain how you replace $y^j$ with $z^j$?\n* In Table 3, BASH performs better than RAC, could you explain the reason?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes Batch-Scheduled Sampling and Reference-Answer-based Correction to reduce the gap between training and inference. The authors conduct sufficient experiments on different datasets to support their claims. Also, the authors have conducted the analysis to visualize the embedding distance between generated and reference responses.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please check the weakness.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is motivated by the fact that there is a difference between the way models are trained and what they do at inference time; this is a well-known observation. During training, the standard approach is to use the reference/ground-truth tokens for the loss. Naturally these are not available at inference time, so only the generated tokens are used to predict the next token---unlike in training. This gap can lead to potential performance costs. The authors introduce two methods to address the issue. One is a more efficient batch version of the scheduled sampling method. The other is a method to integrate a method for self-correction---a fairly creative approach. 

Most reviewers found that there were a number of interesting insights in this paper. The experimental results are particularly solid. One of the reviewers asked for a more detailed analysis of the self-correction approach (or at least more of an explanation for what should happen in certain intuitive cases where it's not clear what the method will do, as spotted by the reviewer). The argument raised is that while the authors' original motivation is specifically to remove the gap between training and inference, in certain cases this need not happen, i.e., the authors' method is producing improved results without strictly removing this gap. While I appreciate the reviewers' thoughtful points, even if this is the case, the fact that the self-correction approach appears to actually deliver strong results makes it worthwhile. In fact this would be an interesting finding on its own.

Ultimately there is enough here that is valuable that I believe the paper is worth accepting. For the camera ready version, the authors are encouraged to provide a bit more commentary on the scenarios brought up above.
Decision: Accept (poster)