Hierarchical Latent Prediction for Language Models

The paper addresses limitations of standard Next-Token Prediction (NTP) in language model pretraining. The authors note that while standard Next-Token Prediction (NTP) lays the foundation of language model pretraining, its teacher-forced training paradigm may not be optimal for long-horizon reasoning and planning. NTP introduces "a training-inference mismatch known as exposure bias: during inference time, the auto-regressive token generation relies on its own previous outputs, leading to compounding errors which could degrade long-range generation quality."

Recent attempts to mitigate this include Multi-Token Prediction (MTP) and Next-Latent prediction (NextLat). However, the authors argue that those auxiliary objectives either have a limited horizon or suffer from compounding error from multi-step rollout. Specifically, NextLat's latent dynamics model predicts one step ahead, meaning the self-predictive learning signal, while effective in inducing local transition consistency, provides only indirect gradient pressure to form representations that capture structure at longer temporal scales.

The authors observe that language incorporates hierarchical structures at multiple scales, from characters, phrases to sentences, paragraphs and discourses, which motivates extending NextLat with hierarchical learning signals at different granularities. The central research question is: can we introduce temporal hierarchy directly into the latent-space self-prediction objective, shaping the representation of a language model transformer to encode multi-scale predictive structure during pretraining?

The paper answers affirmatively by introducing Hierarchical Latent Prediction (HiLP), which introduces an auxiliary higher-level abstract latent to help reduce the error accumulation effect in latent-space rollouts.

The architecture works as follows: "A transformer Gθ produces lower-level latents ht = Gθ(X1:t). A sliding-window attention (SWA) module with window k produces higher-level abstract latents ut = SWA(ht−k+1:t), a deterministic function of the last k lower states. A higher-level dynamics model is trained to predict k steps ahead in this abstract space. This explicit coarser-scale prediction objective encourages lookahead planning beyond single-step latent transitions. A combined NTP head then conditions the next-token prediction on both level latents, taking advantage of abstract lookahead information to improve subsequent predictions."

Four parameterised maps are jointly learned: output head pθ, lower dynamics pψ, upper dynamics pϕ, and combined head pρ.

The total training objective is a weighted sum of five terms: the standard NTP loss, the NextLat (lower-level) transition consistency loss, the KL term, the higher-level transition consistency loss, and the combined NTP loss:

• Lntp = E[−log pθ(Xt+1 ht)] — standard NTP loss (Eq. A)

• Lh = E[SL1(sg[ht+i], ĥt+i)] — lower-level rollout consistency loss with SmoothL1 over d rollout steps (Eq. B)

• LKL = E[DKL(p sg θ(·sg[ht+i]) ‖ p sg θ(·ĥt+i))] — KL divergence between distributions from real vs. predicted latents, decoded through a frozen copy of the NTP head (Eq. C)

• Lu = E[SmoothL1(sg[ut+k], ût+k)] — higher-level transition consistency loss, where the higher-level dynamics model predicts k steps ahead with no token input (Eq. D)

• Lcntp = E[−log pρ(Xt+1 ht, ũt)] — combined NTP loss (Eq. E)

A key design detail is the stop-gradient in the combined NTP path: "We use this stop-gradient in the combined NTP path to keep the token-level representation ht governed by the standard NTP and lower-level transition consistency losses, while allowing the combined head to train the SWA module and its higher-level representation. Without this separation, the combined-head cross-entropy would also update the lower latents through the SWA window, double-counting token-level supervision and creating a competing optimization signal for ht."

At optimality, the losses enforce four conditions:

• (I) pθ(Xt+1 ht) = P(Xt+1 X1:t) — next-token consistency

• (II) pψ(ht+1 ht, Xt+1) = P(ht+1 X1:t+1) — lower transition consistency

• (III) pϕ(ut+k ut) = P(ut+k X1:t) — higher transition consistency

• (IV) pρ(Xt+1 ht, ut) = P(Xt+1 X1:t) — next-token consistency via combined head

Crucially, At inference time, only the standard NTP head is used, the entire hierarchical apparatus serves purely as an auxiliary training signal. The paper emphasizes that the entire hierarchical apparatus can be removed at inference time with no architectural overhead, preserving NextLat's option of being a pure training-time intervention.

1B-parameter models are trained on 100B tokens using 8 × NVIDIA B200 GPUs. Evaluations use "LM Evaluation Harness (Gao et al., 2024) to evaluate the zero-shot performance of the models on HumanEval coding benchmark (Chen et al., 2021), and DataComp for LLMs (Li et al., 2025) to evaluate the models on a set of symbolic and multi-step benchmarks."

Key hyperparameters include: vocabulary size 100,352, training sequence length 8,192, 24 hidden layers, 16 attention heads, latent window size 4, latent lookahead steps 4, and loss weights λntp=1.0, λh=10.0, λKL=0.1, λu=1.0, λcntp=0.5.

On HumanEval pass@1: NTP scores 8.77, MTP 9.21, NextLat 10.58, and HiLP 11.33 — the highest. The authors state that results show that HiLP improves performance from the multi-scale latent prediction.

On code data held-out validation split: HiLP (NTP head) achieves Avg Acc. Tok. of 3.41 and Avg match of 0.592, versus MTP's 2.57/0.545 and NextLat's 3.26/0.580. On Nemotron-ClimbMix data: HiLP achieves 2.59/0.491 vs MTP's 2.11/0.162 and NextLat's 2.57/0.483.

Figure 3 shows latent prediction cross-entropy losses across rollout horizons up to 8 steps ahead. The authors report: As the curves show, HiLP preserves the near-term cross-entropy of NextLat while producing lower future prediction error at longer horizons, suggesting that the hierarchical latent mitigates multi-step error accumulation. The paper states that HiLP has comparable near-term latent prediction loss and lower long-horizon future and excess cross-entropy, indicating more stable multi-step latent rollout.

