Internal Planning in Language Models: Characterizing Horizon and Branch Awareness
Authors: Muhammed Ustaomeroglu
Baris Askin
Gauri Joshi
Carlee Joe-Wong
Guannan Qu
Venue: ICLR 2026 Poster (2026)
Keywords: language model
LLM
deep learning
planning
explainability
interpretability
information theory
mechanistic interpretability
Primary area: interpretability and explainable AI

Abstract: The extent to which decoder-only language models (LMs) engage in planning, that is, organizing intermediate computations to support coherent long-range generation, remains an important question, with implications for interpretability, reliability, and principled model design. Planning involves structuring computations over long horizons, and considering multiple possible continuations, but how far transformer-based LMs exhibit them without external scaffolds, e.g., chain-of-thought prompting, is unclear. We address these questions by analyzing the hidden states at the core of transformer computations, which capture intermediate results and act as carriers of information. Since these hidden representations are redundant and encumbered with fine-grained details, we develop a pipeline based on vector-quantized variational autoencoders that compresses them into compact summary codes. These codes enable measuring mutual information and analyzing the computational structure of the underlying model behavior. Using this framework, we study planning in LMs across synthetic grammar, path-finding tasks, and natural language datasets, focusing on two planning properties: (i) the planning horizon of pre-output computations, and (ii) the extent to which the model considers alternative valid continuations. As a separate downstream use of the same pipeline, we also analyze how decision-relevant information is distributed across layers and earlier prefix blocks when producing next-token predictions. Together, these analyses advance our understanding of planning in LMs and provide a general-purpose pipeline for inspecting internal model dynamics. Our results reveal that the effective planning horizon is task-dependent, that models implicitly preserve information about unused correct continuations, and that predictions draw most on recent computations, though earlier blocks remain informative.

Published as a conference paper at ICLR 2026
INTERNAL PLANNING IN LANGUAGE MODELS: CHAR- ACTERIZING HORIZON AND BRANCH AWARENESS
Muhammed Ustaomeroglu', Baris Askin', Gauri Joshi, Carlee Joe- Joe-Wong, Guannan Qu Department of Electrical and Computer Engineering
Carnegie Mellon University Pinsburgh, PA 15213, USA {mustaome, baskin, gaurij, cjoewong, gqu]@andrew. cmu, edu
ABSTRACT
The extent which decoder-only language models (LMSs) engage in planning, that is, organizing intermediate computations to support coherent long-range gen- eration, remains an important question, with implications for interpretabilityy reli- ability, and principled model design. Planning involves structuring computations over long horizons, and considering multiple possible continuations, but how far transformer- baseed LMs exhibit them without external scaffolds, e.g., chain-of- thought prompting, is unclear. We address these questions by analyzing the hid- of transformer computations, which capture intermediate
den states the results and act as carriers of information. Since these hidden representations are redundant and encumbered with fine-grained details, we develop a pipeline based on qucttr-quantized variational autoencoders that compresses them into compact summary codes. These codes enable measuring mutual information and analyzing the computational structure of the underlying model behavior. Using this frame- work, we study planning LMs across synthetic grammar,
tasks, and natural language datasets, focusing two planning properties: (i) the plan- ning horizon of pre-output computations, and (ii) the extent to which the model considers alternative valid continuations.
same pipeline, we also analyze how As separate downstream use of the information is distributed across layers and earlier prefix blocks when producing next-token predictions. To- gether, these analyses advance our understanding of planning in LMs and provide general-pure purposs pipeline for inspecting internal model dynamics. Our results reveal that the effective planning horizon is task-dependent, that models implic- itly preserve information about unused correct continuations, and that predictions draw most on recent computations, though earlier blocks remain informative.
INTRODUCTION
Language models (LMs) have advanced so rapidly that they now engage in open-ended conver- sation, write functional code, and even solve challenging mathematical problems (Brown, 2020
Achiam et al., 2023 Grattafiori et al., 2024: Touvron et al., 2023). Beyond these raw capabilitics
scaffolding and augmentation techniques further enhance reasoning and planning (Wei et al. 2022
Wang et al. 2024c Sel et al. 2025 Yao et al. 2023 Besta et al.
2024). In parallel, hybrid systems integrating LMs with symbolic planners external tools have achieved state-of-the performance
in embodied and tool-use domains (Zhao el al., 2023 Wang et al. 2024a Shen et al. 2023).
All the above success suggests that LMs have certain capability of "planning ahead," similar to what humans do in generating coherent long speeches and long-horizon problem solving. Yet, the typical training method for LMs, next token prediction, focuses only on predicting the next token, which seemingly suggests LMs are myopic. Motivated by this contrast, this paper seeks to understand the planning behavior of LMs and how is affected by next-token prediction (NTP) vs. multi-token prediction (MTP) in training. As "planning behavior" is a broad term that has many facets, we further narrow our focus to the following two canonical aspects of planning. First, good
"Equal contribution

---

Published as a conference paper at ICLR 2026
planners forward-looking. They structure intermediate computations over receding horizon so that near-term actions serve longer-term objectives, an idea made explicit in model- preeiivve control (MPC) and holdd-ooell--aaae approaches (Morari & Lee [1999] Mayne et al. (000 Ha &
Schmidhuber 2018: LeCun 2022). Second, good planners are branch-aware, before committing, they keep multiple plausible futures "alive," comparing candidates rather than greedily following single line, an ability central in classical MPC and world-model-based control methods (Morari &
Lee 1999; Mayne et al., 2000 Hafner et al. 2019). and closely tied to robustness under paraphrase
and logically related prompts (Lin et al. [025] Ahn & Yin [025 Saxena et al. (2024) and to the
success of search-based scaffolds such as Tree- and Graph- of Thoughts (Yao et al. 2023 Besta
et al. 2024). Given the above aspects about planning, this paper seeks to understand,
To what extent are LMs forward-looking and brunch-aware?
How does the training method (NTP vs. MTP) affect these qualities?
We approach the above questions by studying the internal states and computations of an LM, as recent work shows that transformers internally encode rich, high-level abstractions such belief states, game configurations, and world models that extend far beyond their immediate outputs, and that these abstractions can be partially extracted (Shai et al. [0244] Li et al.] [2023 Pal et al., 2023
Richens et al. (225). This makes the internal representations of LMs natural locus for investigat-
ing how they plan and reason. While not aimed at these specific questions, existing methods have explored internal computations of models, e.g. using probing classifiers to test for linguistic fea- tures in hidden states mechanistic analyses identify circuits and disentangled features within transformers (Hewitt & Liang 2019 Elhage et al.) 2021 Bricken et al.) 2023). However, these only indirectly relate to our questions, and despite important progress, these strategies have clear limitations. Circuit discovery requires heavy manual engineering (Elhage et al. [021 Wang et al.) 2023 Chan al.) 2022, Meng et al., 2022), while probing risks conflating genuine representations
with probe artifacts (Hewitt& & Liang 2019; Voita & Titov] [020] Pimentel et al. 2020: Kunz &
Kublmann 2020; Kumar et al. 22022). These limitations necessitate new approaches towards un- derstanding the internal computations that are automated, scalable, and less susceptible to probing's confounding of representations learned by the probe with those in the model itself.
citommation-thoorettt framework to study how Glamning-relate computations are organized inside LMs that is both automated (no manual circuit engineering) and free from confounding caused by learned probes. Specifically, to avoid probe-induced confound- ing, we compute mutual information (MJ) between learned discrete representations that summarize internal states of the LM. By comparing these MI relationships, we characterize whether internal computations exhibit patterns consistent with forward-looking and branch-aware computation. For example, understand the fooward-looking" characteristic, the MI between the internal states of the prefix and that of future generations can shed light on how much computations inside LMs plan ahead for future tokens. In contrast, probing aims to detect whether a piece of information is present in LM hidden state by training an external model to predict the information from the hidden state, which may introduce additional representational power brought by the external model. The resulting prediction losses are often treated as informal proxies for "how much information" a hidden state contains but they depend on the marginal complexity and scale of the chosen target and can vio- late basic cetemmatioo-thhoretee desiderata such as data processing. As we show in our experiments (App. D.2, such probe-based quantities can be strongly influenced by these nuisance factors and need not track true mutual information even in simple controlled settings. In contrast, computing MI between learned representations gives a confound-resisten information, without introducing additional supervised model whose capacity or objective may obscure the underlying LM computation. Moreover, unlike correlation, which is sensitive primarily to linear relationships under a chosen parameterization, MI captures arbitrary statistical dependence and is symmetric and invariant under invertible reparameterizations. To make the MI calculation scalable to potentially very high dimensional hidden state vectors, we employ a Vector- Quanti-Qu Variational Autoencoder (VQ-VEE) to map collections of block outputs into discrete codes that serve as coarse summaries of internal states (Van Den Oord et al.) 2017). We use VQ-VAE as a practical compressor with discrete codebook and transformer encoder, enabling MI estimation over variable-sized blocks; our validation study (App. A.4) supports this choice. We calculate MI between the coarse summaries instead of the raw hidden states. This step is crucial: fine-graine activations are high-dimensional and redundant, making a detailed direct analysis infeasible, while the compressed codes capture the salient distinctions necessary for comparing computations across layers, positions, and contexts. Beyond these planning analyses, the same pipeline also supports
Contribution 1:
We propose
metric of how much two variables share
2

---

