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

Abstract

Autoregressive language models are usually trained with next-token prediction under a fixed left-to-right target order. This paper studies target-token organization as a training-time design choice across arithmetic reasoning, multi-label text classification, and text generation. Controlled order studies show that the default NTP order can be far from the best-performing tested organization, while fixed hand-crafted orders such as reverse order do not generalize across tasks. To address this, Max(MI(S;t)) is introduced as a deterministic mutual-information-based heuristic that prioritizes information-rich target tokens. For fixed-length outputs, it constructs task-specific prediction orders; for open-vocabulary text generation, the same principle is applied through lightweight prefix augmentation. Experiments with GPT-style models and pretrained Llama/Qwen variants show that Max(MI(S;t)) often improves over standard NTP and simple baselines such as reverse order, TF-IDF, and conditional-probability selection, without changing model architecture or autoregressive decoding.