HiLP adds ≈190–200M auxiliary parameters (draft module) on top of the 1.06B trunk, bringing total to 1.27B, but inference parameters remain 1.06B — identical to all baselines since only the trunk and LM head are used. Training throughput: NTP baseline reaches ≈126K tokens/s/GPU, NextLat retains ≈83% of that throughput, HiLP retains ≈65%, and MTP drops to ≈49%. An input-combination ablation shows that the glu cross mode (gating a projected lower latent with the conditioning latent) reduces the latent predictor cost from 40D2 to 14D2 with default m=4, and approximately halves the combined-head logit computation compared to concat mode.

The paper positions HiLP against:

1. Pretraining objectives beyond NTP (e.g., MTP, future summaries): "these approaches usually operate on a single granularity scale and have limited horizon capability due to compounding error. HiLP differs in two ways: it introduces two latent prediction pathways at different temporal scales, so the coarser pathway supervises long-horizon structure directly rather than through repeated single-step unrolling; and all auxiliary heads are dropped at inference."

2. Latent-space language modeling (Large Concept Models, CALM, Coconut): "All of these approaches change the inference-time generation process to operate on latents, requiring bespoke decoding procedures or latent-to-text decoders. HiLP is the opposite design point: latent prediction is used purely as an auxiliary training signal to shape representations, while inference remains standard token-level autoregressive decoding with zero added latency."

3. Hierarchical models (MegaByte, Byte Latent Transformer, H-Net, hierarchical world models in RL): "HiLP differs in a way that the hierarchy complexity is not consumed at inference time: to our knowledge it is the first to introduce hierarchical latent prediction as a pretraining objective for autoregressive language models, where the hierarchy shapes the representation during training and is then removed entirely."

The conclusion states: "HiLP introduces temporal hierarchy into the latent-space self-prediction objective, shaping the representation of language model transformers to encode multi-scale predictive structure during pretraining. Experiments show that it improves both speculative decoding and language models on tasks that need longer-horizon reasoning and planning."

Future work directions include dynamically choosing the lookahead horizon (currently a manually set hyperparameter) and potentially using the combined NTP head at inference time if we trade speed for accuracy.

By applying HiLP, I can improve language model pretraining in the following concrete ways:

• Add a time-hierarchical auxiliary objective to the LM training loss.

Instead of only predicting the next token or the next latent, I train a lower-level transition predictor and a higher-level abstract-latent predictor that forecasts k steps ahead (e.g., k=4) in a compressed representation built by sliding-window attention over the last k lower-level states. This forces the transformer to encode multi-scale predictive structure during pretraining.

• Inject abstract lookahead into the next-token prediction head during training.

I train a combined head pρ(X t+1 h t, u t) that conditions token prediction on both the current token-level latent and the higher-level abstract latent, while applying stop-gradient on the lower latent. This keeps the token-level representation optimized by the standard NTP loss and the lower transition loss, while the higher-level abstract path learns useful lookahead information. At inference time, this combined head can be discarded.

• Use the higher-level latent predictor as a self-speculative draft model.

Because the higher-level dynamics model outputs a future abstract latent û t+k without multi-step token rollout, I can use it to draft multiple future tokens more accurately, improving acceptance rates and reducing latency in speculative decoding while keeping inference parameter count identical to the base LM.

• Reduce compounding-error / exposure-bias effects for long-horizon generation.

The hierarchical latent prediction directly supervises multi-step future structure instead of relying on repeated single-step unrolling during training. This yields more stable rollouts: the resulting model has lower excess cross-entropy at horizons up to 8 steps, which translates into better long-range code generation and multi-step reasoning.

What the improved AI system can do:

• Better multi-step reasoning and planning:

A 1B-parameter model trained with HiLP achieves 11.33 pass@1 on HumanEval, outperforming NTP (8.77), MTP (9.21), and NextLat (10.58). It can generate coherent multi-step code that requires longer-range dependencies.

• Faster and more accurate generation:

On held-out code data, the HiLP-trained model using its standard NTP head gives 3.41 average accepted tokens and 0.592 average match in self-speculative decoding, versus MTP’s 2.57/0.545 and NextLat’s 3.26/0.580. The system can therefore decode with lower latency while producing more tokens accepted per draft step.

• Longer-horizon prediction without added inference cost:

The hierarchical apparatus is removed at inference time, so generation remains standard token-level autoregression with zero added latency and a 1.06B inference parameter count. The improved system can maintain near-term accuracy while reducing future prediction error over longer rollout horizons.

• More efficient training than multi-token prediction:

HiLP retains 65% training throughput relative to NTP, while MTP drops to 49%. The improved training system can therefore incorporate multi-scale latent supervision at a modest throughput cost, producing a model better suited for reasoning, planning, and speculative-decoding acceleration.

While standard Next-Token Prediction (NTP) lays the foundation of language model pre- training, its teacher-forced training paradigm may not be optimal for long-horizon reasoning and planning. Recent works such as Multi-Token Prediction (MTP) and Next-Latent prediction (NextLat) try to mitigate the problem through predicting multiple future tokens and self-supervised prediction in the latent space. However, those auxiliary objectives either have a limited horizon or suffer from compounding error from multi-step rollout. We introduce Hierarchical Latent Prediction (HiLP), which introduces an auxiliary higher-level abstract latent to help reduce the error accumulation effect in latent-space rollouts. Experiments show that HiLP can lead to longer-horizon coherent belief state representation and demonstrate the effectiveness of our method across coding and multi-step reasoning benchmarks, and offers more speculative decoding efficiency.