Published as a conference paper at ICLR 2026
practical diagnostic: localizing where next-token decision information resides across layers and earlier prefix blocks. While our framework is general and could be applied to other applications in deep learning, in this work, we it to analyze planning in transformer models, and we additionally demonstrate the localization diagnostic as a method application.
Contribution 2: Understanding the planning of LMs. Using the cntemmtiion-thoorede frame- work, we analyze LMs' ability to be (i) forward-looking, (ii) branch-aware, and (iii) as an applica- tion of our framework, how layers and earlier prefix blocks contain next token information, across a symbolic task, a structured reasoning problem, and natural text (S (:::
Horizon of the Plan ($ 3.1 how far ahead a model plans before producing its next token.
Branching in the Plan 3.2 :: to what extent LM internally considers alternative responses.
Diagnostic of the information in the computational history ($ 3.3 where decision-relevan
information about the next token is concentrated across layers and earlier prefix blocks.
Our results show that planning is task-contingent and weakly modulated by the training objective loss, i.e., MTP versus NTP. When the task demands a longer horizon, the LM's pre-output states retain information about tokens beyond the immediate next step. In contrast, in locally syntactic settings, this dependence concentrates near the next token. Training with MTP loss modestly reduces purely myopic behavior. Internally, the model encodes alternatives to the produced answer, and the strength of this branching awareness of the plan varies with task difficulty and model quality. Finally, next-token decisions draw most strongly on the last layers and the most recent token indices.
Related work.
Despite the recent advances in LM planning abilities ( (Wang et al. 2024cc Sel et al.,
2025 Wei et al.
2022; Yao et al. [2023 Besta et al. 2024; Zhao et al., [2023] Wang et al. 2024a Shen et al. 2023), ,, recent studies highlight that significant challenges remain, and current approaches
to LM planning still fall short of fully addressing complex reasoning and decision-making tasks. For
example, Lin et al. (2025) show that models can produce conflicting answers under logically related
prompts despite local plausibility, which shows lack of long- mon-hoorzon planning: Ahn & Yin (025)) )
Saxena et al. (024)): Momennejad et al,] (2023); Wang et al. (2024b) highlight inconsistencies in the model's 's outputs and struggles with planning tasks. Taken together, these findings underscore that
understanding whether and how planning arises in LMs is not only an open empirical challenge, but
also central to both their interpretability and the principled design of future model architectures.
To interpret LMs, some approaches treat the model as a black box and design tasks or benchmarks to gauge reasoning, robustness, or generalization abilities at behavioral level (Srivastava al. [023] Liang et al.) (2022). While such evaluations provide useful insights, they miss several perspectives that can be gained by examining the internal mechanisms of the model. In contrast, mechanistic in- terpretability seeks to reverse-engineet transformer computations into human- uumaa--nderrsandabll parts, treating the residual stream as the main information pathway and attention heads as separate com-
ponents that pass information along (Elhage et al. [2211; Olsson, 2022; Cunningham et al. 2023 Bricken et al., 2023 Hewitt & Liang] 2019] Dunetsky et al., 2024] Lindsey et al.) 2025).
How-
ever, circuit discovery requires significant manual engineering (Wang et al. 2023 Chan al. 2022
Meng et al.] 2022). Beyond empirical tools, mathematically grounded perspectives also shed light
on transformer and LLM interpretability (Liu et al. [0033] Ahn al. 2023 Ustaomeroglu & Qu 2023 Cia et al. (2024). yet these approaches are often criticized for lacking a one-to-one correspon- dence with experimental results, since their proofs typically rely on strong assumptions that may not hold for the highly non-smooth LM architectures. In contrast to them, our method ($ 2 en- ables the study of LMs both behaviorally and structurally, by examining their internal mechanisms, without the need for its reliance on information theooreti tools, our approach can only capture aggregate phenomena, providing insights in an average sense rather than interpreting individual input prompts.
circuit discovery or strong assumptions on the LM. Still, due
different line of interpretability work views LM hidden states as structured representations that can be pprobed"-probingg refers to training lightweight models, often linear classifiers, to read out specific information from hidden states in order to test what the model represents internally. Using probing, researchers have shown that transformer hidden states encode structured belief-state and world-model-like information (Shai et al. 2024| Gurnee & Tegmark, 2024 Hazineh et al.) 2023). Other works demonstrate that a single hidden state can carry information about multiple future tokens (Pal et al. 2023 Wu et al.) and that probing can reveal the underlying algorithms LLMs use to solve tasks (Allen-Zhr 2024). Aithough there are some mathematical foundations for them (Xu
3

---

Published as a conference paper at ICLR 2026
et all), standard accuracy- based probing has been criticized for combining what the probe can learn with what the representation actually encodes, making the results sensitive to probe capacity, data size, and hyperparameters (Hewitt & Liang 2019 Pimentel et al.) 2020). Further. high probe scores often come from exploiting superficial linear context cues rather than genuine structural knowledge (Kunz & Kuhlmann 2020) probing can even reveal features that a model does not use for its task |Ravichander et al. 2021 Kumar et al.] (0022) Consequently, some critiques motivate adopting ceermaationtheorretee lenses (Voita & Titov, 2020; Diego-Simón et al., 2025). Similar some of these wetemmaaionnthhooeeee approaches (Tishby & Zaslavsky [205]] (Swwartzr Ziv & Tishby| 2017, Skean al., 2025) Voita et al.] 22019), our method makes use of information theory. For a more
detailed discussion of related work, please refer to App.
2
METHOD
Given prefix sequence of length T, X1:7 [z1, (2., IT). the decoder-only MM, generates subsequent tokens {ōr++ T = 1,2,... autoregressively. We denote the hidden activation at position & after the 1th transformer layer by ht E R where d is the hidden dimension of the model.
In our analysis, we treat the outputs h of individual transformer blocks as the fundamental units of computation We posit that block's output is informative about the computation performed
within that block (Pal et al.] 2023 Shai et al.] 2024 Li et al., 2023 Elhage et al., 2021; Lind- sey et al. 2025). We use them as a proxy for the model's internal processing since each block's
computation contributes to the overall function of the model. To analyze how planning emerges in the model's internal dynamics, we examine collections of layer-token indices S and aggregate the corresponding block outputs as Gs (h. tf) E S), representation for the computations within its block. While Gs provides fine- -raiira representation of block computations, it often contains unnecessary details for our planning analysis. Hidden states may encode small variations due to token positions, normalization shifts, or attention noise, factors that are critical for exact out- put reconstruction but irrelevant for studying planning structure. Moreover, MI estimates between such high-dimensional vectors have high variance and are hard to find. To address both issues, construct coarse, discrete representations of Gs. These abstractions preserve the structural de- pendencies required for our analysis while filtering out irrelevant micro-level variability, yielding more stable and interpretable MI estimates. We depict the proposed method in Fig. [ In the first step, we train a model that maps each collection Gs to a coarse discrete representation Zs- Using the trained model, we then obtain dataset-wide codes for the target hidden states and use these to estimate mutual information between the coarse representations of LM's internal computations.
Measuring mutual information between coarse representations. To investigate how planning is realized within the model, we use MI as principled way to quantify relationships between the coarse representations of different computations, corresponding to Steps 2 and 3 of Fig. [ Suppose focus on two hidden state blocks, G. and GB. which random variables depending on the data distribution and the trained model parameters. We denote their coarse, discrete representations by ZA E A and Zn E ZB where Z and Eg are the finite alphabets of possible codes. For realizations Za E ZA and 26 E ZB. P(za: 26), p(za), and p(z6) refer to the joint probability of ZA and Zg, and the marginal probabilities of ZA and ZB, respectively. Mutual information, which measures how much knowing Z. reduces uncertainty about 7B (and vice versa), is defined as,
1 AZ ZB) > Emsz,,
EZA 20 p(za: log p(za)p(zs)
p(Za.
(I)
Since the true data distribution is inaccessible, we approximate by replacing p(za). p(as), and p(ze: 26) in Eq. I with their empirical estimates from the datasets, obtaining by counting code pairs as in Step 2 of Fig. We analyze the finite- sample error of this bount-based MI estimator (and bias-corrected variant) in App. E We do not interpret MI as a direct measure of planning; rather, use MI comparatively. We further define normalized mutual information (nMI) metric as
((ZA:ZB)
AMI(Z.:: ZB)
Imax
where Imax denotes the maximum value across the set of MI calculations performed in the same ex- perimental analysis for comparison. For our purpose of LM planning investigation, the nMI metric
We use the terms "hidden state" and "(hidden) activation" interchangeably to refer to these outputs.
4

---

Published as a conference paper at ICLR 2026
Step Train VQ-VAEs
Step Save (Z, 2n) statistics
Language Model
Prefix
G
Output
U
Encoder
Discretization
Z
Z
Gg
Encoder
Decoder
G
Gs.
Encoder
Discretication
Codebook vector
Step 3: Calculate the mutual information, (z) using the statistics
The proposed method. Step 1 (training): For frozen LM M. hidden states from selected transformer blocks Gy are passed through a VQ- encoder, which maps each block to latent vector and then to a discrete codebook index Zs [K], providing coarse summaries of internal computations. Step 2 (analysis): For two sets of hidden-state blocks, Gs, and Gsg we apply the trained encoder and codebook to the dataset to obtain discrete variables ZA and Zg and collect their empirical co-occurrence counts. Step 3: Using these statistics, we estimate joint and marginal distributions p(za, 26); p(za), and p(za) to compute mutual information I(Z ZB) (and its normalized variant). Then, we analyze how information is shared between different components of the model's computation.
Figure 1:
sufficient since analyses depend relative, rather than absolute, comparisons between of computations. Accordingly, all claims are about within-setting trends, not about nMI as an absolute planning score. Absolute values can be misleading for three reasons: (i) the estimated MI of coarse representations is bounded above by the logarithm of the alphabet size, making it inherently de- pendent on the experimental choice; (ii) any coarse-grasd coarse-graining graining necessarily discards some information, which is advantageous for our setting as discussed but removes any meaning of absolute MI values related the original hidden states; and (iii) answering our three research questions ($ requires only relative comparisons of MI instead of relying absolute values.
Acquiring coarse representations. To estimate MI between the high-level representations of hid- den state blocks, Gs, we need to find a mapping from Gs to the discrete codes, as depicted in Step 1 of Fig. [ The size of target index set S and its corresponding activation outputs G may vary depending on the analysis. For instance, when examining all block outputs across the prefix tokens, S spans every layer and grows with the prefix length T. Since neural networks effective at learning representations, we employ an encoder-decoder neural network architecture: the encoder maps the potentially long and variable-sized Gs to compact representation, and the decoder aims to recover G5 from it. Leveraging the effectiveness in handling variable-length inputs and learning rich representations (Lin et al.) 2022), we adopt transformer-bas based encoder and decoder models. Specifically, we employ Vector- Quabtizeuadt Variational Autoencoder (VQ-VAEs; Van Den Oord et al. 2017). VQ-VAE framework consisting of encoder-decoder network with trainable embed- ding codebook between them to discretize the encoder's output.
Coneretely, for each family of index sets S used in our analyses, we train separate single VQ-V24 1 consisting of an encoder E, a codebook {ex J jkm1 C RRd and a decoder D. Given block Gs. the encoder produces a latent vector TS = E(Gs) E IRd,, We then quantize by nearest neighbor in the codebook.
kt arg min ||rs ex 11 KE[K]
Zs = k*
is Ck
and reconstruct GS D(fs). During the training of VQ-VAE model, gradients are propagated through the quantization step using the ttraight-through estimator (Van Den Oord et al. 2017). The overall training objective combines reconstruction, standard pector-uantizztion, and two regular- ization terms that promote informative and well-used codes:
= Lres + Xy frg + Acoe Less + >emt Lest-
In Eq. Lrs is a reconstruction loss that encourages Gg to match Gs. Lvg is the quantization and commitment loss in Van Den Oord et al. 2017) tying encoder outputs to discrete codes. In
5

---

Published as a conference paper at ICLR 2026
addition to vanilla VO-VAE's training objective we introduce Leos, cosine-similarity penalty that pushes codebook embeddings {ex} to be diverse, and Leu.. an entropy term that discourages collapse to small subset of codes. These losses encourage representations assigned to different codes to diverge so that they are not only compact but also discriminative?? This enables embeddings to spread out in the latent space and ensures that codes capture distinet features of the data making them informative and discriminative. Details about the formulas and implementation are given in App. A.I and App. A.2
Overview of the pipeline. Fig. I summarizes the procedure. For frozen pretrained LM M, we run prefixes X1:T, select block index sets S, and extract the corresponding hidden-state blocks Gs- For each family of blocks, we then train a VQ-VAE so that its encoder maps each Gs to a discrete code Zs E []]. where K' is the codebook size (Step 1). Instead of fine-grained activations, we work with these high-level representation codes that serve as compact summary of complex internal computations in G. Importantly, we treat these codes not as semantically meaningful entities, but representational summaries that allow to find information relations. Therefore, the estimated mutual information depends not only on the underlying relationship between the hidden state them selves, but also how they are encoded to high-level codes. Our use of VQ-VAE ensures that the learned representations are nontrivial since the decoder achieves meaningful reconstructions, and the cosine similarity penalty enforces diversity among codebook embeddings, which is validated by the reconstruction quality and codebook similarity results presented in App. [A.] A In Step 2, we apply the trained encoder and codebook to the full dataset to obtain codes ZA and Zp for two com- pared block collections GA and GB. count their co-occurrences. Then, we using the empirical joint and marginal distributions, in Step 3, we find AZ ZB) and its normalized variant nMI (Z ZB), used as measures of how information is shared between different parts of the model's computa- tion. Furthermore, we validate that our framework can find closer relation a known underlying distribution, in representative experimental setting. Due to space limits, we refer the reader to App. A.4 for the details and results of the validation experiment. Experiments discussing probing- based method (Xu et al.) and a diagnostic illustrating why probing have confounding effects are presented in App. D.5
While framework is general and could be applied to other applications in deep learning, in this work we use to analyze how transformer models implement planning in practice, focusing on dimensions that both theoretically meaningful and empirically informative.
3 ANALYSIS OF PLANNING CAPABILITY
We leverage our method to investigate two core planning-related aspects 33. how much model's prefix computation forward-lood lookid--ooinggggg about future tokens, and (3.2) whether preserves alternative correct continuations- and, as application of the same framework, (3.3) how relleat information is encoded across layers and earlier prefix blocks when producing next-token decisions. In our analyses, use cechitecturr-mmtcche GPT-3 Small (Brown 2020) models employing rotary position embeddings (Su et al. 2024). We report results from multiple seeds.
Training objectives: next-token vs. multi-token. In our analysis, we examine two variants of LM training objectives: MNNP is trained with the standard next-token prediction (NTP) loss (LNTP), which optimizes conditional likelihood of the immediate next token, and MMTP is trained with the multi-token prediction (MTP) loss (LMTP), which encourages models to align their hidden states with predictions across multiple forthcoming tokens. Following Glocekle al. (2024) for MTP implementation, we use shared transformer layers and separate heads for the next tokens during training. During inference, both models use standard autoregressive next-token generation. Intu- itively, NTP prioritizes local boken-by-token consistency, whereas MTP incentivizes models to form short-horizon plans and reduces strictly myopic behavior (Nagarajan et al. 2025 Bachmann Na-
garajan 2024). Both losses are described in detail in App. B.l
Datasets. Our experiments span three datasets chosen to cover distinct aspects of planning. (i) A cortext-free grammar (CFG) dataset (Allen-Zhu, 2024 emphasizes token-level rules and local
PIn short, our VQ-VAR setting with the additional cosine loss combines the regularizing and contrastive
approaches for representation learning, e.g., (Chen et al. 2020 Gao et al. 202T

---

= {M t= 1,...,,7 T: t= 1,... I.,,,--- where H encodes the entirety of the model's prefix computation that has effect future ken generations. Training separate VQ-VAE VAE over H. we get the LM's pre-output computation Published as a conference paper at ICLR 2026
844 84+2
Res ... Ryreet !*
Layer Number
L-1
23
14
22
Xq
X741
Block Number
Solution
Sturt node
Solution
NI 1th layer hidden state at
x prefix input
f model outpur
Goal node
Other
Decay
Figure 2: The visualizations of experimental settings. Left: Horizon of the plan and branches in
the plan experiment ($ 3.1 $ 3.2). Middle: Information in the computational history experiment (.3... Colors in the left and middle panels denote the target variables. Right: An illustration of simplified sample in PF task. Graphs and token numbers are randomly generated except for the start (1) and goal node (2). sample prompt is "19 23, 13 2, 11 8 4 17, 23 9 2. 8, 17 132, 14 23, 23 7, 23 :" and correct responses are "1 11 8 2" or " 17 13
coherence. Tokens in CFG carry intrinsic meaning, much like words in natural language, and are governed by both local and global syntactic constraints. The model's task is to learn the grammar and predict valid next tokens that obey its rules. As controlled substitute for natural language, CFG allows us to adjust vocabulary size and task difficulty, providing a clean environment for scientific study. (ii) A path-finding (PF) task (see Figure right for illustration) requires models generate valid paths between a fixed start and goal node in graphs with varying sizes and edge structures. The node vocabulary consists of 28 unique tokens, each representing node identity. Unlike in CFG, these tokens do not carry intrinsic semantics and the meaning of a node token arises only through the of edges it participates in. To construct language-model inputs, we linearize the graph by writing edges as pairs of adjacent node tokens separated by commas in a random order to prevent the model from exploiting positional shorteuts (e.g.. "u v, 1, ....). "). This edge-list representation ensures that the model must infer connectivity and graph structure rather than rely on token identity alone. Solving PF thus requires composing multiple reasoning steps over these relational inputs, analogous to chaining lemmas form proof, introducing natural flavor of long-hon hong-horizon planning. We construct two variants, PF-Short and PF-Long, which require finding paths of length 4 and 6 (including start and goal), respectively. (iii) natural language dataset, Teb (Gokaslan et al., 2019), reflects real-world distributional modeling. Here the model predicts continuations consistent with natural text. These datasets cover a spectrum from low-level symbolic structure to high-level naturalistic text. We refer the reader to App. B.2 for details.
3.1 HORIZON OF THE PLAN
In autoregressive generation, predictions for tokens beyond the immediate next one depend on both the prefix and the model's own generated outputs, which raises the question of how much computa- tion the model allocates to tokens beyond the immediate prediction. We specifically ask:
How much do LMs plan about future tokens before deciding on the immediate next token?
To answer this question, we analyze two sets of activations. First, for a prefix of length T, we take the final-layer hidden states of the autoregressively generated tokens, h for 2 0. As h fully specifies the distribution over ZT+7+1, we treat it as the model's decision state and use its code 24_, concise summary of the strategy employed for that token, obtained from trained VO-VAE. Second, we collect all prefix activations across layers:
1-1) 1-...,I-111
E TT(L-11)xd
summary, and denote its high-level code as Z1:L-1 H and h;_,, are respectively shown in red and blue colors in Figure 2 (left). We address our question by comparing how much the prefix

---

Published as a conference paper at ICLR 2026
computation's s summary Z1:f-1 tells about the model's decision state at the generated tokens.
We measure the nMI between 21:7-1 and -/+7 for 720 as defined in Eq. 2 with normalization Imax = ( 1:T' (Z): Z/+n)) hn If this ratio remains significantly above zero for large т, then the pre-output computation Z1:7- carries information on the strategy to produce ÖT+r, indicating the model's initial prefix processing is not merely myopic but encodes forward-looking structure that persists into later generations. Conversely, if the ratio quickly decays to zero, then prefix com- putations primarily support only the immediate next prediction, with little evidence of long- horizon dependence between prefix computation and later decision states.
Context-free grammar results. We generate sentences of length [16, 67] from CFG rules, and for each sample select prefix length I E [8, 24] uniformly at random before, with the model generating the remaining continuation. For MNNP and MMIP, Figure reports the nMI between 2 and Z/et across T 2 0. Because CFG tokens have intrinsic meaning and are significantly governed by local syntactic constraints, it is expected that prefix computations correlate most strongly with the first few generated tokens. Indeed, Z1:7-1 retains the highest information about the decision state for r = 0, and by -- = 10 the nMI drops to roughly one- one-fifth of its initial value. This indicates that the pre-output computation encodes short-horizon plan tied mainly to the next few tokens. We observe the same rapid decay of nMI when scaling the LM to 0.3B parameters (App. D.I MMTP exhibits a similar pattern, with a slightly slower decay for T < 10, suggesting that the MTP loss does not substantially extend the horizon of pre-output computation in model well- trained on CFG.
Path finding results. On PF-Short and PF-Long, we trained models on NTP and MTP losses. Because the start and end nodes are fixed tokens, the challenge lies in predicting the correct interme- diate nodes (T = 1, 2 for PF-Short and r 1, 3,4 4 for PF-Long) that connect them. We therefore focus our analysis on the intermediate positions and compute nMI as defined in Eq. Figure 3b shows nMI across the generated path as well as the accuracy of correctly finding the whole path. For PF-Short, we see similar nMI trends across generated tokens for both MMTP and MNTP, which both attain high accuracy. Interestingly, prefix computations encode more information about the second intermediate node than about the first. A likely explanation is that there is less uncertainty about the final token given the edge information and that the model spares more pre-output computation to find the last token. This resembles strategy work backwards from the goal, as human solving this task might do. Similarly, for PF-Long, in (ig both MNNT and .MMIP prefix computations encode a comparatively high amount of MI about the future tokens, unlike the steady decay seen in CFG. This suggests that pre-output computations are not limited to the next token but also embed plans for upcoming positions, reflecting deliberate allocation of capacity toward future outputs. This finding is aligned with previous research on training LMs to exploit this reverse-solving approach, which has been shown improve performance (Bachmann & Nagarajan 2024). Taken together, these results indicate that LMs trained with both NTP and MTP loss can exhibit non-myopic behav- ior when trained on tasks that demand it. Lastly, in the PF-Long experiment, we observe that the nMI of MMTP is slightly more uniform across 7 than that of MNTP. which is likely explanation for its higher accuracy by better predicting the earlier (harder) tokens along the generated path.
CFG Task: Normalized MT vs. T
PF-SShort: aMI vs.
PF-Long: nMI vs.
NTP
MIP Less
0.5
0.5
BMI
0.5
0.4
0.4
0.4
0.- NTP Loss 0.92
02-
Loss (Acc:0.50)
0.0
-2-
MTP Losk (Anc 0.87)
0.0
-2-
MTP Loss (Ac::0.85)
(a) Context-free grammar
(b) Path finding
Figure 3: nMI results between the prefix summary codes and the last hidden state codes of generated tokens for CFG (a) and PF (b) tasks. nMI decays fast in CFG, consistent with short-range depen- dence, while PF maintains or even increases nMI beyond r = 1, consistent with longer-horizon predictive dependence of prefix computation on later decision states.
8

---

Published as a conference paper at ICLR 2026
Table I: Mean + std for I(ZH; / Edecay) metric and accuracy values branches in the plan experiment ($ (.2... Models encode information about unchosen correct branches more strongly than unrelated decoys, indicating branch awareness prefix computations.
PF-Short
PF-Long
Model
IiZziZal 7(ZH (rress)
Accuracy
I(ZziZa))
Accuracy
T(Z.H:
MNTP
7.60 + 0.78
0.92 + 0.03
1.45 + 0.01
0.60 + 0.01
MMTP
6.29 + 0.17
0.88 + 0.05
1.82 + 0.27
0.85 + 0.02
Overall, the dependency between pre-output summaries Z 1:T and decision states Z confirms that the planning horizon is task-dependent. On CFG, nMI drops quickly with T (short, local plan- ning); PF, it stays high and can peak beyond T = 1, allocating compute to later steps.
3.2 BRANCHES IN THE PLAN
In many tasks, there can be multiple correct continuations for a given query. For exam- ple, a math problem may admit different solution strategies, may allow several equally valid completions. This raises the question of whether LMs, when producing one correct answer, also implicitly encode information about other possible cor- rect answers. Specifically, for the subset of samples the model solves correctly, we ask:
natural language prompt
Does the model consider alternative correct answers when generating the next token?
To investigate this question, we use the PF datasets, where each sample is constructed to have two correct paths and one decoy path (see App. B.2.2). By design, the correct paths and the decoy path do not share any common nodes, ensuring no correlation among them. Similar to 3.1 we summarize prefix computations H into high-level codes Z
1:T I: using a VO-VAE encoder. In addition, we train a separate VQ-VAE to encode an entire path to find codes for the alternative correct path (Zatt) and the decoy incorrect path (Z H and the generated path are visualized in Figure 2 (left) with red and green color, respectively. We then measure the MI between the prefix and alternative solution versus that between the prefix and decoy path through the ratio I( I(Z 1:: Z I ( [Z ZZeccry): A ratio above one means that the prefix computations encode more information about the alternative path than the decoy, indicating branch awareness pre-output planning.
1:T
Table reports this ratio together with path-finding accuracy. On PF-Short, both MNTP and MMTP yield ratios well above and high accuracy, showing that the prefix retains information about unused correct branches on easier instances. On PF-Long, the ratios are smaller yet remain above which indicates attenuated but persistent branch information as path length grows. .MMP achieves both higher accuracy and a larger ratio than MNIP, suggesting that model which solves the task more reliably also maintains richer branch-aware computation in its prefix states. Because the correct and decoy paths disjoint, trivial overlap is ruled out explanation.
nMI over Blocks and Layers for Generated Tokens at r and
Conditional nMI
0.57
0.57
0.56
0.57
1.00
0.57
0.56
0.56
0.58
0.70
7
0.30
6
0.53
0.33
0.53
0.53
0.87
6
0.53
0.52
0.54
0.53
0.66
6
0.31
Layer Number
0.45
0.44
0:45
0.31
0.38
0.38
0.38
0:38
0.85
0.38
0.38
0.38
0.38
0.53
4
0.32
3-
0.38
0.36
0.36
3
0.38
038
0.53
3
0.32
0.35
0.35
0.34
0.35
0.36
0.30
0.34
0.34
0.33
0.34
0.41
0.34
0.34
0.33
0.34
0.37
1
0.04
4
3
2
5
4
3
Block Number (k)
Block Number (k)
Figure 4: nMI across blocks and layers, and conditional nMI. Left: nMI between the hidden state block codes and the token decision state code at 0. (B
(B;; 24). Middle: nMI between
block codes and the last-layer decision code at 7 1, mMI Zf+1). In both heatmaps, nMI is higher for recent blocks (small k) and final layers (high E). Right: Conditional nMI for the 1st block, nMI(Zf (ff(((t--11 24 Z7.) showing that most the dependence between the 1st block and the generated token at 0 is attributable to the final prefix position T.

---

Published as a conference paper at ICLR 2026
3.3 INFORMATION IN THE COMPUTATIONAL HISTORY
In an autoregressive LM, the final hidden state at the end of the prefix determines the next token, yet the computation producing it spans layers and earlier positions. This raises a natural question: Which earlier layers and prefix blocks remain informative about the next-token decision state?
We view this primarily as a diagnostic. It does not tell how improve the model, but it quan- tifies how concentrated next-token decision information is in recent vs. earlier computation, which can help design choices that alter attention span or depth- wise context allocation. To answer the question, we use an 8-layer decoder-only Transformer trained on WpenWe with NTP objec- tive. We estimate the nMI between the codes of the model's 's decision state for the generated tokens and of blocks of prefix hidden state computations within the LM. Similar to $ 3.1 we obtain summary of the strategy employed for token Tr+r+1, denoted 24+- with a vQ- VAE trained over the last layer of hidden states. We also partition each layer's prefix hidden states into contiguous, non-overlapping blocks of length 16 and train a VQ-VAE to acquire the code B where the block index E {I...., }2} denotes the ht block from the end of the prefix (i.e., k = corresponds to the most recent 16 time steps before generation). In Figure 2 (middle), sample block and sample last hidden state are illustrated. We quantify the dependence between the codes of each computa- tion block and the decision state by measuring nMI between B and Z4+, across all layers é and blocks k as defined in Eq. [ with Imax= = max (B
K Z4 24 ( k.km. Heatmaps in Figure (left and middle) show the results for 7= 0 and 1 (see App. B for T >>> 1).
We observe that last layer computations retain the most information about the decision state of both immediate and future tokens which is aligned with the results of Pal et al. (2023) and consistent with the design choice of assigning longer attention spans to higher layers (Sukhbaatar et al. 2019, observe clear recency effect across all layers: the codes of the most recent blocks, i.c., final time steps of the prefix, computations (21), exhibit the highest nMI (Eq. 2 with the decision state of the generated tokens, and the nMI decays over earlier blocks in the prefi. Although analysis along both axes indicates that the LM primarily relies on the most recent computations, we still observe appreciable nMI in lower layers (small E) and earlier blocks (large k). This suggests that LMs retain information from earlier parts of the prefix when generating new tokens, instead of relying only on the most recent computations.
Beltagy et al. 2020). Furthermore, along the block axis,
We also define conditional normalized mutual information, nMI App. B.5.2 for details), where Z 15:T5-1-1 denotes the codes from layer E spanning prefix posi- tions from 15 to 7 This quantity, reported Figure4 measures how much additional information about the decision state Z4 is captured by earlier time steps in the 1st block, beyond what is already contained in the final prefi position Z The resulting values are approximately 0.3, which is much lower than the unconditional nMI values of the 1st block observed in the left heatmap. This gap suggests that the majority of the dependency between the 1st block and the decision state arises from the final prefix token at position T. This finding further reinforces our earlier conclusion: the LM primarily relies on the most recent computations when generating new tokens.
Z4 Z4.)) (see
4 CONCLUSION
Drawing on VQ-VAE, we develop an information -eerrration--hhoretiicc pipeline that compresses hidden-state trajectories into discrete codes and uses them to compute MI across an LM's computation. This lens measures (i) how prefix computations inform future decision states, (ii) whether models retain in- formation about alternative continuations, and (iii) how decision-relevant information is distributed across layers and prefix blocks. We evaluate models trained with standard NTP versus MTP losses on a synthetic CFG task, PF tasks, and natural text; we find no consistent differences between NTP- and MTP-trained models. The results are strongly dask-dependent: in CFG, nMI between the prefix and the decision state decays quickly, consistent with largely myopic computation; in path-finding the prefix retains substantial information about later steps and alternative correct paths, indicating stronger branch awareness. Across settings, we observe recency effect, decision states depend most on late layers and recent blocks, while earlier activations remain measurably informative. Overall, LMs exhibit internal planning, but its extent and form vary with the task and training objective; our VO-VAE-M framework provides an automated way to study these behaviors. Future work could extend this analysis to reasoning models and test architectural changes that promote planning.
10

---

Published as a conference paper at ICLR 2026
Acknowledgements. This work was partially supported by the US National Science Foundation under grants CCF 2045694, CNS-2112471, CPS-2111751, ONR 100014-23-1-2149 to GJ and US Department of Energy under grant DESC0025652 to CJW. In addition, the work supported by Pennsylvania Infrastructure Technology Alliance, NSF Grants 2154171, CAREER Award 2339112 12, NSF Award 251 2512805, CMU CyL.ab Seed Funding GQ. This work used PSC Bridges-2 GPU at Pittsburgh Supercomputing Center through allocation CIS250226 from the Advanced Cyberinfras- tructure Coordination Ecosystem: Services Support (ACCESS) program, which supported by US National Science Foundation grants #2138259 #2138286, #2138307, #2137603, and #2138296.
Reproducibility statement. We describe our method relevant appendices. Our experimental analyses presented in $ with additional details likewise included the appendices. To ensure reproducibility, we provide the full code, tools to acquire the datasets, and instructions in the supplementary material. In addition, include README file that explains how to run each experiment separately.
2 with further details provided in the
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08777 2023.
Jihyun Janice Ahn and Wenpeng Yin. Prompt-reverse inconsistency: LLM eeff-inconsistence be- yond generative randomness and prompt paraphrasing. In Second Conference on Language Mod- eling, 2025. Uttphhtss://openeeeet
net mett/orrmPPid=yfRkKe.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614-456500 2023.
Zeyuan Allen-Zhu. ICML 2024 Tutorial: Physics of Language Models, July 2024. Project page:
htttps://physicee allen zhu. com/
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Forty-first International Conference on Machine Learning, 2024. URL httpss://penrrevier net/
forum?1d=762q8Wkk
Iz Beltagy, Matthe E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint 2rXvv:2004.05150, 2020.
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gian- inazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. In Proceedings of the AAAI conference artificial intelligence, volume 38, pp. 17682 17690, 2024.
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Con- crly, Nick Tumner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https:/ttransformer-
Tom et al. Brown. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
Lawrence Chan, Adrià Garriga-Alonso, Nicholas Goldwosky-Dill, Ryan Greenblatt, Jenny Nitishinskaya, Ansh Radhakrishnan. Buck Shlegeris, and Nate Thomas. Causal scrub- bing, method for rigorously testing interpretability hypotheses. Al Aligument Fo- rum, 2022. http:://wwaaaaaaaaaa.....ooooooa causal oeg.

---

Published as a conference paper at ICLR 2026
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Hal Daumé III and Aarti Singh (eds.), Pro- ceedings of the 37th International Conference Machine Learning. volume 119 of Proceed- ings of Machine Learning Research, pp. 1597-1607. PMLR, 13-18 Jul 2020. URL https:
//proceeding ngs. ml www.r.pprss//se / 19/ httpaa///vv11//eenn.n.11 c.tp:0eenn22nnmmm
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen- coders find highly interpretable features in language models. arXiv preprint arXiv:2309.086MM) 2023.
Pablo J Diego-Simón, Simón, Emmanuel Chemla, Jean-Rémi King, and Yair Lakretz. Probing syntax in large language models: Successes and remaining challenges. arXiv preprint 03X2:12508.03211, 2025.
Jacob Dunefsky, Philippe Chlenski, and Neel Nanda. Transcoders find interpretable Ilm feature circuits. Advances in Neural Information Processing Systems, 37:24375-24410, 2024.
Nelson Elhage, Neel Nanda, Catherine Olsson, and et al. A mathematical frame- work for transformer circuits. Transformer Circuits Thread, 2021. https://trnsformmrr
Cheng Gao, Yuan Cao, Zihao Li, Yihan He, Mengdi Wang, Han Liu, Jason Matthew Klusowski, and Jianqing Fan. Global convergence in training large-scale transformers. The Thirty- eighth Annual Conference on Neural Information Processing Systems, 2024. URL https:
Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE: Simple contrastive learning of sen- tence embeddings. In Marie-Francing Moens, Xuanjing Huang. Lucia Specia, and Scott Wen- tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 6894-6910, Online and Punta Cana, Dominican Republic, November Associatin for Computational Linguisties. doi: 10. 10.11865//1222221emmml 1wtp:2000mmlpmmmmii.55222222 https::/aclanthholoo wwtp://...coooo.o.ooo........ emtppll.mmmmmmmmaan 552/
wemmlpmmain.5555 m5in.552. URL
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Syn- naeve. Better & faster large language models via multi-token prediction. arXiv preprint arXi::2404.19737. 2024.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. 7/5kylion001 gi wti.ttu..oo/ooooo wtt..b...ooopeeeeeeettttooooooo
http:
2019.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, and et al. The llama 3 herd of models, 2024.
hollhttss
abs/2407 2271 21783
Wes Gurnee and Max Tegmark. Language models represent space and time. In The Twelfth Interna- n
tional Conference on Learning Representationss 2024. URL
Eorum?id=
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603. 2019.
Dean S Hazineh, Zechen Zhang, and Jeffery Chiu, Linear latent world models in simple transform- ers: A case study on othello-gpt -ppt. arXiv preprint arXiv:2310.07582 2023.
John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Confer- Natural Language Processing (EMNLP-IJCNLP),, pp. 2733-27433 Hong Kong, China, November 2019. Association for Computational Linguistics. doi: (0.18653/v1/D199 1275. URL https: antho -ot.
12

---

Published as a conference paper at ICLR 2026
Abhinav Kumar, Chenhao Tan, and Amit Sharma. Probing classifiers unreliable for concept removal and detection. Advances in Neural Information Processing Systems, 35:17994-18008, 2022.
Jenny Kunz and Marco Kuhlmann. Classifier probes may just learn from linear context features. In
Donia Scott, Nuria Bel, and Chengqing Zong (eds.), Proceedings of the 28th International Confer- ence on Computational Linguistics. pp. 5136-5 146, Barcelona, Spain (Online), December 2020. International Committee on Computational Linguistics. doi: 10. 10.18653/11/2000oooongg m50 wwt. URL https: http://aaaaaathooooooooo wttp:///oooorrg.20000220/
450/
Yann LeCun. path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open
Review, 62(1):1-62, 2022.
Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wat- tenberg. Emergent world representations: Exploring sequence model trained on a synthetic task. In The Eleventh International Conference Learning Representations, 2023. URL https:
opwprevieeewe..e/////e httpps://ppeneeiieee
net/ htt////orumddddeeeeeee Tc2vT
Perey Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. arXiv preprint arXiv:2211.09) 2.09110,,,11,,, 2022.
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transform- ers. Al Open, 3:111-1331 2022. ISSN 2666-6510. doi: http:///b..rr///..000//.a.oo 2022010.001. URL https::/nat http://wwwww/.. http::::/wwww.siienee.m www.sciieneeirect..o.o cww comm/scleence/atiicll//..i $2666651022000146
Zhenru Lin, Jiawen Tao, Yang Yuan, and Andrew Chi-Chih Yao. Existing llms are not self-consistent for simple tasks. arXiv preprint arXiv:2506.1877 2025.
Jack Lindsey, Wes Gurnee. Emmanuel Ameisen, Brian Chen, Adam Pearce, Nicholas L. Turner, Craig Citro, David Abrahams, Shan Carter, Basil Hosmer, Jonathan Marcus, Michael Sklar, Adly Templeton, Trenton Bricken, Callum McDougall, Hoagy Cunningham, Thomas Henighan, Adam Jermyn, Andy Jones, Andrew Persic, Zhenyi Qi, T. Ben Thompson, Sam Zimmerman, Kelley Rivoire, Thomas Conerly, Chris Olah, and Joshua Batson. On the biology of large language model. Transformer Circuits Thread, 2025. URL httpp:://taansoorrrrre pub/2025/atee http://ti....iiiooooogggp/////aa
ci riuits.
html
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representa- tions, 2023. URL http:::tt///cc hattps:://penreviiie
net m 4F
David Q Mayne, James B Rawlings, Christopher V Rao, and Pierre OM Scokaert. Constrained model predictive control: Stability and optimality, Automatica, 36()) 814, 2000.
Kevin Meng, David Bau, Alex J Andonian, and Yonatan Belinkov. Locating and editing factual asso- ciations in GPT. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL htttts:///ppeeeeeeeeee net//forum? dimmmiidi==-- hee.rrmmmemd---hhhee6e hWAAS6F44
Ida Momennejad, Hosein Hasanbeig, Felipe Vieira Frujeri, Hiteshi Sharma, Nebojsa Jojic, Hamid Palangi, Robert Ness, and Jonathan Larson. Evaluating cognitive maps and planning large language models with cogeval. In Thirty-sevnnth Conference on Neural Information Processing Systems, 2023. URL httt:ttt////peneeeeeee
fort/forum? wtt.:ummmmddieeeeeeeeeeee
Manfred Morari and Jay Lee. Model predictive control: past, present and future. Computers & Chemical Engineering, 23:667-682, 1999. htt.::ttp::s/attt semant wommmnttcccchlaa... org/
(CorpusD::192677)
Vaishnavh Nagarajan, Chen Henry Wu, Charles Ding, and Aditi Raghunathan. Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction. In Forty-second International Conference on Machine Learning, 2025. URL https: httpa//////peeeeeee..eeee
forum? d= 0SyHMmkd.
13

---

Published as a conference paper at ICLR 2026
Catherine al. Olsson. In-context learning and induction heads. Transformer Circuits Thread, 2022. https:/transforme-cinnottt
anduction-heads/nne html.
Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: Antici- pating subsequent tokens from single hidden state. In Jing Jiang, David Reitter, and Shumin Deng (eds.), Proceedings of the 27th Conference on Computational Natural Language Learn- ing (CoNLL), pp. 548-560, Singapore, December 2023. Association for Computational Lin- guistics. doi: 10. 11211///1/2223..oo 1.37. URL con 1-1111... 1
Liam Paninski. Estimation of entropy and mutual information. Neural computation, 15(6):1191- 1253, 2003.
Stefano Panzeri and Alessandro Treves. Analytical estimates of limited sampling biases in different information measures. Network: Computation in neural systems, 7(1):87, 1996.
Tiago Pimentel, Josef Valvoda, Rowan Hall Maudslay, Ran Zmigrod, Adina Williams, and Ryan Cotterell. cnermatiionthhoretiii probing for linguistic structure. Dan Jurafsky, Joyce Chai, Na- talie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 4609-4622, Online, July 2020. Association for Computational Linguistics. doi: 10. 10.18653/ 10/55/3/112222cccc/ awl--aaclmmm main-
URL https: httpp:::/aaclaathoooogooo orq/
a020 mainn 420/
Abhilasha Ravichander, Yonatan Belinkov, and Eduard Hovy. Probing the probing paradigm: Does probing accuracy entail task relevance? In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pp. 3363-3377, 2021.
Jonathan Richens, Tom Everitt, and David Abel. General agents need world models. s. In Forty-second International Conference on Machine Leurning, 2025. URL htttss://oppereeiiiiiiie
Forum? d=d) TownNi Xt.
Yash Saxena, Sarthak Chopra, and Arunendra Mani Tripathi. Evaluating consistency and reasoning capabilities large language models. In 2024 Second International Conference on Data Science and Information System (ICDSIS), pp. 1-5. IEEE, 2024.
Bilgehan Sel, Ruoxi Jia, and Ming Jin. LLMs can plan only if we tell them. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //pp///vvieeeeeeeeee
net/forum? (o
Adam Shai, Paul M. Riechers, Lucas Teixcira, Alexander Gictelink Oldenziel, and Sarah Marzen. Transformers represent belief' state geometry in their residual stream. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL htLltttpss///peeeeeee
Sett/Eorumnid=yIBB
Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, and Yueting Zhuang, Hugging- gpt: Solving ai tasks with chatgpt and its friends in hugging face. Advances in Neural Information Processing Systems, 36:381 154 38180, 2023.
Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via informa- tion. arXiv preprint arXiv: 1703: 000100 2017.
Oscar Skean, Md Rifat Arefin, Dan Zhao, Niket Nikul Patel, Jalal Naghiyev, Yann LeCun, and Ravid Shwartz-Ziv. Layer by layer: Uncovering hidden representations language models. In Forty- second International Conference Machine Learning, 2025. URL httl:httss/:/ppeeevvii
net net/Eoznn Futt/Eorm??l id WGXb7Udv TU
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, and et al. Beyond the imitation game: Quanti- fying and extrapolating the capabilities of language models. Transactions on Machine Learn- ing Research, 2023. ISSN 2835-8856. URL https:/ httt:::///pperevieew /pwpreeeeeeww.. ww..reviie.rerr net. fo rum? ? j Featured Certification.
Jianlin Su. Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
14

---

Published as a conference paper at ICLR 2026
Sainbayar Sukhbaatar, Édouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In Proceedings of the 57th Anmual Meeting of the Association for Compu- national Linguistics, pp. 331-335, 2019.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pp. 1-5, leee, 2015.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Alessandro Treves and Stefano Panzeri. The upward bias in measures of information derived from limited data samples. Neural Computation, 7(2):399 407, 1995.
Muhammed Ustaomeroglu and Guannan Qu. A theoretical study of (hyper) self-attention through the lens of interactions: Representation, training, generalization. In Forty-second International Conference on Machine Learning, 2025. URL htt.h::p//// https::/openrev ew net. f formmm id= WQVRILHbbD
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
E Voita and Titov. Information-thheortic probing with minimum description length. In EMNLP 2020-2020 Conference on Empirical Methods in Natural Language Processing, Proceedings of the Conference, pp. 183-196, 2020.
Elena Voita, Rico Sennrich, and Ivan Titov. The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Confer- ence on Natural Language Processing (EMNLP-IJCNLP), pp. 4396-4406, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19..4442 URL
Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language mod- els. Transactions on Machine Learning Research, 2024a. ISSN 2835-8856. URL https: hwwp://wwww.eeeeoooonneeeeeeeeee
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations, 2023. URL https:
//openzeview.c wwt./eee/t/eooo for.tt/fooummmm wtw.rummmd.dNossSSS5So..
Siwei Wang. Yifei Shen, Shi Feng, Haoran Sun, Shang-Hua Teng, and Wei Chen. ALPINE: Un- veiling the planning capability of autoregressive learning in language models. In The Thirty- eighth Annual Conference on Neural Information Processing Systems, 2024b. URL https: 7/ //opneevieewwo..mm foret/foorm?
Xinyi Wang, Lucas Caccia, Oleksiy Ostapenko, Xingdi Yuan, William Yang Wang, and Alessandro Sordoni. Guiding language model reasoning with planning tokens. In First Conference on Lan- guage Modeling, 2024c. URLpttps::0 htttp:://openreveeee net/ coooummidiyy 91 FRRVM
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
Wilson Wu, John Xavier Morris, and Lionel Levine. Do language models plan ahead for future tokens? In First Conference on Language Modeling.
Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints, In International Conference on Learning Represen- fations.
15

---

Published as a conference paper at ICLR 2026
Shunyu Yao. Dian Yu, Jeftrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimban. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances in neural information processing systems, 36: 10011889--1122) 2023.
Wenhao Zhao, Qiran Zou, Rushi Shah, and Dianbo Liu. Representation collapsing problems in vector quantization, In Neurips Safe Generative Al Workshop 2024, 2024. URL https://
openreview net forum? Sarpummmd==aoiTTee-
Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models as commonsense knowledge for large-scale task planning. Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
Yongxin Zhu, Bocheng Li, Yifei Xin, Zhihua Xia, and Linli Xu. Addressing reprrssttonnnnn collapse in vector quantized models with one linear layer, 2025. URL http:://ariivoor//aaatst
2411. 02038
A METHOD DETAILS
In this section, we explain the details of our method and implementation.
A.1 VQ-VAE DESIGN AND IMP.EMENTATION
Let hf E Rd denote the hidden activation at token position t after Transformer layer t. We view a computation block as a of layer-token indices S S (1....,TI (I..... 7) and write the corresponding activations
Gs ( (L)
The goal is to map the variable-shaped Gs to discrete code 2s E [K] with vector-quantized variational autoencoder (AE- These codes serve as coarse summaries that represent the com- putation block and make information-theoreii analysis feasible. We measure dependencies between computations using codes, for example A(ZA; ZB) for two blocks GA and GB. as in the main text
$
Architecture overview. We train separate VO-VAEs for target S sets used in our analyses. Each VO-VAE has an encoder that takes Gs and outputs latent vector rs Rd., codebook {ex [e C 124, and a decoder D that reconstructs Gg from the selected codebook vector. Quanti- zation uses nearest neighbors
b** arg min ||rs - - kE[k]
k
is
and the straight-througe estimator for backpropagation is used for the gradient flow in training.
Design rules by input structure. our experiments, we select various S depending on the anal- ysis, described in $ Even within the same experiment, S can vary along the token axis because different samples may have different lengths. We design the encoder E and decoder D according the structure of the hidden state blocks, which span the layer dimension, token index dimension, and hidden state vector dimension. Below we summarize the design specializations of E and D to explain how we handle challenges the structure of the hidden state blocks.
Multiple layers the same token index. For each layer & E L, we first process the sequence { with transformer encoder. The resulting representations are then stacked across the layer axis. An MLP maps the stacked representation from dimension [ down to de in the encoder. In the decoder, starting from fs. follow exact reverse operations, e.g., dimensionality expansion instead dimensionality reduction with MLP.
Multiple token indices at the same layer. We treat {he hET as a sequence and process with a transformer encoder. To handle variable sequence lengths, the encoder either crops the sequence to a fixed window Teme or appends m learned sentinel vectors and uses their outputs to form Ts- In the end, concatenate representations across multiple time indices and apply additional
16

---

Published as a conference paper at ICLR 2026
MLP for further processing to get the encoder output. In the decoder, first upsample the time dimension by mapping is with an MLP to a maximum original sequence of length with d-dimensional vectors, then refine per-token reconstructions with transformer encoder.
No token axis (single vector) or after reducing it as above. When the input does not include the token dimension, when has been reduced as described above, we use only MLPs in both the encoder and decoder. The encoder outputs single summary vector, and the decoder reconstructs the corresponding Gg shape.
A.2 TRAINING AND LOSSES VQ-VAE
Training setup. We train each VQ VAE on hidden state blocks Gs extracted from fixed, pre- trained LM M over datasets used in our analyses. For each mini-batch. we sample sequences X1-T, feed M with imputs to get hidden activations h4, assemble the target blocks Gg according to the experiment. Then, we optimize encoder E, codebook {ex 1 and decoder while keeping frozen. This follows the standard VQ-VAE pipeline (Van Den Oord al. 2017) adapted to variable- shape transformer hidden states.
Objective. Given an input block Gs and encoder output rs E(Gs) E Rd,, quantize by nearest neighbor
k arg min [rs kE[K]
ex
is Ek
and reconstruct Gs = D(ss) with ttraight-through gradients for the quantizer 2017). The loss combines four parts
Van Den Oord et al,
Lree data fidelity
Xy Lrg
Acos Lens
Nent Lent
quantization and commitment
codcbook diversity
anti-collapsc
Reconstruction uses mean-squared error over the entries of Gs.
Lrs = |Gs - G 112-
Quantization and commitment follow Van Den Oord et al. (017) with the straight-through estimator
Lvg [sg[rs] -
- c + 3 l|rs sg[ex+
where sg[:] soopggradient and 3 balances codebook learning and encoder commitment. Codebook
diversity encourages distinct codes by discouraging cosine similarity among embeddings,
(e,, Te:lla ||e
Loos
Σ ifj
K(K -
which spreads codebook vectors and yields more discriminative high-level codes, consistent with our use of codes as discussed main text $ To avoid mode collapse of VQ-VAE (Zhao et al. [204]]
Zhu et al.) 2025), we further uses anti-collapse entropy term which employs soft assignments computed from distances to the full codebook. Let
Ph exp(-al|rs - cx
== exp( -a Irs e,1
then we minimize the negative entropy
Lent - Σ Pk logPk.
H(p)
which promotes higher-entropy assignments during training and discourages mode collapse.
Stabilizing codebook usage. VQ-VAE training can suffer from codebook collapse (Zhao et al., 2024 Zhu et al., 2025). We adopt three complementary heuristics. (i) Dead-code reset: codes that have not been the argmin k* for fixed window are reinitialized to random encoder outputs from the current batch. (ii) Codebook dropout: with probability 0.1 we temporarily mask a random subset of codebook vectors when forming nearest-neigee neighbt distances in the forward pass of training, which encourages redundancy avoidance and better exploration. (iii) Entropy regularization: the Len term above explicitly pushes assignments away from degenerate peaks.
17

---

Published as a conference paper at ICLR 2026
Optimization details. We use the Adam optimizer for training. Inputs every token index and layer index are normalized before encoding.
Evaluation metrics. For quantitative checks of representation quality and nontrivial code usage, report: (i) nRMSE between inputs and reconstructions. Let H denote vectorized view of Gş and H that of Gs. With ZB E [K] the selected code and H D(ez,). ), define
|| H12
nRMSE
|| |112
(ii) Codebook geometry: cosine similarity among codebook vectors {ex ] and its distribution, which reflects diversity. (iii Usage distribution: empirical frequencies of ZH.
Rationale for the objective. The reconstruction term certifies that codes retain task-relevant infor- mation about Gs, the VQ term ties encoder outputs to discrete codes, the cosine penalty promotes discriminative summaries that support our information-theoretic analyses, and the entropy term plus stabilization heuristics sustain broad code usage throughout training.
A.3 VQ-VAE TRAINING SETUP ACROSS EXPERIMENTS
Across all experiments, each VQ-VAE is trained once per dataset/task and representation type. This design ensures that the quantization mechanism remains fixed within each study, avoiding con- founding effects that would arise from retraining VQ-VAE models at different mutual information estimates, In the planning horizon experiments, each underlying dataset is equipped with its own hierarchical pair of VO-VAEs (VO-VAE, and VQ-VAE²) for quantizing prefix hidden-state blocks, along with its own last-layer VQ-VAE. Thus, while the architectural pattern is shared, the mod- els themselves are trained separately for each task. In the branching experiments, again train VQ-VAE, and VQ-VAE₂ for prefix hidden states, together with dedicated path-sequence VQ-VAE. In the history experiments, Text representations. Once trained, all VQ-VAE models are reused for every mutual-information estimation within their respective experimental setting.
train one VQ-VAE for prefix blocks and one for last-layer OpenWeb-
With respect to training budgets, the VQ-VAE models require only moderate optimization steps. In the planning horizon experiments, VO-VAE,1 and VQ-VAE.2 (used for prefix hidden states) are trained for 15,000 optimization steps cach, corresponding to 0.96M sequences in total, while the last-layer VQ-VAE is trained for 10,000 steps (0.32M sequences). The path-finding version of this experiment uses the same budgets. These settings apply uniformly to both next-token and multi- token prediction evaluations. In the branching experiments, VQ-VAE,1 and VQ-VAE,2 (for prefix states) are trained for 20,000 (0.64M sequences) and 15,000 steps (0.48M sequences), respectively, and the path-sequence VQ VAE is trained for 25,000 steps (0.80M sequences). Overall, this experi- ment uses between 0.32M and 0.96M sequences per VO-VAE depending on the role. In the history experiments, we train two VO-VAEs, one for prefix hidden-state blocks and one for last-layer hidden states, each for 10,000 optimization steps, corresponding to 0.64M sequences per model.
Choosing hyperparameters for our VQ-VAE models requires some initial effort, but not because the models are inherently sensitive. Instead, the difficulty arises from the fact that our VQ-VAE objective includes several additional components beyond the conventional formulation, such as cosine-push regularization, entropy-based loss, and codebook-reset mechanisms. These additions interact in nontrivial ways, and as a result the first working configuration for the first experiment takes time to establish simply because there is no pre-existing recipe for our variant of the objective. Once this initial configuration is found, however, the process becomes easier. We consistently find that once a reasonable setting is chosen for one model, closely related values work across all others with little to no retuning. The only exception is the codebook size. Starting from very small values, increasing the codebook size improves VQ- VAE performance, but only up to two possible limits. The first limit is model size. If the VQ VAE is too small, increasing the codebook size provides little benefit, and making the model larger restores the usefulness of larger codebooks. The second limit is the inherent complexity of the task. After certain point, increasing the codebook size and the model size no longer leads to meaningful improvements. For these reasons, our experiments use codebook sizes ranging from 64 (for quantizing single token in single-layer hidden state) up to 1024 (for
18

---

Published as a conference paper at ICLR 2026
Normalized Theorstical vs Estimated Mutual Information by Codebook Size
Cadebook 128
Codsbook uize 256
512
Normalbnn bormalbed Estimated MI
Codebook size 1924
Purfect prediction
0.a
Normalized Theoretical
Figure 5: The comparison of our normalized latent mutual information estimations with different codebook sizes.
quantizing all prefix hidden states). For completeness and reproducibility, all hyperparameters can be accessed in the configurator files provided with the code.
A.4 VALIDATION ON REPRESENTATIVE EXPERIMENTS
To validate our use of vQ-vAE-derived codes for estimating between computations, we conduct three experiments below with different difficulty levels.
A.4.1 VALIDATION EXPERIMENT WITH KNOWN MUTUAL INFORMATION
We define discrete variables (A, B). Let A be uniform on a domain A with |A] 1024, and let
Pe
if b a otherwise.
PPB B = A a)
1023
The I (A; B) admits a closed-form expression that depends on Pa. We treat (A, B) as the ground- truth coarse variables that our method aims to recover.
To mimic transformer computations, we associate each a E A with hidden-state tensor G E R 110 capturing all block outputs except the final layer trained GPT-3 small sized model given a random prefix of length 10. Also, we associate each possible value of B with G E R 768 the final-layer hidden state at a single position. These tensors are redundant surrogates for (A, B) in the sense that they only contain 10-bit information, i.e., log, 1024, despite their huge dimensions, which match the largest hidden-state structures used in our experimental analysis.
We then apply our pipeline: train VQ- VQ-VAE models on GA and GB to obtain discrete codes Z and Zn. which serve as coarse summaries of the corresponding computations. Finally, we compare I (Z ZB) to the ground- tround-truthh I (A; B) across values of pe and across codebook sizes K. Figure 3 reports both I (A; B) and I (Z ZB) normalized within each K. For every K, the codes preserve the ordering induced by the ground-truth, and I (Z 7B) approaches the ground truth K in- creases. We do not train to recover (A, B), so (A; B) is not supervised target. Rather, I (A: B) is reference- process M/ that comes from the constructed latent variables defining the generative story behind (GA, GB).
proocesse
A.4.2 VALIDATION EXPERIMENT (HARDER SETTING WITH MULTIPLE SIMILAR PREFIXES FOR EACH LABEL.
We construct more challenging variant of App. A.4.I For each A we create set of 16 surrogate hidden-state tensors IR 10 Concretely, we take a random length- 10 prefix and generate 15 additional prefixes by changing single token random index, then the trained
19

---

Published as a conference paper at ICLR 2026
Normalized Theoretical vs Estimated Mutual Information by Codebock Size
Codelaskk
Oodabook aize
Codabonk atre 256
Codebooiz
Normalizud Estimatad M1
0.8
Codebook
Codcboolz size 2048
Perfect prediction
0.4
6.0
9.8
Normalized Theoretical MI
Figure 6: Harder validation experiment with similar prefixes for each label. Each label maps to 16 surrogates created by one-token edits. Normalized I (Za; Zn) remains order-consistent with reference-prose procccs MI across K, despite injected noise and within-class variability.
LM collect the corresponding hidden state block outputs. Similarly, for each b we produce 16 final-layer hidden states in 7666 using the same single-token perturbation strategy. Therefore, this experiment consists of 24 possible hidden state block GA and 2 possible final-layer hidden state GB. This setting reflects families of inputs that differ at one token yet retain essentially the same semantics, which exactly the kind of fine detail our coarse summaries should compress away. To further increase difficulty. we add Gaussian noise during VQ-VAE training on hidden states. We then compute normalized mutual information between the learned coarse variables, treating the 16 surrogates for fixed label equiprobable when forming empirical distributions. The results in Figure show that for every codebook size K, the codes preserve the ground-truth ordering across Pu, and I (Z ZB) tightens toward the reference as K increases.
A.4.3 VALIDATION EXPERIMENT (HARDEST SETTING WITH MULTIPLE UNRELATED PREFIXES)
We now remove within-class similarity entirely. For each a A we sample 16 independent length- 10 token sequences uniformly at random and collect their R 10x x hidden-state tensors. For each we likewise sample 16 independent sequences and extract the corresponding 707 final- layer states. Thus, the 16 surrogates that share label need not be close in vector space. This setting emulates many-to-on mappings from highly diverse inputs to same high-level code. This setting resembles natural language settings with dissimilar texts having similar meanings. Again, this experiment consists of 2 possible hidden state block GA and 2 possible final-layer hidden state Gg-
We repeat the same pipeline and probability treatment as App. A.4.2 estimating normalized between the learned coarse variables while assuming the 16 surrogates per label are equiprobable. Figure 7 reports the results. Although this regime poses 16 equally likely outcomes per input and eliminates structure exploitable by local similarity, our method still recovers normalized MI that preserves the correct ordering across p.. with improved agreement at larger K. This regime would be quite challenging for a probing-type method to recover any meaningful information due one- to-many mapping with totally unrelated targets.
B DETAILS OF EXPERIMENTAL ANALYSES
We provide details of experimental settings and analyses here.
20

---

Published as a conference paper at ICLR 2026
Normalized Theoretical vs Estimated Mutual Information by Codebook Size
Codelaskk
Oodabook aize
Codabonk aiee 256
Codebooiz
Normalizud Estimatad M1
0.8
Codebook
Codchoolz size 2048
Perfect prediction
0.
6.0
9.8
Normalized Theoretical MI
Figure 7: Hardest validation experiment with fully independent surrogates. Despite the lack of within-class proximity, normalized (Za; BB) preserves the ordering induced by the reference- process MI, and improves as the codebook grows.
B. LOSS FUNCTIONS
Let V be the vocabulary. sequence is (21...., Fr) with prefixes 2gt- The model M has embedding E. E Vixd and shared unembedding and width d. Hidden states are Hf Hg for l = I,..., and t = I....,7; I' write he for the final-layer state. We use for probabilities and log for natural logarithms.
R V|xd The transformer has layers
Next-token prediction (NTP). With teacher forcing, the baseline objective maximizes the con- ditional likelihood of the immediate next token from final-layer state hr. use a shared unembedding U E Vixd
T- Σ
Po (#e+1 agg)) softmax(U M/) 2011
CNIP(
log Pe (# xg()
Multi-token prediction (MTP). To align internal states with short-horizon futures, we equip the model with I' output heads Gloeckle al. (2024). Head y E {l..... } implemented trans- former layer p(t). Rd + I2d that takes h. as input and is followed by the same shared unembedding U:
Po (gg))= softmax (U (l))))))))))
(h+))
The training loss sums the cross-entropy across positions and horizons:
T-r
LMTP(
Σ
Σ log Po [age)-
All heads contribute equally. At test time decoding remains autoregressive; the auxiliary heads affect only the training signal and can optionally be used for speedups.
In our experiments, when we use MTP loss in CFG task, we used r = 4. As for, PF-Long and PF-Short tasks we used r
B.2 DATASET
Here explain how generated the dataset and how used them in LM training.
21

---

Published as a conference paper at ICLR 2026
B.2.1 CONTEXT-FREE GRAMMAR DATA GENERATION DETAILS
The specific formal details of our CFG are not central to the claims of this paper. We use it only as controlled source of long, purely syntactic sequences. Readers who want a formal presentation and broader context can consult the CFG-based data generation reference Allen-Zhul (2024) on which our CFG data is based on. Here we provide the details of our specific CFG parallel to the naming convention in Allen- (2024).
We construct a layered grammar that expands strictly downward from a single start symbol to a set of terminals. The layers are organized as one start symbol, then four intermediate nonterminal layers of size three each, and finally a terminal inventory of sixty four symbols, summarized as the tuple (I, 4, 4, 4, 64). For every nonterminal we define small menu of alternative productions; each nonterminal has between three and five alternatives, and each alternative rewrites the parent into two three symbols drawn only from the next layer. This acyclic layout fixes the depth of any derivation and guarantees termination.
Once the grammar is fixed, each sequence is produced by starting the start symbol and repeatedly expanding pending nonterminals until only terminals remain. Whenever a nonterminal has several alternatives, one is chosen at random. Because expansion always moves one layer down, sequences are long but tightly controlled in length. In the configuration used in the paper we generate four million sequences. The shortest contains 16 terminals and the longest 67 (the number of terminals corresponds to the sequence length), measured without begin or end markers. The resulting corpus contains 4M sequences (150M tokens)
B.2.2 PATII-FINDING DATA GENERATION DETAILS
PF-Short is built to be small but nontrivial
task. In every instance the start node is labeled and the goal node is labeled 2. We first lay down two correct trunks from the start to the goal. Each trunk has exactly two internal nodes, so a correct answer has the form 1, a, b, 2. We then lay down two decoy trunks that also begin at the start but end at fresh terminals that are not the goal. Decoys are the same length as the correct trunks, and the interiors of all trunks are disjoint. Up to this point the graph contains the start, the goal, two correct routes that meet only those endpoints, and two look-alike alternatives that never reach the goal.
To make the local neighborhood around each trunk less revealing, we sprinkle branches on top of the trunks. We do this in one layer: for each eligible trunk node we independently add a small number of fresh children drawn from Poisson distribution with mean one, and cap the number of branch children per node at two. We never branch from the goal or from decoy terminals, and for the start node reduce its branch allowance account for the four trunk edges already attached to it, so its total degree remains bounded. If the instance would exceed the node budget of 28 distinct labels we discard and resample.
Before writing the example to disk we randomly relabel every internal node using a permutation of the labels through 28, while keeping and fixed. We also shuffle the edge list. The model never sees coordinates or orders, only an unordered set of undirected edges and an answer slot. A simplified typical prompt looks like
1 9, 7 2, 9 13,37,52. 2;
and either 5 or 1 37 2 accepted as correct, In our experiments the PR-Short training split contains 16M samples generated with the recipe above. We also have a similar sized dataset to train corresponding vqvaes and a smaller validation data. Please also refer to Figure (right) for a simplified network example.
PF-Long extends the horizon while keeping local clutter modest. We again fix the start to and the goal to 2. We lay down two correct trunks, each with four internal nodes, so correct answer has the form 1, a, b, c, d, 2. We then add one decoy trunk of the same length that begins at I but ends at a fresh non-goal terminal; interiors are disjoint. To add distractors we apply two-layer branching pass: a light first layer around trunks followed by a very light second layer, with at most one branch child per parent. We never branch from the goal or from decoy terminals, and we reduce the start's branch allowance to respect the degree cap. If a sample would exceed the budget of 28 distinct node labels we discard and resample. Finally we randomly permute the internal labels within 3 through
22

---

Published as a conference paper at ICLR 2026
28, keep and fixed, remap the paths, and shuffle the edge list. The training split contains 16M PF-Long examples generated with this recipe.
B.2.3 NATURAL LANGUAGE (OPENWERTEXT)
Finally, we include a large-scale natural language dataset (OpenWebTexx) to examine whether our information- theorretic analyses extend to unconstrained real-world data. Unlike the synthetic CFG or path-finding finding settings, this dataset reflects distributional structure from natural language and allows us to test whether planning signals identified in controlled tasks also emerge under standard pretraining conditions.
B.2.4 HOW THE GENERATED DATA ARE USED TO TRAIN THE LANGUAGE MODELS
For the CFG and OpenWebTee corpora we run conventional pretraining in separate setups. For CFG, first generate many standalone sequences that satisfy the grammar. We then randomly shuffle these CFG sequences and concatenate them end to end to form single text stream for that corpus. Training samples are taken by choosing a random starting offset in the CFG stream and slicing fixed-length block equal to the model context window; the loss is computed on all tokens the block except any optional boundary markers.
For Open WebTeet, apply the same pretraining recipe but within its own corpus only: documents are shuffled and concatenated to form an Open WebText stream, blocks are sampled at random off- sets, and next-token prediction is used with causal masking. There is no mixing or interleaving between CFG and Webxet at any stage.
For the path-finding data we use supervised finetuning rather than corpus streaming. Each example is a self-contained prompt and answer: the input is the unordered edge list written as camma-separated node pairs followed by colon, and the target is any valid shortest path from the fixed start to the fixed goal. We concatenate prompt and answer for teacher forcing, mask the prompt so the loss is applied only the answer span, and we do not concatenate different path-finding examples split one across blocks. Under this scheme we use both phort and PF- PF-Long, whose correct answers have lengths four and six respectively when counting the start and goal.
B.3 DETAILS PLANNING HORIZON EXPERIMENT 3.
Details about VQ-VAE. The VQ-VAE used to quantize the prefix representation H RTxLxd employs 256-entry codebook (ZH {0, 1..., 2555), ) while each future final-layer state hy Rd is quantized with a 64-entry codebook (Znz... E {0, .,..,.,6)))
For the CFG task, the dimension of H is on average 16 x 11 x 768 = 135,168, and the dimension of hū is 768. The mean nRMSE values of VQ-VAE encoding into summary codes is 0.47, and that of encoding h. is 0.21. Representative codebook similarity and usage plots from VQ-VAE training to encode H for the CFG task are provided in Figure and Figure Also, those of VQ-VAE to encode ht. for the CFG task are provided in Figure 10 and Figure]
For the PF-Short task, the dimension of is on average 70x 768 = 591,360, and the dimension of h5 is 768. The mean nRMSE values of vQ-VAE encoding H into summary codes is 0.75, and that of encoding h. is 0.48.
For the PF-Long task, to reduce dimensionality, subsample every three layers when encoding И due to the limited model size, and we discard token indices corresponding to comma (.) tokens. The dimension of H is average 54 x 4 x 768 165,888, and the dimension of h is 768. The mean nRMSE values of VQ-VAE encoding into summary codes 0.59, and that of encoding h 0.49.
We emphasize that these are the results of huge quantization mappings, e.g., from 591,360 real numbers 256 discrete labels.
23

---

Published as a conference paper at ICLR 2026
16
1rae werters 316
90110 1.00
00
04
sector PdRcc
Figure The codebook similarities for VO-VAE encoding H in CFG task.
Figure 9: The codebook usage for VQ VQ-VAE encoding in CFG task.
B.4 DETAILS BRANCHES IN THE PLAN EXPERIMENT 3.2
Details about VQ-VAE. For the VQ-VAE. details encoding in PF-Short and PF-Long tasks, please, refer to the relevant part of App. B.3 since the same models are used.
Encoding the paths. The paths themselves does not have vector representation and it is chal- lenging use a unique class each possible path since in PF-Long, there can be 28 paths. Although PF-Short is relatively easier, reduce the number of possible codes to 512 by training VAE over a learned embedding vector for each of 28 tokens. Then we obtain the high-level codes of full path using the trained VQ-VAE. For PF-Long, from VO-VAE and each intermediate token (using the token't 's index itself) along the response path. Then, take the mean of these values over the nodes the path get a full-path MI result.
calculate the MI between the H's code obtained
24

---

Published as a conference paper at ICLR 2026
10
Usee 68 mu wortters
02
-04
Codebeec viecter
Figure 10: The codebook similarities for VQ-VAE encoding H in CFG task.
300
10
20
50
Eodebook Vestor
Figure 11: The codebook usage for VQ-VAE encoding in CFG task.
B.5 DETAILS OF THE INFORMATION IN THE COMPUTATIONAL HISTORY EXPERIMENT
In this section we provide details of the information in the computational history experiment.
B.5.1 NMI RESULTS IN S B FOR LONGER-HORIZON GENERATED TOKENS (HIGHER T)
We provide the results in Figures 13
B.5.2 DETAILS OF THE CONDITIONAL NMI ESTIMATES
To estimate I(24- - 24 Z4), we used the identity:
I(Z4-16.7 74 z = I677-1677 74) - I(74 -II(7-15-T-1-2 747,
15.7;
(4)
which holds because
bィ -4-15:1 12/ 15 2f}.
The first term the right hand side (4). is already estimated and shown in the Figure 4 (k block estimates). As for the second term, find approximate value follows. First, because of
25

---

Published as a conference paper at ICLR 2026
nM1 over Block and Layer for the Token Generated at t=
0.54
0.54
0.55
0.55
0.55
0.55
0.56
0.56
0.56
0.57
0.57
0.56
0.57
Layer Number (0)
6
0.51
0.51
0.51
0.31
0.51
0.51
0.53
0.52
0.53
0.53
0.53
0.53
0.53
0.97
5
0.44
0.44
0.45
0.45
0.40
0.40
0.39
0.38
0.38
0.38
0.30
0.38
0.38
0.38
0.38
3
0.40
0.39
0.38
0.39
0.38
0.38
0.38
0.36
0.38
0.38
0.38
0.38
0.36
0.85
2
0.35
0.33
0.34
0.34
0.35
0.35
0.31
0.34
0.34
0.31
0.34
0.34
0.33
0.31
0.41
14
13
12
11
10
9
4
2
Block Number (k)
nMI over Block and Layer for the Token Generated at
7
0.54
0.53
0.56
0.55
0.35
0.56
0.36
0.55
0.57
0.57
0.56
0.56
0.58
0.70
Layer Number (
0.52
0.52
0.53
0.45
0.46
0.44
0.45
0.41
0.44
0.45
0.44
0:45
0.45
0:45
0.45
4
0.41
0.40
0.39
0.39
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.53
3
0.39
0.35
0.34
0.35
0.35
0.35
0.31
0.34
0.34
0.35
0.34
0.35
0.36
0.36
0.50
1
0.34
0.34
0.33
0.35
0.34
0.34
0..34
0.344
01.34
0.34
0.34
0:33
0.34
0.37
14
13
12
11
10
9
7
4
2
Block Number (k)
nMI over Block and Layer for the Token Generated at ==
7
0.55
0.57
0.57
Layer Number (2)
0.52
0.51
0.51
0.52
0.52
0.52
0.52
0.52
0.52
0.52
0.53
0.53
0.51
5
0.46
0.45
0.45
0.44
0.44
0.44
0.45
0.44
0.44
0.45
0.45
0.45
0.45
0.53
4
3
0.40
0:40
0.38
0.39
0.39
0.38
0.38
0.38
0.38
0.39
0.38
0.38
0.38
0.48
2
0.35
0.35
0.35
0.35
0.35
0.31
0.34
0.34
0.31
0.34
0.35
0.35
0.37
0.45
0.34
0.33
34
0:34
14
13
11
9
7
2
Block Number (k)
Figure 12: Heatmap of normalized mutual information over block and layer indices for 0, 1,2.
Please refer to the caption of
the underlying task's symmetry (it is Wehtext, pretraining task), we have translational invari- ance,
1 2+) I 2++1).
(5)
Then, we approximate the right-hand side of this equation by
I(Zi-147: k4+1) ~ I(Z 27+1).
which relies the observation that in natural language, the incremental information contributed by extending the context diminishes: as sentences become longer, the mutual information between past. tokens and future ones grows only logarithmically.
B.5.3 DETAILS ABOUT BLOCK VQ-VAE.
The vQ VAE used to quantize gth layer tth block representation hf 102:7:10 18-2-1 E 16xd em- ploys 1024-entry codebook (Bf (0, 1, 1,... 1023)), while each future final-layer state h. Rd
26

---

Published as a conference paper at ICLR 2026
nM1 over Block and Layer for the Token Generated at t=
0.53
0.55
0.55
0.55
0.55
0.56
0.56
0.56
0.57
0.56
0.56
0.56
0.57
0.62
Layer Number (0)
6
0.51
0.52
0.51
0.31
0.52
0.52
0.52
0.52
0.32
0.52
0.53
0.53
0.53
0.38
5
0.44
0.45
0.45
0.45
0.41
0.39
0.38
0.39
0.38
0.39
0.30
0.39
0.38
0.38
0.38
3
0.39
0.39
0.38
0.38
0.38
0.38
0.38
0.36
0.38
0.38
0.38
0.38
0.36
0.45
2
0.35
0.34
0.34
0.33
0.31
0.31
0.33
0.34
0.31
0.34
0.33
0.33
0.31
0.36
14
13
12
11
10
9
4
2
Block Number (k)
nMI Block and Layer for the Token Generated at
7
0.54
0.54
0.54
0.54
0.35
0.55
0.36
0.55
0.56
0.57
0.57
0.56
0.57
0.60
Layer Number (
0.51
0.52
0.46
0.45
0.45
0.41
0.45
0.45
0.45
0.44
0:45
0.14
0:45
0.45
0.48
4
0.41
0.40
0.39
0.38
0.39
0.39
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.42
3
0.39
0.34
0.35
0.34
0.35
0.31
0.31
0.34
0.34
0.31
0.34
0.34
0.35
0.36
0.40
1
0.33
0.34
0.34
0.34
0.34
0.34
0.33
0.344
0..34
0.34
0.33
0:33
0.33
0.35
14
13
12
11
10
9
7
4
2
Block Number (k)
nMI over Block and Layer for the Token Generated at T
7
0.55
0.55
0.57
0.57
Layer Number (e)
0.51
0.51
0.51
0.52
0.52
0.51
0.53
0.53
0.52
0.53
0.52
0.53
0.51
0.56
5
0.46
0.45
0.45
0.44
0.45
0.45
0.45
0.45
0.44
0.44
0.45
0.45
0.45
0.47
4
3
0.40
0.39
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.38
0.42
2
0.35
0.35
0.35
0.35
0.31
0.31
0.35
0.34
0.35
0.34
0.35
0.35
0.35
0.39
0.34
0.34
0.34
34
0:34
14
13
11
9
2
Block Number (k)
Figure 13: Heatmap of normalized mutual information over block and layer indices for 3, 4, 5.
Please refer to the caption of
quantized with a 64-entry codebook (244,- (0, 10.,,...., 63)). Recall this experiment is on NLP (OpenWebTes) WebText pretraining). The dimension of h4 167-161-441-100 is on average 16 x 512 = 7680, and the dimension of he is 512. The overall VQ-VAE for the blocks has 4M parameters in total. Repre- sentative codebook similarity and usage plots from VQ-VAE training to encode hf 192.7:10 168 for the (OpenWebTees) WexText) task are provided Figure [ and Figure 15 For additional details on VO-VAE trainings on WebTett (NLP pretraining) data please refer to App. B.6
B.6 ON THE SCALABILITY OF OUR PIPELINE
Recall that to obtain the code Z we use encoder whose input is h4 + To study this, trained VQ-VAES of increasing sizes on OpenWebText with a fixed codebook size of 64 vectors. OpenWebText is widely used for NLP pretraining, but the data lacks strong coherent structure across samples. As a result, it is natural, and fact expected, that reconstruction errors remain relatively high.
27

---

Published as a conference paper at ICLR 2026
Cosine Similarities Belween Used ViAAAEI Codebook Vectors
Used vectors: 729
100
Used Vestor 1008
100
300
400
700
Used Veccor Indes
Figure 14: The codebook similarities for VQ-VAE VAE encoding 1o in OpenWebText task.
1>
Figure 15: The codebook usage for VQ-VAE encoding 16 169 in Open Weh'ewe task.
Figure 16 reports the normalized nRMSE function of model size. This hints that scaling the encoder can improve representation quality despite the inherent noisiness of the dataset.
C RELATED WORK IN DETAIL.
We provide longer discussion of the Related Work the main text $
Language model planning. In order to improve LMs' reasoning and planning abilities, re- searchers have developed scaffolding and augmentation techniques, including Chain-of-Though prompting (Wei et al.) 2022). planning tokens (Wang et al.) 2024c (Sel et 2025), and structured inference methods such as Tree-of-Thoughts and Graph-of-Thoughn (aao et al. 2023 Besta et al.) 2024). In parallel, hybrid systems that integrate LMs with symbolic planners or external tools have achieved of-te-of-the-arr performance in embodied and tool-use domains (Zhao et al. 2023 Wang et al.) 2024a] Shen al. 2023). Despite these advances, recent studies highlight that significant challenges remain, and current approaches to LM planning still fall short of fully addressing com- plex reasoning and decision-making tasks. For example, Lin al. (025) show that models can
28

---

Published as a conference paper at ICLR 2026
0.91
0.90
0.89
0.88
30
0.87
202 SIM
0.86
6.0
6.5
7.0
7.5
8.0
log10/Number parameters)
Figure 16: The nRMSE vs VQ-VAE encoder parameter count plot.
produce conflicting answers under logically related prompts despite local plausibility, which shows lack of long horizon planning, although short-horizon planning works fine; Ahhn Yin (025) re- veal discrepancies when asking model what is correct vs what incorrect; and[SSaxena et al. (2024) demonstrate that even repeated queries yield inconsistent outputs. Beyond inconsistency, some stud- ies find that out-of-the-box models struggle with even simple planning tasks (Momennejad et al. 2023). Complementing these empirical findings, theoretical analyses suggest that autoregressive models may face fundamental limits in their planning ability (Wang al. 2024b). Taken together, these findings underscore that understanding whether and how planning arises in LMs is not only an open empirical challenge, but also central both their interpretability and the principled design of future model architectures.
Behavioral and Mechanistic Interpretability. Some interpretability (explainability) methods consider the LM as black box, and design tasks or benchmarks to gauge reasoning, robustness, or generalization abilities (Srivastava et al. 2023 Liang et al. (002)). In contrast to them, mechanis- tic interpretability seeks to reverse-engineer transformer computations into human uumaa--nnerstandablle parts, treating the residual stream as the main information pathway and attention heads separate components that pass information along (Elhage et al.] 202T]. Using this approach, researchers arise from characteristic two-
have explained key phenomena in LMs: in-context learning head circuit that appears during a training "phase change," with causal support from perturbation studies (Olsson 2022); many neurons polysemantic, but sparse autoeneoders (SAEs) place them with more interpretable, monosemantic features that enable finer-grained causal anal- ysis (Cunningham et al. 2023 Bricken et al. 2023); models implicitly encode structural linguistic properties such as syntax (Hewitt & Liang 2019). Building on earlier approaches, transcoders, ap- proximate dense MLPs with wider sparse layers, separating circuits into input-invariant and input- dependent components, and matching or surpassing SAES in sparsity, faithfulness, and interpretabil- ity (Duncisky et al., 2024). Frontier-scale case studies use cross-layer transcoders trace multi-step reasoning and other behaviors, illustrating how these tools can audit mechanisms, not just features, in modern LLMs (Lindsey et al.) (2025). Overall, mechanistic interpretability offers concrete tools for uncovering how LLMs compute, though these methods are still developing.
Mathematical Perspectives. Beyond empirical tools, mathematically grounded perspectives also illuminate transformer/LLM interpretability: aitomata-theorett analyses show self-attention can implement finite-state algorithms (Liu et al. (023);; optimization views prove in-context learning corresponds to (preconditioned) gradient steps (Ahn et al. 2023) and transformers are mutual inter- action learners (Ustaomeroglu & Qu 2025); meal-field theory gives global convergence guarantees at scale (Gao et al. 2024).
Probing. different line of interpretability work views LLM hidden states as structured represen- tations that be "probed", probing refers training lightweight models, often linear classifiers,
29

---

Published as a conference paper at ICLR 2026
read out specific information from hidden states in order to test what the model represents internally. Using probing, researchers have shown that transformer hidden states encode structured belief-state and world- model-like information (Shai et al.) 2024] Gurnee & Tegmark] 2024 Hazineh et al. 2023). Other works demonstrate that a single hidden state can carry information about multiple future tokens (Pal al. 2023) and that probing reveal the underlying algorithms L.MSs use to solve tasks (Allen-Zhu) 2024).
However, standard accuracy-basee probing has drawn a critique for combining what the probe can learn with and what the representation actually encodes making the results sensitive to probe ca- pacity, data size, and hyperparameters Hewitt & Liang (immetel al. (2020). Further, high probe scores often come from exploiting superficial linear context cues rather than genuine struc- tural knowledge (Kunz & Kuhlmann, 2020p. Probing can even reveal features that model does not use for its task (Ravichander et al. 2021, Kumar et al. 2022). Consequently, some critiques moti- vate adopting unformation-theowe approaches illustrate this perspective. The original Information Bottleneck framework (Tishby & Zaslavsky 2013; Shwartz-Ziv & Tishby] 2017) casts learning as trade-off between compressing
lenses (Voita & Titov 2020; Diego-Simón et al. 2025). Several
input representations and preserving predictive information about outputs. Complementary work by Voita et al. (2119) examined how information flows across transformer layers under different training objectives, revealing that tasks like language modeling, masked language modeling and machine translation induce distinct patterns of information loss and reconstruction. More recently, Skean al. (2025) demonstrated that intermediate layers of LLMs often yield stronger representa- tions than the final layer, using unified metrics based on entropy.
D EXTRA EXPERIMENTS WITH LARGER MODELS AND COMPARISON WITH PROBING
the In this section we (i) scale experimented language model and repeat the Horizon the Plan analysis,
(ii) add probe-based baselines, and (iii) explain why probe performance can be confounded in our setting.
D.1 SCALING TO A LARGER MODEL
We scale the basdd- decoder-only architecture with Rotary Positional Encoding (Su et al,, 2024 to 24 layers and demoueel 1024 (> 0.3B parameters), than twice the size used in the main experiments. We rerun the Horizon of the Plan experiment on the CFG dataset from Section [.] The 0.3B model attains near-perfect sequence-complltion accuracy Our method scales smoothly with larger language models.
Figure repports the normalized mutual information between the prefix summary codes and the last- layer hidden-state codes at generated positions t + r. The nMI curve decays rapidly as 7 increases, matching the trend in smaller models. On CFG, this is consistent with myopic behavior rather than long-horizon planning, which aligns with the task's syntactic structure.
D.2 PROBING EXPERIMENTS
Setup. Probing fits a supervised predictor from internal states a target and uses generalization performance as a proxy for whether the information is linearly or simply recoverable. We evaluate probes in the same setting as $D.I]
Let L be the number of layers. For prefix ending at time t, denote by
Hpet
(he, E R 8 1:L, i t}
the block of hidden states across all layers and prefix positions. For a future offset r 2 1, denote the generated token by 7 and the last- hidden state by h/ Rdd. We train separate probes for each r. Since the prefix length can vary across samples, we use zero padding for shorter samples.
"Sequence-completion accuracy is the fraction of prefixes for which the model completes CFG-valid continuation.
30

---

Published as a conference paper at ICLR 2026
CFG Task: Normalized MI vs.
1.0
0.8
nMI
0.6
0.4
0.2
Normalized MI between prefix summary codes and the codes of the last-layer hidden state at future position + r on CFG for the 0.3B LM. The rapid decay in nMI with 7 replicates the main-text trend at larger scale.
Figure 17:
Token prediction probe. A two-layer MLP o, maps H pprtt to distribution over the 32 CFG tokens. The loss is cross-entropy,
tok
E
log Po. (Zetr Hprer))
and we report accuracy.
Hidden-state regression probe. two-layer MLP t/r predicts the future last-layer hidden state,
[[Ib,, (Hoprrt 11++1130
We report normalized MSE,, measured by dividing the MSE values by the highest MSE in that setting.
Probing method: token accuracy
Probing method: token accuracy
0.25
0,300
Mean Accuracy
Accurace
0.75
0.23
0,083
0300
0L000
0005
3000
2000
4020
5000
Traming iter.
Probing method: last layer
Probing method: last layer
Moan Inlidation normalized MSE
hidden state normalized MSE vs
IDean hidden state normalized MSE vs Iter.
Validation normalzzed MSE
2.98
0.600
0.625
Mm pommMyy MM
0.00
0050
0,23
Training
Figure 18: Probe baselines on CFG. Top left: token accuracy vs. future offset Top right: mean token accuracy vs. training iteration. Bottom left: hidden-state regression normalized MSE, vs. 7. Bottom right: normalized mean MSE vs. training iteration.
31

---

Published as a conference paper at ICLR 2026
Findings. Token probes decisively above random sampling only for O; beyond the imme- diate next token the accuracy does not change much and not yield a clear picture for interpretation. Hidden-state regression probes close normalized MSE, across т. Even with nonlinear two-layer MIPs and per-r training, the results are inconclusive about the structure of long-horizon computa- tion and planning in the language model. This reflects a practical limitation of probe-style supervi- sion when the input pree and the target hi+ are both high dimensional.
Relation to linear or higher-capacity probes. A linear probe is special case of the ML.P used here, so the nonlinear results upper bound what linear probe can extract under the same inputs and targets. high-capacity probes are used, supervised training can memorize dataset idiosyncrasies and fit mappings from prer the target that the LM itself does not implement, yielding optimistic scores without evidencing a mechanism in the model. These effects blur the link between probe performance and information flow.
We include probes as a diagnostic comparison, not as an estimator or validation of our MI measure. Probes answer different question how well a chosen supervised predictor recover hand- specified target from -nd their performance can be sensitive to predictor capacity and target complexity. In our setting, probing does not yield stable interpretable signal about long-horizon structure across т, whereas our MI analysis measures statistical dependence between learned com- pressed summaries. See Appendix D.3 for an illustration of why probe scores can be confounded.
D.3 WHY PROBING CAN BE CONFOUNDED AND A SIMPLE ILLUSTRATION
Target-variance sensitivity. Prior work (Voia & Titov 2020 that probe performance can be dominated by the marginal complexity of the target rather than by information shared with the source.
Diego-Simón et al. 2025) has noted
Let X
Onif 10] so (XX log, 10 3.32 bits. Define
Y22 {o
fX
with prob. 0.5, with prob. 0.5.
Yi = 0,
Then Y() = 0 while r(X; Y2) = (X) RO 1.66 bits. A probe achieves zero error on Y1 yet struggles on Y2. which would wrongly suggest that contains more information about Y1 than Yg- The discrepancy arises because probe loss is governed by the target's marginal variability, not by the actual information shared with X. Furthermore, as discussed in the previous section, higher capacity probes can blur the line between dataset idiosyncrasies and actual language model mechanisms.
Why our method avoids this pitfall. Our method trains two VQ- VAE encoders E1 and Ey on prefix and future hidden states using only reconstruction losses, with no labels task supervision. We then estimate mutual information between Z (E, (h1): ): (22)) and interpret it comparatively in all tasks. Our conclusions are comparative within fixed experimental setup: we hold the encoders, codebooks, and MI estimator fixed and report normalized MI between the same representation types while varying only the variables of interest. For example, when comparing the MI between the summary of pre-output computations and the summaries of the first vs. second output token, every pipeline component is identical, so any MI shift must arise from the underlying interaction among the selected variables. Finally, by obtaining unsupervised, high-level representations, we reduce the effect of any blur induced by target-specific idiosyncrasies in the raw hidden states. Our method coarsens the space and the interaction among the underlying variables is not probe- inbee-idduedd short- cuts.
D.4 PREDICTIVE U-INFORMATION CAN BE ARBITRARILY DISTORTED BY RESCALING
Xu et al. define predictive y-information (xXX + Y) from squared-loss prediction problem over a Gaussian linear family. In the setting of their Proposition 1.5, this reduces to
-X) (n(cov(Y)) R².
32

---

Published as a conference paper at ICLR 2026
where R is the coefficient of determination of the optimal linear regression of Y X Unlike Shannon mutual information, this quantity is sensitive to the marginal scale of and can changed arbitrarily by rescaling the target, even when the underlying information content is fixed.
To illustrate, let X NO 1) and e N/ 1) independent. Define
Y1 = X +e, Y2 a Y = a(X +),,
for some scalar a # 0. Since Y2 is invertible linear transformation of Y, Shannon mutual information invariant:
(XX;Y)=1((;3))
Now consider I, in the Gaussian linear setting Xu al. We have
VYr = Var(X) + Var(c) = 2,
Var(Y2) = a Var(Y)) = 2a²
The optimal linear predictor of Yτ from X is proportional to X for both k: = 1, and the squared correlation is
R R2 Cov(X,Yi))
1
Var(X)Var(Y))
1.2
since scaling Y does not change correlation. Therefore,
[x Y1) tr(Cov((Y)) R = 2
while
L(X - 12) = tr(Cov((2)) R2 2n² =
By choosing |a] arbitrarily large or small, we can make I,(X Y2) arbitrarily larger smaller than (XX + Y1). despite the fact that Yi and Y2 contain exactly the same Shannon information about X. This demonstrates that predictive y-information, in this common probe setting, fails to satisfy a basic information-theont theoretic property: invariance under invertible transformations of the target (equivalently, does not obey the data-processing inequality with respect to deterministic rescalings of Y).
Why our method avoids this scaling pathology. In contrast, our VQ-VAE approach estimates Shannon mutual information between learned discrete codes I (E, (hi); (2())). Because mutual information is invariant under invertible transformations of each argument and satisfies the data- processing inequality, rescaling reparametrizing the underlying continuous hidden states before encoding cannot arbitrarily inflate or deflate the measured dependence. Once the encoders and codebooks fixed, changes in our estimated MI across conditions reflect differences in shared structure between the representations, rather than arbitrary choices of units or target scaling.
D.5 I, MUTUAL INFORMATION PLOTS
In addition the probing results in Fig. 18 investigate another probing approach with different
mutual information variant. We reproduce the experiment in App. D via I, (H..--1, :, defini- tion from Xu et al. As it is seen from the results, Fig. the I, [1 1: (/4-)) definition is not helpful either. Also, test validity of In with easiest validation test from A.4. The results
are seen in Fig. 20 As is seen I, fails on our validation test as well.
E FINITE SAMPLING ERROR
We quantify the finite-sampn sample error when estimating the mutual information between two discrete random variables X and Y over two different MI estimation approaches.
Let X and Y take values in (1...., .Kx and {1...., ........ Ky }, with joint pmf Piş = Pr(X i, j), marginals E, più and pj := E; Pij. and true mutual information
Piji.
T(X:)) ΣΣ Pig log Pú
P.P;
'Equivalently, Tx(X + Y) is the total variance of Y times the fraction of variance explained by the hest. linear predictor.
33

---

Published as a conference paper at ICLR 2026
Probing method: ly[H- -hh vs
Probing mean MSE across t iter.
168
2:25
2.30
230000
Mean MSE.
2.15
M M M
2.10
210000
2.05
1000
2000 Training iter.
3000
4000
5000
Figure 19: Reproduction of experiment in App. via Ty-
Normalized Theoretical Unnormalied Probing
Probing -nh)
Normalzed Prohing k0H
Probiing
Probing Resula
Perfect socdiction
Probing Beeall
0.4
Normalized Theoretical
Figure 20: Plots to check validity of I, the easiest validation experiment in App. A.4.1
From N Lidd samples (X,, Y2) drawn from (X,Y) Y), define the cell counts
Ng Σ 1(x,= =j N
Nij.
N2
Nijs
and the empirical probabilities P.j == No/N, P. = N./N, PÚ = N2/N. The standard plug-in (maximum-likelimhe mutual information estimator is
Kx Ky
i(x;Y) Σ
Pig log Pij prD3
(6)
Asymptotic bias and variance of the plug-in MI. Under standard regularity assumptions (in particular, Pij >> for all j, which are satisfied in our experiments), an exact local expansion plus
delta-method argument gives the following large-N behavior (Panzeri & Treves 1996 Treves &
Panzeri 1995 Paninski 2003]:
bias
= - ((x:V)
(Kx - 1)(Ky -1) +0
1n2
N2
Crar (X, Y)
1
Var Var(f(X:Y))
+0
N
N2
where Cvar(X, 0 is constant that depends only the true joint distribution Pij- Conse- quently, the mean-squared error (MSE) satisfies
MSE()(X;Y)) = -
CMSE XX, Y)
+0
1
99)
N
N2
34

---

Published as a conference paper at ICLR 2026
for another constant (MME(XXY) O. In other words, the plug-in MI estimator has
bias O (N)
Var = =o())
MSE o (N))
A simple first-order bias correction that cancels the leading O(1/N) term in is
+ (Kx
() - 1)
(10)
2N In2
Then
(X): 0
(11)
N22
(ar(Ipr(X;Y)) Cvar(X, Y)
1
+0
(12)
N
N2
MSE (X) = CMSE(X,y)
1
+0
(13)
N
N2
i.e, the MSE still scales O(1/N) but with smaller constant prefactor Treves & Panzeri 1995
Paninski-type estimators and worst-case MSE upper bounds. Paninski (2003) constructed family of linear estimators for discrete entropy, and hence for mutual information via [X;; (X) + H(Y) H(X, , Y), by explicitly optimizing a rigorous upper bound on the worst-case MSE over all discrete distributions with given alphabet size. The resulting "best universal bound" (BUB) estimators (X; Y) satisfy bounds of the form
CBUB(Kx, x, N)
sup E [ ((x0))--(x;;))) I(X:)))
(14)
N
for an explicit constant CBUB(K' (Kx, Kr,N) that can computed from (Kx, Ky, N); particular, the worst-case MSE is also upper bounded by a constant times 1/N.
What we use this work. In our experiments, we computed mutual information using both (i) the plug-in estimator (6) typically with the Panzeri- Treves bias correction (TO). and (ii) Paninski's BUB-type estimators assembled from the corresponding entropy estimators (Paninski 2003). In
all cases tested, both approaches produced essentially identical values of 7(X;Y). Since our available sample sizes N are very large compared to the effective domain sizes Kx and Ky (we have easy access to CFG, Path Finding, and NLP samples), the (1//N) finite-sample errors of both methods are very small. In most of our experiments we therefore report the mutual information estimated using the plug-in estimator.
THE USE OF LARGE LANGUAGE MODELS (LLMs)
We used LLMs in the following ways.
We polished sentences and corrected the grammar.
For visualization (figures, plots) of the results, made LLM to write the visualization code.
For the experiment codes we got LLM help, e.g., tab completions, function generation given the prompt, and we checked all LLM generated code line by line to ensure it works as intended.
At the beginning of the project, we used an LLM help identify relevant literature and read the suggested papers most pertinent to our work.
35


Reviews:
[{"review_id": "AaeKKZoXcB", "reviewer": "Reviewer_Cb1t", "summary": "This paper studies planning capabilities of LLMs through an information-theoretic lens. The authors develop a framework using Vector-Quantized Variational Autoencoders (VQ-VAE) to compress the hidden states of the LLM into discrete codes. This allows them to compute mutual information (MI) estimates between different \"computational blocks\"/hidden states. \nUsing this approach, they analyze three key aspects of planning: (i) the planning horizon (how far ahead models plan before generating tokens), (ii) branch awareness (whether models internally represent alternative valid continuations), and (iii) computational history dependence (which earlier computations inform current predictions). \nThe framework is evaluated across synthetic grammar tasks, path-finding problems, and natural language (OpenWebText), comparing next-token prediction (NTP) versus multi-token prediction (MTP) training objectives.", "questions": "- For each experiment, do you train one VQ-VAE per dataset/task (e.g. one for CFG prefix, one for PF prefix, one for PF paths, one for OpenWebText blocks), or per model checkpoint?\n\n- Roughly how many sequences and optimization steps were needed to train your VAEs?\n\n- What are the overall conclusions for NTP versus MTP?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- I think the idea behind the paper is quite novel: thee use of information-theoretic measures via VQ-VAE compression to study planning is a creative methodological tool\n\n- Also, the specific focus on quantifying planning through three different characteristics is exciting. Namely, via horizon, branching, and history dimensions. This provides a more structured framework for understanding LM internal computations. \n\n- The experimental analysis is comprehensive, including controlled synthetic tasks (CFG, path-finding) to natural language. \n\n- The path-finding task with disjoint correct and decoy paths is particularly well-designed for testing branch awareness. I really liked the idea. \n\n- This work addresses an important open question about whether and how LMs engage in planning-like computations. The finding that planning behavior is task-contingent has implications not only for model design, but also for training strategies. \n\n- I believe the framework itself could also be valuable for broader interpretability research, not just planning.", "weaknesses": "While I really enjoyed this paper, I think it also has room for improvement:\n\n- The main text simply says \"we train a VQ-VAE\", but since this is a key aspect of the method, I would expect more information about it in the main text. However, crucial aspects are only in the Appendix. \n\n- Given the details in the appendix, the number of hidden states to store is huge, which suggests that the framework is non-trivial since the VQVAE is difficult to train. \n\n- The experiments were carried using only GPT-3 Small models (12M-202M parameters). Findings may not generalize to larger models where emergent capabilities differ. It also becomes more difficult to apply the framework for larger models.\n\n- The core pipeline (freeze LM -> sample block -> train VQ-VAE -> quantize -> estimate MI -> ...) is convoluted. A figure summarizing the pipeline/steps of the full framework would help a lot. Right now the reader keeps jumping between sections in the main text and appendix to figure out what is the full framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2E5xAdYa4B", "reviewer": "Reviewer_T5EQ", "summary": "The paper examines whether decoder-only language models exhibit planning behavior. It formalizes planning through three abilities: forward-looking horizon (how far future information is encoded), branch awareness (whether multiple plausible continuations are internally represented), and statefulness (how much earlier computations influence later predictions). Using an information-theoretic framework, the authors compress hidden states with a VQ-VAE and compute normalized mutual information between prefix and future representations to quantify these properties. Experiments across grammar, path-finding, and natural text tasks show that planning behavior increases with task complexity, while models remain mostly short-sighted and multi-token prediction provides only minor gains.", "questions": "The questions are related to the weaknesses:\n1. Can the authors provide a more rigorous theoretical analysis linking mutual information to actual planning or reasoning ability, rather than relying on intuitive correlation?\n2. Could future experiments include larger-scale and more capable reasoning models to test whether the observed phenomena generalize beyond small synthetic settings?\n3. Can the proposed framework offer more concrete guidance or insights on how to enhance models’ reasoning or planning capabilities?\n4. Could the authors further analyze how the framework applies to reasoning behaviors under Chain-of-Thought (CoT) or Tree-of-Thought (ToT) prompting settings?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Please find the strengths below:\n1. The paper provides a novel perspective by examining whether decoder-only language models possess planning capabilities from an information-theoretic viewpoint, making the study a fresh and meaningful contribution to the understanding of planning in LMs.\n2. The use of VQ-VAE to discretize hidden states and compute normalized mutual information across layers offers a new theoretical approach for analyzing information flow, which could potentially be applied to other research problems.\n3. The paper presents intuitive figures and maintains a clear logical flow, helping readers grasp the key ideas and results effectively.", "weaknesses": "Please find the weaknesses below:\n1. Intuitively, mutual information captures only statistical correlations rather than causal relationships, and the paper does not provide theoretical justification that higher MI truly reflects stronger planning or reasoning ability; this assumption may thus be overly simplified.\n2. Although the experiments cover diverse tasks, they mainly rely on small-scale models and synthetic settings, without evaluation on more recent large language models. In addition, the VQ-VAE discretization process may introduce information loss or noise, yet the paper lacks a rigorous analysis or quantification of its impact on measurement accuracy.\n3. The multi-token prediction objective yields only marginal improvements, suggesting that the proposed framework, while informative, offers limited practical gains in enhancing model performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yjZgWmczgy", "reviewer": "Reviewer_cc54", "summary": "For interpretability purposes, it is useful to be able to estimate the mutual information (MI) between two distinct blocks of activations in a transformer model. The authors propose to do this by quantizing these activations using a specially trained Vector-Quantized Variational Autoencoder (VQ-VAE), then measuring the MI between the resulting discrete codes. This novel interpretability technique is then used to measure the presence of planning in transformer models in both toy settings and natural language.", "questions": "- How canonical are the three aspects of planning you study here (forward-looking, branch-aware, stateful)? Are you claiming that possessing all three properties is either necessary or sufficient for being a good planner?\n- One advantage of linear probes is that they are cheap and simple to implement. How expensive is the VQ-VAE to train? Does it scale to the kinds of LLMs that may be used in practice (billions of parameters)?\n- What temperature was used for sampling in the experiments?\n- I strongly suggest adding a comparison to linear probe baselines. I.e., wherever $I(E_1(h_1), E_2(h_2))$ is measured, instead fit a linear probe from $h_1$ to $h_2$ and report $R^2$.\n- In the introduction, I recommend fleshing out the discussion of MI vs probes and citing [1].\n- In \"training objectives: next-token vs multi-token\", I recommend citing [2], which shows that cross-token gradients induce forward planning even with a next-token objective.\n- I like the sanity check performed in A.3, but having a discrete domain with support only 0-1 OOMs larger than codebook size is maybe too artificial. Are there continuous distributions for which the theoretical MI can be computed and compared with the estimated MI via VQ-VAE? Alternatively, what if the support is exponentially larger than codebook size (which is what happens in practice, because there are exponentially many possible input sequences)?\n\n[1] Xu et al. \"A theory of usable information under computational constraints\". ICLR 2020.\n\n[2] Wu et al. \"Do language models plan ahead for future tokens?\" COLM 2024.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- As far as I'm aware, the use of VQ-VAE to estimate the mutual information between high-dimensional continuous activations is novel.\n- Several distinct aspects of planning are investigated in both toy settings and natural language.\n- The authors acknowledge that their absolute MI estimates may be misleading, and focus on relative comparisons.\n- The authors provide a basic sanity check to validate their MI via VQ-VAE technique (A.3).", "weaknesses": "1. I am not fully convinced by the authors' philosophical argument that MI via VQ-VAE is more appropriate than existing linear probe techniques. \n\t- To fix notation, say that we care about two model activations $h_1,h_2\\in\\mathbb{R}^d$. We can either train two VQ-VAE encoders $E_1,E_2$ and estimate the mutual information $I(E_1(h_1);E_2(h_2))$, or we can train a probe $\\phi\\colon\\mathbb{R}^d\\to\\mathbb{R}^d$ and measure the $\\ell_2$ loss $||\\phi(h_1)-h_2||_2^2$. \n\t- My understanding is that *both* of these, up to some constant factor, are valid lower bounds for the quantity of interest $I(h_1; h_2)$. The former is a lower bound by the data processing inequality, while the latter is by [1, Prop 1.5].\n\t- The extent to which these bounds are loose is due in the former case to the information discarded by $E_1,E_2$, and in the latter case to the inability of the probe architecture to represent the true conditional probability $P(h_2\\mid h_1)$. It's not clear to me why the former approach is better or more principled than the latter.\n\t- I'm not sure I understand what is meant by the \"confounding effect\" of learned probes mentioned in the introduction. Both the VQ-VAE approach and the probing approach require training separate auxiliary models, and the quality of the lower bound depends on the representational capacity of the auxiliary model in either case.\n\t- I would argue that in the neural network setting, *both* approaches are somewhat philosophically dubious. It is often the case that $h_2$ is a deterministic function of $h_1$ (for example, in the setting of Section 3.1, where $h_1=h_{1:T}^{1:L-1},h_2=h_{T+\\tau}^L$, this is true if sampling temperature is zero). In this case $I(h_1;h_2)=H(h_2)$, telling you nothing about the relationship between $h_1$ and $h_2$.\n2. Experimental results have no comparisons with baselines such as linear probes.\n\t- Since the theoretical/philosophical discussion is somewhat unclear, an empirical comparison would give much stronger evidence for the suitability of MI via VQ-VAE vs probes for interpretability.\n\t- As it stands, the paper applies an unproven method to novel toy settings without established baselines. Without proper baselines, it is hard to be confident that these results are correct or meaningful.\n3. MI between past and future model activations is a fairly weak information-theoretic condition, and is e.g. always present in the zero-temperature autoregressive setting. In my opinion, \"planning\" is a somewhat misleading term for this, as it implies intentionality in the model which may not really be present.\n\n [1] Xu et al. \"A theory of usable information under computational constraints\". ICLR 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3ejWmcjlRn", "reviewer": "Reviewer_rbc9", "summary": "This paper investigates whether decoder-only language models (LMs) engage in internal planning, defined along three axes: forward-looking (planning beyond the next token), branch-aware (considering multiple valid continuations), and stateful (reusing earlier computations). The authors propose an information-theoretic framework that compresses high-dimensional hidden states into discrete codes using a modified Vector-Quantized Variational Autoencoder (VQ-VAE), enabling mutual information (MI) estimation between different parts of the model’s internal computations. They apply this pipeline across synthetic (CFG, path-finding) and natural language (OpenWebText) tasks, finding that planning behavior is task-dependent. Stronger evidence of non-myopic computation is observed in structured reasoning tasks than in syntactic or natural language settings.", "questions": "- I am somewhat confused about the authors’ definitions of planning properties: \"Stateful\" involves reusing earlier computations, but isn’t this also a manifestation of being \"forward-looking\"? Could the authors provide more intuitive examples to clarify the distinction between these two concepts?\n- Regarding the use of VAE for compressing information to obtain distributions:\n  - Could the authors briefly explain why this method is unlikely to introduce external information? Is it because the two positions being compared are compressed by two independent VAEs? Is it possible that the learned encoder and codebook inadvertently encode external information, leading to misleading conclusions?\n  - Could the same MI patterns arise in a model that does not \"plan\" in any meaningful sense?\n- Based on the observed patterns of LM planning across different tasks, could the authors propose any improvements to the LM architecture or algorithms, such as a blueprint for enhancement?\n- If earlier blocks retain \"nontrivial information,\" as claimed, why doesn’t ablating or perturbing them significantly degrade performance? Such a test could potentially strengthen the argument for statefulness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper is well-motivated. Previous studies on the internal mechanisms of LMs often rely on linear probing, which is susceptible to confounders and may lead to unreliable conclusions. The proposed method avoids these confounders, making the analysis of internal planning in LMs more reliable.\n- The use of mutual information over compressed discrete codes provides a principled and scalable approach to analyzing high-dimensional activations.\n- The experimental design covers a thoughtful range of tasks—from highly structured symbolic problems to natural language—allowing for nuanced insights into when and how planning emerges.\n- The validation experiment in Appendix A.3 reasonably demonstrates that the VQ-VAE compression preserves meaningful statistical dependencies.", "weaknesses": "- Overall, the empirical findings are somewhat incremental. For instance, the observation that LMs plan more in path-finding tasks and exhibit myopic behavior in CFG tasks is somewhat expected. Additionally, the proposed framework does not offer a mechanistic explanation of how such planning is implemented, which limits the significance of the findings.\n- In my view, using VAE-based methods also has drawbacks compared to linear probing:\n  - Interpretability is relatively poor. Since the latent space of VAEs is inherently difficult to interpret, even if changes in MI are detected, it remains challenging to explain the underlying mechanisms behind these changes.\n  - The VQ-VAE introduces its own set of hyperparameters and design choices (e.g., codebook size, cosine penalty), and the sensitivity of the results to these choices is not thoroughly explored.\n- The reliance on normalized MI makes it difficult to assess the absolute strength of planning signals. Moreover, the paper does not convincingly demonstrate that the observed MI reflects causal planning rather than passive statistical correlation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)