Multiple Token Divergence: Measuring and Steering In-Context Computation Density

TL;DR

This work introduces Multiple Token Divergence (MTD), a practical measure of in-context computational effort for language models defined as the KL divergence between the full model's next-token distribution and a shallow, auxiliary MTP predictor. Unlike latent-space bottleneck methods, MTD operates directly on output distributions and can be computed from pre-trained models without further training. The authors additionally propose Divergence Steering, a decoding technique that interpolates between the full-model and MTP distributions (with optional fixed-entropy projection) to bias generation toward more or less computationally dense tokens. Empirically, MTD correlates with problem difficulty and lower MTD aligns with higher reasoning accuracy on math benchmarks, while Divergence Steering can enhance creativity in algorithmic and writing tasks. The approach offers a lightweight, non-invasive tool for analyzing and shaping the computational dynamics of language models, with potential use in dynamic compute allocation, solution convergence signaling, and intrinsic motivation for open-ended learning.

Abstract

Measuring the in-context computational effort of language models is a key challenge, as metrics like next-token loss fail to capture reasoning complexity. Prior methods based on latent state compressibility can be invasive and unstable. We propose Multiple Token Divergence (MTD), a simple measure of computational effort defined as the KL divergence between a model's full output distribution and that of a shallow, auxiliary prediction head. MTD can be computed directly from pre-trained models with multiple prediction heads, requiring no additional training. Building on this, we introduce Divergence Steering, a novel decoding method to control the computational character of generated text. We empirically show that MTD is more effective than prior methods at distinguishing complex tasks from simple ones. On mathematical reasoning benchmarks, MTD correlates positively with problem difficulty. Lower MTD is associated with more accurate reasoning. MTD provides a practical, lightweight tool for analyzing and steering the computational dynamics of language models.

Figures (18)

• Figure 1: Comparison between the architecture of a PHi model (left) and of a MTP model (right).
• Figure 2: Divergence Steering on a $K$=3 simplex with temperature curve for $p$, geodesic interpolation between from $m$ to $p$ and beyond, and projection onto distributions with a fixed entropy of $H(p).$
• Figure 3: Distributions corresponding to Figure \ref{['fig:simplex_diagram']}. Geodesic interpolation $s_\alpha$, and the entropy of the resulting distribution (top). The same distributions projected onto the surface with fixed entropy, $\hat{s}_\alpha$ (bottom).
• Figure 4: Normalized PHi or MTD loss of the four different model types on each of the five tasks. Only in-context language learning (ICLL) requires sophisticated in-context computation. This is reflected by the scores, with the exception of the MTD model without access to the latest embedding, which assigns high MTD also to the memorized programs task (see the discussion in Sections \ref{['sec:phi_vs_mtd']} and \ref{['sec:exp_pfa_on_different_tasks']}). Bootstrapped mean with $95\%$ confidence intervals across 8 runs.
• Figure 5: Partial correlation of PHi or MTD loss with the complexity of the modelled PFA, controlling for NLL. Also here, MTD without latest embedding access is the outlier.