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

Chun-Hao Yang; Bo-Han Feng; Tzu-Yuan Lai; Yan-Yu Chen; Yin-Kai Dean Huang; Shou-De Lin

cs.CLOct 31, 2025

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

Authors: Chun-Hao Yang, Bo-Han Feng, Tzu-Yuan Lai, Yan-Yu Chen, Yin-Kai Dean Huang, Shou-De Lin

Organizations: National Taiwan University

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.

Figures & tables

Figure 1: An example of learning information-rich tokens first in MLC task. Prioritizing information-rich tokens ( insult label) during training reduces prediction uncertainty and improves subsequent label prediction, whereas left-to-right NTP can be harder to learn due to a lack of emphasis on informative tokens and cumulative uncertainty.

Figure 2: Rank-based comparison of target-token organization strategies. For (a)–(c), we exhaustively evaluate all

4!=24

orders for four output digits or labels; NTP ranks poorly among these orders, while Max(MI(

S;t

)) consistently approaches or matches the optimum, and reverse order only works well for addition. For (d), where reverse ordering does not apply to natural language generation, we instead expose one selected word earlier by prefix duplication; Max(MI(

S;t

)) achieves the lowest perplexity, outperforming NTP, TF-IDF, and conditional-probability baselines. For experimental setup, see Appendix B. For separate illustrative training-curve experiments showing performance discrepancy across training orders, see Figure 9 .

Figure 3: Overview of Max(MI(

S;t

)) . We first compute MI scores between the current source and candidate targets. Fixed-length tasks use the selected targets to form a reordered output sequence, while open-vocabulary TG uses the selected word for per-sentence prefix augmentation.

Task	Best Order	Best Acc.	NTP	Reverse	Max(MI( $S$ ; $t$ ))
3-digit Addition	$C_{4}C_{3}C_{2}C_{1}$	99.99%	85.09%	99.99%	97.46%
2-digit Multiplication	$C_{3}C_{2}C_{4}C_{1}$	98.22%	70.15%	95.35%	97.38%
4-digit Logarithm	$C_{1}C_{4}C_{3}C_{2}$	98.21%	95.99%	96.63% †	96.63% †
3-digit Greatest Common Divisor	$C_{1}C_{3}C_{2}$	95.80%	91.45%	92.59% †	92.59% †
3-digit Multiplication	$C_{5}C_{4}C_{3}C_{6}C_{2}C_{1}$	99.66%	71.43%	88.15%	96.23%
2-digit Chicken-Rabbit Problem	$C_{4}C_{2}C_{3}C_{1}$	92.12%	34.92%	45.04%	89.48%

Table 1: GPT-based model performance across different arithmetic tasks. Token prediction strategies: (1) NTP : standard left-to-right target order. (2) Reverse : reversed NTP. (3) Max(MI(

S

;

t

)) : prioritizes output digits with the highest mutual information with the source. † : Identical predicted token sequences for both strategies.

Figure 4: Arithmetic data format with reordered digits. The source expression is fixed, while output digits are rearranged according to the selected prediction order.

Task	NTP	Reverse	Max(MI( $S$ ; $t$ ))
3-digit Addition	99.3%	99.8%	100.0%
2-digit Multiplication	94.1%	95.4%	96.7%
4-digit Logarithm	89.1%	89.0% †	89.6% †
3-digit Greatest Common Divisor	94.0%	93.7% †	94.3% †
2-digit Chicken-Rabbit Problem	84.7%	96.9%	99.4%
Average	92.2%	95.0%	96.0%

Table 2: Qwen-Math performance across arithmetic tasks. NTP : standard left-to-right order; Reverse : reversed NTP order; Max(MI(

S

;

t

)) : prioritizes information-rich output digits. † : Identical predicted token sequences for Reverse and Max(MI(

S

;

t

)).

Figure 5: Long-sequence multiplication with Qwen-Math. Each task multiplies two

d

-digit numbers. All strategies degrade as output length grows, but Max(MI(

S

;

t

)) shows better robustness before the longest settings collapse.

