adu2021 Predicting the Order of Upcoming Tokens Improves Language Models
Core Concept
Token Order Prediction (TOP) is a lightweight auxiliary training objective that teaches models to rank upcoming tokens by proximity. Instead of predicting exact future tokens (which is difficult), TOP uses learning-to-rank loss to teach the model to understand token sequencing. This simple auxiliary objective improves performance across mathematics, coding, and standard NLP benchmarks while requiring only a single additional unembedding layer.
Architecture Overview
- Learning-to-Rank Loss: Rank tokens by distance rather than exact prediction
- Single Extra Layer: Minimal architectural overhead vs. multi-token prediction
- Auxiliary Objective: Combined with next-token prediction during training
- Task Generalization: Benefits across math, code, and general NLP
- Efficiency: No increase in inference latency or model size
Implementation Steps
Stage 1: Design Token Order Prediction Objective
Formulate the learning-to-rank loss for token ordering.
# Token order prediction: learning-to-rank objective
import torch
from torch import nn
import torch.nn.functional as F
class TokenOrderPredictionHead(nn.Module):
"""Predict ordering of upcoming tokens"""
def __init__(
self,
hidden_dim: int,
vocab_size: int,
num_future_tokens: int = 5
):
super().__init__()
self.hidden_dim = hidden_dim
self.vocab_size = vocab_size
self.num_future = num_future_tokens
# Projection from hidden state to ranking scores
self.ranking_head = nn.Linear(hidden_dim, vocab_size)
def forward(
self,
hidden_states: torch.Tensor,
future_tokens: torch.Tensor
) -> torch.Tensor:
"""
Compute ranking scores for upcoming tokens.
Args:
hidden_states: [batch, seq_len, hidden_dim]
future_tokens: [batch, seq_len, num_future]
Returns:
ranking_scores: [batch, seq_len, vocab_size]
"""
# Project hidden states to token scores
scores = self.ranking_head(hidden_states) # [batch, seq_len, vocab_size]
return scores
class TopLoss(nn.Module):
"""Learning-to-rank loss for token ordering"""
def __init__(self, loss_type: str = "listwise"):
"""
Args:
loss_type: "listwise", "pairwise", or "pointwise"
"""
super().__init__()
self.loss_type = loss_type
def forward(
self,
predicted_scores: torch.Tensor, # [batch*seq_len, vocab_size]
target_tokens: torch.Tensor, # [batch*seq_len, num_future]
target_positions: torch.Tensor # [batch*seq_len, num_future] - distance from current
) -> torch.Tensor:
"""
Compute learning-to-rank loss.
The model should score tokens higher if they appear soon.
"""
if self.loss_type == "listwise":
return self._listwise_loss(predicted_scores, target_tokens, target_positions)
elif self.loss_type == "pairwise":
return self._pairwise_loss(predicted_scores, target_tokens, target_positions)
else:
return self._pointwise_loss(predicted_scores, target_tokens, target_positions)
def _listwise_loss(
self,
predicted_scores: torch.Tensor,
target_tokens: torch.Tensor,
target_positions: torch.Tensor
) -> torch.Tensor:
"""
Listwise loss: maximize probability of correct ranking.
Tokens appearing sooner should have higher scores.
"""
batch_size = predicted_scores.shape[0]
loss = 0
for i in range(batch_size):
scores = predicted_scores[i] # [vocab_size]
tokens = target_tokens[i] # [num_future]
positions = target_positions[i] # [num_future]
# Get scores for target tokens
token_scores = scores[tokens]
# Closer tokens should have higher scores
# Create preference pairs
for j in range(len(tokens)):
for k in range(j + 1, len(tokens)):
# Token at position j should have higher score if closer (smaller position)
if positions[j] < positions[k]:
# score[j] should be > score[k]
margin = token_scores[j] - token_scores[k]
else:
# score[k] should be > score[j]
margin = token_scores[k] - token_scores[j]
# Hinge-like loss
loss += F.relu(1.0 - margin)
return loss / batch_size
def _pairwise_loss(
self,
predicted_scores: torch.Tensor,
target_tokens: torch.Tensor,
target_positions: torch.Tensor
) -> torch.Tensor:
"""
Pairwise loss: compare tokens at different distances.
"""
batch_size = predicted_scores.shape[0]
# Get scores for target tokens
batch_idx = torch.arange(batch_size).unsqueeze(1)
token_scores = predicted_scores[batch_idx, target_tokens] # [batch, num_future]
# Create pairs: closer token vs farther token
# Closer should have higher score
loss = 0
for dist_close in range(target_positions.shape[1]):
for dist_far in range(dist_close + 1, target_positions.shape[1]):
score_close = token_scores[:, dist_close]
score_far = token_scores[:, dist_far]
# Loss: margin between close and far
margin_loss = F.relu(1.0 - (score_close - score_far))
loss += margin_loss.mean()
return loss
def _pointwise_loss(
self,
predicted_scores: torch.Tensor,
target_tokens: torch.Tensor,
target_positions: torch.Tensor
) -> torch.Tensor:
"""
Pointwise loss: score tokens relative to position.
"""
batch_size = predicted_scores.shape[0]
# Get scores for target tokens
batch_idx = torch.arange(batch_size).unsqueeze(1)
token_scores = predicted_scores[batch_idx, target_tokens]
# Target: closer tokens should have higher scores
# Invert positions (smaller distance = larger target score)
max_pos = target_positions.float().max()
position_targets = (max_pos - target_positions.float()) / max_pos
# Regression loss
loss = F.mse_loss(token_scores, position_targets)
return loss
Stage 2: Integrate TOP with Next-Token Prediction
Combine TOP auxiliary objective with standard language modeling.
# Combined training with TOP auxiliary objective
class LanguageModelWithTOP(nn.Module):
"""LM with top-k token order prediction"""
def __init__(
self,
hidden_dim: int = 4096,
vocab_size: int = 32000,
num_layers: int = 32
):
super().__init__()
# Standard transformer LM
self.transformer = TransformerLM(hidden_dim, vocab_size, num_layers)
# Next-token prediction head
self.lm_head = nn.Linear(hidden_dim, vocab_size)
# Token order prediction head
self.top_head = TokenOrderPredictionHead(hidden_dim, vocab_size)
self.top_loss_fn = TopLoss(loss_type="pairwise")
def forward(
self,
input_ids: torch.Tensor,
future_token_ids: torch.Tensor = None
) -> Dict:
"""
Forward pass with both NTP and TOP.
Args:
input_ids: [batch, seq_len]
future_token_ids: [batch, seq_len, num_future] (for TOP)
Returns:
ntp_logits, top_scores
"""
# Get hidden states
hidden_states = self.transformer(input_ids)
# Next-token prediction
ntp_logits = self.lm_head(hidden_states)
# Token order prediction (if provided)
top_scores = None
if future_token_ids is not None:
top_scores = self.top_head(hidden_states, future_token_ids)
return {
"ntp_logits": ntp_logits,
"top_scores": top_scores,
"hidden_states": hidden_states
}
class CombinedTrainingLoop:
"""Train with NTP + TOP objectives"""
def __init__(
self,
model: LanguageModelWithTOP,
lr: float = 1e-4,
top_weight: float = 0.1
):
self.model = model
self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
self.ntp_loss_fn = nn.CrossEntropyLoss()
self.top_loss_fn = TopLoss(loss_type="pairwise")
self.top_weight = top_weight
def extract_future_tokens(
self,
input_ids: torch.Tensor,
num_future: int = 5
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Extract future tokens and their distances.
Returns:
future_token_ids: [batch, seq_len, num_future]
future_positions: [batch, seq_len, num_future]
"""
batch_size, seq_len = input_ids.shape
future_tokens = torch.zeros(
batch_size, seq_len, num_future,
dtype=torch.long, device=input_ids.device
)
future_positions = torch.zeros_like(future_tokens)
for i in range(seq_len):
for j in range(num_future):
if i + j + 1 < seq_len:
future_tokens[:, i, j] = input_ids[:, i + j + 1]
future_positions[:, i, j] = j + 1
return future_tokens, future_positions
def train_step(self, batch: Dict) -> Dict:
"""Single training step with NTP + TOP"""
input_ids = batch["input_ids"]
attention_mask = batch.get("attention_mask")
# Extract future tokens for TOP
future_tokens, future_positions = self.extract_future_tokens(input_ids)
# Forward pass
outputs = self.model(input_ids, future_tokens)
ntp_logits = outputs["ntp_logits"]
top_scores = outputs["top_scores"]
# Next-token prediction loss
ntp_loss = self.ntp_loss_fn(
ntp_logits.view(-1, ntp_logits.shape[-1]),
input_ids.view(-1)
)
# Token order prediction loss
top_loss = 0
if top_scores is not None:
top_loss = self.top_loss_fn(
top_scores.view(-1, top_scores.shape[-1]),
future_tokens.view(-1, future_tokens.shape[-1]),
future_positions.view(-1, future_positions.shape[-1])
)
# Combined loss
total_loss = ntp_loss + self.top_weight * top_loss
# Backward
self.optimizer.zero_grad()
total_loss.backward()
self.optimizer.step()
return {
"total_loss": total_loss.item(),
"ntp_loss": ntp_loss.item(),
"top_loss": top_loss.item() if isinstance(top_loss, torch.Tensor) else top_loss
}
def set_top_weight(self, weight: float):
"""Adjust TOP weight (useful for curriculum learning)"""
self.top_weight = weight
Stage 3: Evaluation on Benchmarks
Test TOP across diverse tasks.
# Benchmark evaluation
class TOPEvaluator:
"""Evaluate TOP improvements"""
def __init__(self, model_with_top: LanguageModelWithTOP):
self.model = model_with_top
def evaluate_on_benchmarks(self) -> Dict:
"""
Evaluate on standard benchmarks.
Reproduction of paper results.
"""
benchmarks = {
"arc_challenge": self.evaluate_benchmark("arc_challenge"),
"hellaswag": self.evaluate_benchmark("hellaswag"),
"mmlu": self.evaluate_benchmark("mmlu"),
"gsm8k": self.evaluate_benchmark("gsm8k"),
"math": self.evaluate_benchmark("math"),
"humanevals": self.evaluate_benchmark("humanevals"),
"mbpp": self.evaluate_benchmark("mbpp")
}
return benchmarks
def evaluate_benchmark(self, benchmark_name: str) -> Dict:
"""Evaluate on single benchmark"""
# Load test set
test_set = self.load_benchmark_data(benchmark_name)
correct = 0
total = len(test_set)
for example in test_set:
prompt = example["prompt"]
answer = example["answer"]
# Generate
generated = self.generate(prompt, max_length=256)
# Check correctness
if self.check_correctness(generated, answer, benchmark_name):
correct += 1
accuracy = correct / total
return {
"benchmark": benchmark_name,
"correct": correct,
"total": total,
"accuracy": accuracy
}
def generate(self, prompt: str, max_length: int = 256) -> str:
"""Generate with model"""
tokens = self.model.transformer.tokenizer.encode(prompt)
tokens = torch.tensor(tokens).unsqueeze(0)
for _ in range(max_length):
outputs = self.model(tokens)
logits = outputs["ntp_logits"][:, -1, :]
next_token = logits.argmax(dim=-1)
tokens = torch.cat([tokens, next_token.unsqueeze(1)], dim=1)
return self.model.transformer.tokenizer.decode(tokens[0])
def check_correctness(self, generated: str, reference: str, benchmark: str) -> bool:
"""Check if generation matches reference"""
if benchmark in ["gsm8k", "math"]:
# Extract final number
gen_num = self.extract_number(generated)
ref_num = self.extract_number(reference)
return gen_num == ref_num
else:
# Exact match or substring
return reference in generated
def extract_number(self, text: str) -> float:
"""Extract final number from text"""
import re
numbers = re.findall(r'-?\d+\.?\d*', text)
if numbers:
return float(numbers[-1])
return None
def load_benchmark_data(self, name: str):
"""Load benchmark data"""
# Placeholder: in practice, load from HuggingFace
return []
Practical Guidance
Hyperparameters
- TOP Weight: 0.1-0.3 relative to NTP loss (higher = more TOP emphasis)
- Future Tokens: 5 is optimal (further tokens have diminishing signal)
- Loss Type: Pairwise generally works best
- Training Duration: TOP converges quickly; add after 10-20% of pretraining
When to Use TOP
- Improving language model performance across diverse tasks
- Math and coding tasks (TOP particularly helps here)
- General pretraining where you want better few-shot performance
- Resource-constrained training (minimal overhead)
When NOT to Use
- Models already performing optimally on target tasks
- Very long-context scenarios (future token extraction becomes expensive)
- Real-time systems where any overhead matters
Performance Expectations
- Math benchmarks: +2-4% improvement
- Coding benchmarks: +1-3% improvement
- General NLP: +0.5-2% improvement
- Computational overhead: <5% training time increase
Design Insights
TOP works because predicting exact future tokens is a harder auxiliary task than necessary. By reformulating as a ranking problem ("which tokens appear first?"), the model learns token sequencing patterns without the regression difficulty. This is a goldilocks auxiliary objective—harder than next-token only, but easier than exact multi-token prediction.
Reference
Predicting the Order of Upcoming Tokens Improves Language Models. arXiv:2508.19228