Training LLMs Beyond Next Token Prediction -- Filling the Mutual Information Gap

TL;DR

This paper proposes a new training approach for large language models that focuses on predicting information-rich tokens, aiming to improve performance and efficiency over traditional next-token prediction methods.

Contribution

It introduces a principled method for selecting target tokens during training, enhancing model performance and theoretical understanding beyond conventional next-token prediction.

Findings

Improved performance on arithmetic tasks

Enhanced multi-label classification accuracy

Better natural-language generation quality

Abstract

Optimizing training performance in large language models (LLMs) remains an essential challenge, particularly in improving model performance while maintaining computational costs. This work challenges the conventional approach of training LLMs using next-token prediction (NTP), arguing that by predicting information-rich tokens during training, there is a more effective way to train LLMs. We investigate the impact of the proposed solution in three kinds of tasks for LLMs: arithmetic, multi-label classification of text, and natural-language generation. This work offers a principled approach to optimizing LLM training, advancing both model performance and theoretical understanding of the target-token selection strategies.