Model	Dataset	Best Order	Best Acc.	NTP	Reverse	Max(MI( $S$ ; $t$ ))
Llama-3.1-8B-Instruct	ToxicComment	$C_{3}C_{1}C_{2}C_{4}$	85.81%	84.22%	82.59%	85.14%
Llama-3.2-1B-Instruct	ToxicComment	$C_{1}C_{4}C_{3}C_{2}$	77.20%	74.65%	74.48%	75.04%
Qwen2.5-3B-Instruct	ToxicComment	$C_{1}C_{3}C_{4}C_{2}$	78.85%	76.24%	70.98%	78.76%
Qwen2.5-1.5B-Instruct	ToxicComment	$C_{2}C_{1}C_{3}C_{4}$	74.26%	67.87%	68.90%	71.96%
GPT-2	ToxicComment	$C_{3}C_{1}C_{2}C_{4}$	88.63%	88.02%	88.17%	88.15%
GPT-2	PaperAbstract	$C_{4}C_{1}C_{2}C_{3}$	88.73%	88.52% †	88.23%	88.52% †

Table 3: Maximal accuracy in MLC across models and English datasets. NTP : standard left-to-right label order; Reverse : reversed NTP order; Max(MI(

S

;

t

)) : prioritizes information-rich label tokens. † indicates identical predicted label sequences.

Figure 6: MLC data format with reordered labels. The input text is fixed, while output label decisions are rearranged according to the selected prediction order.

Model	NTP	Reverse	Max(MI( $S$ ; $t$ ))
Llama-3.1-8B †	81.52%	80.16%	82.82%
Llama-3.2-1B †	72.84%	72.39%	74.61%
Qwen2.5-3B †	72.73%	71.92%	74.39%
Qwen2.5-1.5B †	68.08%	70.39%	70.18%
GPT-2	75.65%	69.42%	79.70%
Average	74.16%	72.86%	76.34%

Table 4: Maximal accuracy in MLC on the Chinese ToxicComment dataset. NTP denotes the standard left-to-right label order; Reverse denotes the reversed NTP order; Max(MI(

S

;

t

)) prioritizes information-rich label tokens. † indicates instruction-tuned models.

Figure 7: TG data format with per-sentence prefix augmentation. The selected word is inserted before each sentence between [START] and [END] .

Model	NTP	TF-IDF	P( $t\|S$ )	Max(MI( $S$ ; $t$ ))
GPT-2	18.98	16.28	15.60	14.59
Llama-3.2-1B	8.19	7.51	7.27	6.95
Average	13.58	11.90	11.44	10.77

Table 5: Perplexity on WikiText-2. Lower is better. NTP denotes the standard unaugmented baseline; TF-IDF, P(

t|S

), and Max(MI(

S

;

t

)) use the same prefix-augmentation format with different word-selection criteria.

Metric	NTP	TF-IDF	P( $t\|S$ )	Max(MI( $S$ ; $t$ ))
ROUGE-1	0.2838	0.3298	0.3301	0.3313
ROUGE-2	0.0862	0.1093	0.1091	0.1102
ROUGE-L	0.2127	0.2547	0.2534	0.2554
Average	0.1942	0.2313	0.2309	0.2323

Table 6: ROUGE scores on XSUM using Llama-3.2-1B-Instruct. Higher is better.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Task Name	Data Format
3-digit Addition	$A_{1}A_{2}A_{3}+B_{1}B_{2}B_{3}$
3-digit Addition	$=C_{1}C_{2}C_{3}C_{4}$
2-digit Multiplication	$A_{1}A_{2}\times B_{1}B_{2}$
2-digit Multiplication	$=C_{1}C_{2}C_{3}C_{4}$
4-digit Logarithm	$\log_{10}{(A_{1}A_{2}A_{3}A_{4})}$
4-digit Logarithm	$=C_{1}.C_{2}C_{3}C_{4}$

Appendix

Table 7: Data format across arithmetic tasks for fine-tuning Qwen-Math model. Mathematical signs are removed for training minGPT. The meaning of the numbers in 2-digit CR: the number of heads is

A_{1}A_{2}

, the number of legs is

B_{1}B_{2}

, the number of chickens is

C_{1}C_{2}

, and the number of rabbits is

{C_{3}C_{4}}

. Abbreviations: GCD — greatest common divisor, CR — chicken-rabbit problem.

Figure 8: Running time of the preprocessing pipeline on WikiText-2 and XSUM datasets across data sampling rates.

Task	Metric(s)	NTP	TF-IDF	P( $t\|S$ )	Max(MI( $S$ ; $t$ ))
CoLA	Matthew’s Corr	0.4039	0.4411	0.4325	0.4803
SST-2	Accuracy	0.9266	0.9197	0.9243	0.9266
MRPC	F1 / Accuracy	0.8711 /0.8064	0.8703/ 0.8137	0.8707/ 0.8137	0.8527/0.7892
STS-B	Pearson-Spearman Corr	0.8518 / 0.8498	0.8444/0.8438	0.8465/0.8465	0.8391/0.8398
QQP	F1 / Accuracy	0.8679/0.8994	0.8687/ 0.9010	0.8663/0.8996	0.8688 /0.9008
MNLI	Accuracy	0.8262	0.8222	0.8235	0.8288

