Efficient Training of Language Models with Compact and Consistent Next Token Distributions

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Sathe, Ashutosh, Sarawagi, Sunita
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866929408178126848
author Sathe, Ashutosh
Sarawagi, Sunita
author_facet Sathe, Ashutosh
Sarawagi, Sunita
contents Maximizing the likelihood of the next token is an established, statistically sound objective for pre-training language models. In this paper we show that we can train better models faster by pre-aggregating the corpus with a collapsed $n$-gram distribution. Previous studies have proposed corpus-level $n$-gram statistics as a regularizer; however, the construction and querying of such $n$-grams, if done naively, prove to be costly and significantly impede training speed, thereby limiting their application in modern large language model pre-training. We introduce an alternative compact representation of the next token distribution that, in expectation, aligns with the complete $n$-gram distribution while markedly reducing variance across mini-batches compared to the standard next-token loss. Empirically, we demonstrate that both the $n$-gram regularized model and our approximation yield substantial improvements in model quality and convergence rate compared to existing methods. Furthermore, our approximation facilitates scalability of gains to larger datasets and models compared to the straightforward $n$-gram regularization method.
format Preprint
id arxiv_https___arxiv_org_abs_2407_02819
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Efficient Training of Language Models with Compact and Consistent Next Token Distributions
Sathe, Ashutosh
Sarawagi, Sunita
Computation and Language
Machine Learning
Maximizing the likelihood of the next token is an established, statistically sound objective for pre-training language models. In this paper we show that we can train better models faster by pre-aggregating the corpus with a collapsed $n$-gram distribution. Previous studies have proposed corpus-level $n$-gram statistics as a regularizer; however, the construction and querying of such $n$-grams, if done naively, prove to be costly and significantly impede training speed, thereby limiting their application in modern large language model pre-training. We introduce an alternative compact representation of the next token distribution that, in expectation, aligns with the complete $n$-gram distribution while markedly reducing variance across mini-batches compared to the standard next-token loss. Empirically, we demonstrate that both the $n$-gram regularized model and our approximation yield substantial improvements in model quality and convergence rate compared to existing methods. Furthermore, our approximation facilitates scalability of gains to larger datasets and models compared to the straightforward $n$-gram regularization method.
title Efficient Training of Language Models with Compact and Consistent Next Token Distributions
topic Computation and Language
Machine Learning
url https://arxiv.org/abs/2407.02819