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

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Cen, Zhepeng, Liu, Yao, Zeng, Siliang, Chaudhari, Pratik, Rangwala, Huzefa, Karypis, George, Fakoor, Rasool
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866913658157662208
author Cen, Zhepeng
Liu, Yao
Zeng, Siliang
Chaudhari, Pratik
Rangwala, Huzefa
Karypis, George
Fakoor, Rasool
author_facet Cen, Zhepeng
Liu, Yao
Zeng, Siliang
Chaudhari, Pratik
Rangwala, Huzefa
Karypis, George
Fakoor, Rasool
contents 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 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_14655
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Bridging the Training-Inference Gap in LLMs by Leveraging Self-Generated Tokens
Cen, Zhepeng
Liu, Yao
Zeng, Siliang
Chaudhari, Pratik
Rangwala, Huzefa
Karypis, George
Fakoor, Rasool
Machine Learning
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 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.
title Bridging the Training-Inference Gap in LLMs by Leveraging Self-Generated Tokens
topic Machine Learning
url https://arxiv.org/abs/2410.14655