Bridging the Training-Inference Gap in LLMs by Leveraging Self-Generated Tokens

Zhepeng Cen, Yao Liu, Siliang Zeng, Pratik Chaudhari, Huzefa Rangwala, George Karypis, Rasool Fakoor

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.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

We submit the (deanonymized) camera-ready version of the manuscript. Here are the minor changes:
- In appendix A.3 (page 18~19), we add new comparison results between SCS (online scheduled sampling) and BASH (offline scheduled sampling) in terms of final performances and computation overhead.
- In appendix A.4, we add more descriptions on the failure case of reference-answer augmented correction (RAC) and why RAC fails on those cases. We also mention it in main text.
- Correct other typos.

Assigned Action Editor: ~Frederic_Sala1

Submission Number: 3495