Appendix

Table 8: Comparison of preprocessing strategies on the GLUE benchmark using GPT-2. NTP denotes the standard unaugmented baseline; TF-IDF, P(

t|S

), and Max(MI(

S

;

t

)) use the prefix-augmentation format with different word-selection criteria.

Panel of Figure 9 — Figure 9: Training curves of different target token orders with significant differences (not displaying all permutations of target token orders due to complexity). (a) 3-digit Addition ${A_{1}A_{2}A_{3}+B_{1}B_{2}B_{3}=C_{1}C_{2}C_{3}C_{4}}$ on NanoGPT (0.09M parameters). (b) 2-digit Multiplication ${A_{1}A_{2}\times B_{1}B_{2}=C_{1}C_{2}C_{3}C_{4}}$ on GPT-2-mini (2.67M parameters). (c) Multi-label classification on Qwen-2.5-1.5B-Instruct using ToxicComment dataset, whose labels [ $toxic,obscene,insult,identity\_hate$ ] are retained and marked as [ ${C_{1},C_{2},C_{3},C_{4}}$ ]. Each curve represents a distinct token order. (d) Text generation on GPT-2-small (137M parameters) using WikiText-2 dataset. Abbreviations: Acc – Maximal accuracy at fixed iteration, Rank – the rank of Acc among all permutations of target token orders.

Autoregressive Language Modeling

Autoregressive language modeling predicts sequences one step at a time, using prior tokens to generate text, speech, code, and other data.

See all 128 papers

OrderGPT: Generation Order as a Learnable Modeling Choice
Oct 19, 2026Chenhe Yang, Yuhuan Lyu, Tianyang HuAutoregressive Modeling Autoregressive Generation
SAGA Lab, The Chinese University of Hong Kong, Shenzhen
Decoupling Token Roles in Autoregressive Pretraining
Sep 27, 2026Suqin Yuan, Runqi Lin, Kevin Qinghong Lin +4Language Model Pretraining Autoregressive Language Modeling
University of Sydney · University of Oxford · Southeast University
Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models
Jun 22, 2026Jiawei Xu, Minghui Liu, Aakriti Agrawal +2Masked Diffusion Models Language Model Decoding
University of Maryland, College Park · University of California, Los Angeles.

Oct 19, 2026J/K move · Enter open · S save

OrderGPT: Generation Order as a Learnable Modeling Choice

Chenhe Yang, Yuhuan Lyu, Tianyang Hu

SAGA Lab, The Chinese University of Hong Kong, Shenzhen

Autoregressive Modeling Autoregressive Generation

Open paper AskPDF arXiv

Generative modeling of discrete data has two popular paradigms: autoregressive models such as GPT, which factorize a sequence in a left-to-right (L2R) order, and masked diffusion models, which are essentially any-order autoregressive modeling. Rather than competing methods, these are the two extreme points of a single axis -- the generation order. For structured discrete domains without a natural traversal, both extremes may be suboptimal, yet the intermediate choice -- a single well-adapted order -- remains largely unexplored. We introduce OrderGPT, which builds on Any-Order GPT to learn a task-level generation order without canonical-order supervision. Our key observation is that predicting variables from randomly sampled reveal prefixes produces recurring attention patterns in the backbone that provide a model-dependent signal of predictive dependencies. We aggregate this signal into a reciprocal dependency graph and use a spectral teacher to derive continuous generation priorities. A lightweight OrderMLP integrated into the GPT backbone learns to predict these priorities from attention, incorporating order learning into the architecture within a single training run. Training gradually shifts toward the learned order, which is then frozen for subsequent autoregressive training and generation. The same mechanism adapts across modalities: adaptively it recovers near-L2R dependency structure in language, learns anatomically organized factorizations on human pose, and discovers a factorization on tabular data that places statistically and semantically related fields nearby, while improving downstream modeling metrics across all domains. The code is available at https://github.com/Chenhe-yang/OrderGPT/.

Abstract

Autoregressive Language Modeling

Explore similar work

OrderGPT: Generation Order as a Learnable Modeling Choice

Decoupling Token Roles in Autoregressive Pretraining

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models