cs.AISep 21, 2026
SaveFew-Shot Demonstrations Elicit the Use of In-Context World Representations in LLMs
Organizations: The University of Tokyo · Harvard University · Prior Computers
Abstract
Large language models (LLMs), when acting as agents, are expected to take observed data in context, infer the latent state space underlying the world, and leverage it for downstream prediction. However, prior work demonstrated that LLMs struggle to use representations learned in context on a graph tracking task, where the model needs to construct a representation of the graph governing data generation process and use it for subsequent predictions. In this paper, we show that extending this to few-shot settings, where each demonstration is generated from a different world with either the same or different graph topologies, enhances its prediction on 6 models from 4 model families. To understand this improvement, we linearly probe a low-dimensional world representation that encodes graph information in the hidden states. Notably, we find that few-shot demonstrations relocate the world representation and increase its predictive use. Specifically, for each model, these world representations shift in directions nearly orthogonal to their original subspace, and interventions on these representations selectively impair performance more than interventions on other subspaces. Consistent with this insight, we show that few-shot demonstrations with observations from different worlds improve performance on ARC-AGI-1&2, web agent tasks, and Othello. Our findings elucidate the role and internal mechanisms of few-shot demonstrations in in-context world modeling. More broadly, our work advances our understanding of how LLM agents learn from in-context observations and provides implications for their further improvement.
Figures & tables
Figure 1: (a) Conceptual Framework of In-Context World Modeling. An LLM is provided with random walks generated on a graph topology, where each node is assigned a fruit, and instructed to move two steps to the right. We probe the internal representations to determine whether the model has constructed the latent graph. A linearly decodable subspace of these representations is termed a world representation. See Figure 11 for the actual prompt. In this paper, we treat cucumbers as fruits. (b) LLMs Learned to Use World Representation with Few Shot Demonstrations. Few shot demonstrations consisting of random walks, task instructions, and answers from worlds with the same topology but different assignments of fruits to nodes improve task performance. Interventions on the world representation subspace cause larger performance degradation in the few shot setting than in the zero shot setting relative to interventions in random directions. The results are from Qwen3-4B-Base on a Lattice Grid. (c) Few-Shot Demonstrations Shift World Representations. As the number of shots increases, the graph representation shifts orthogonally to the representation subspace. The left shows a 3D PCA of hidden states of the nodes, and the right shows a 2D PCA of mean hidden states over the random walks in the test set. The results are from Figure 6 .
Figure 2: Illustration of Few-Shot Prompt. See Figure 11 for the prompt example.
Figure 3: Few-Shot Performance. Accuracy of three Qwen3 models on the two-step-right prediction task across five graph topologies. Each point reports accuracy on 300 test samples with 0, 1, 2, 4, 8, 16, or 32 few-shot demonstrations. Dotted gray lines indicate random guessing over the candidate nodes for each topology. Lepori et al. (2026) corresponds to the zero-shot case.
Figure 4: Probe Performance and Normalized Dirichlet Energy (DE). Probe AUROC and normalized DE of hidden states across shots and five graph topologies for layer 20 of Qwen3-4B-Base. Gray lines show random feature probe baselines.
Figure 5: (a) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another, averaged over five topologies. (b) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts, averaged over the five topologies. (c) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces, shown separately for each topology. (d) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts, shown separately for each topology. All results are from layer 36 of Qwen3-4B-Base.
Figure 6: (Left) PCA of Shift in Hidden States of a Random Walk. PCA of Lattice Grid at layer 36 of Qwen3-8B-Base (Right) PCA of Shifts in Hidden-State Centroids across Random Walks. PCA of 300 test samples per shot for Lattice Grid at layer 20 of Qwen3-8B-Base. Points represent individual random walks and triangles indicate their means. In each panel, PCA is fit jointly to all activations shown.
Figure 7: Intervention Results. Graph subspace erasure at iteration 20 in three Qwen3 models. Interventions target node subtokens in the query’s random walk at every Transformer block output. The random intervention uses an orthogonal subspace with matched rank and perturbation norm per node. Points show the mean difference in the probability of correct answer token reduction between graph and random interventions, normalized by each query’s unperturbed probability, across 300 test samples. Probabilities are normalized over candidate answer sequence likelihoods. Positive values indicate greater degradation from graph erasure. Horizontal bars show 95% bootstrap confidence intervals obtained by resampling test samples 5000 times.
Figure 8: (a) ARC-AGI and Othello with Different Worlds across Shots. In ARC-AGI, the world is the rule that maps inputs to outputs. In Othello, the world is the board state. (b) Performance across Different Numbers of Shots on ARC-AGI-1&2, Mind2Web, and Othello.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Illustration of World Models in LLMs.
Figure 10: Illustration of Five Graph Topologies.
Figure 15: Few-Shot Performance where Demonstrations are from Different Graph Topologies. Accuracy of three Qwen3 models on the two-step-right prediction task across five graph topologies. The demonstrations are randomly sampled from four topologies other than the query topology, with each topology represented equally when the number of shots is 4, 8, 16, or 32. We report performance separately for each graph topology in the query. Each point reports accuracy on 300 test samples with 0, 1, 2, 4, 8, 16, or 32 few-shot demonstrations. Dotted gray lines indicate random guessing over the candidate nodes for each topology.
Figure 16: Probe Performance and Normalized Dirichlet Energy (DE). Probe AUROC (Top) and normalized DE of hidden states (Bottom) across shots and five graph topologies for layer 20 of three Qwen3 models. Gray lines show random feature probe baselines. Each condition is evaluated on 300 test samples.
Figure 17: Node Distance Geometry Relative to Zero-Shot Setting. Node distance geometry relative to zero-shot setting using probe projected representations (Top) and full hidden states (Bottom) across shots and five graph topologies for layer 20 of three Qwen3 models. Colored lines compare the same query across shots, while gray lines compare different query families with the same topology but different node vocabulary assignments. Lower values indicate more similar relative node geometry. Each condition uses 300 test samples.
Figure 18: Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. Results are from three Qwen3 models on five graph topologies.
Figure 19: Representation Subspace Overlap Results. Overlap between probe subspaces across shot counts. Results are from three Qwen3 models on five graph topologies.
Figure 20: PCA of Shift in Hidden States of a Random Walk Across Different Number of Shots. In each panel, PCA is fit jointly to all activations shown. Results are from three Qwen3 models on five graph topologies.
Figure 21: PCA of Shifts in Hidden-State Centroids across Random Walks Across Different Number of Shots. PCA of 300 test samples per shot. Points represent individual random walks and triangles indicate their means. In each panel, PCA is fit jointly to all activations shown. Results are from three Qwen3 models on five graph topologies.
Figure 22: Angle Between the World Representation Subspace and the Shot Vector. Angles to the corresponding probe subspaces, shown separately for each topology. Results are from three Qwen3 models on five graph topologies.
Figure 23: Angle Between the Shot Vector and the Task Vector. The angle between the two vectors is in degrees from to . Results from Qwen3-4B-Base on five topologies.
Figure 24: Representation Subspace Distance Results. Normalized projection distance between probe subspaces at consecutive tested shot counts, shown separately for each topology. Results are from three Qwen3 models on five graph topologies.
Figure 25: Intervention Results across Intervention Iterations. Plotted points represent the average paired difference in this degradation between the 16-shot and 0-shot settings across 300 samples in test set. The top is iterative nullspace projection targeted to graph structure, and the bottom is random orthogonal subspace intervention matched in rank. For each prompt condition, relative degradation is computed per test sample as the drop in the model’s assigned probability for the correct answer after intervention, normalized by the probability before intervention. Positive values indicate that intervention causes a larger relative performance drop under 16-shot prompting than under 0-shot setting.
Figure 26: Probe Performance during Intervention Iterations. At each cumulative intervention step (4, 8, 12, 16, and 20), a new linear probe is trained on the intervened residual stream activations and evaluated on a held-out test set. Plotted points indicate the average of test AUROC across the layers 2–36 for Qwen3-4B-Base and Qwen3-8B-Base, and layers 2–40 for Qwen3-14B-Base.
Figure 27: Few-Shot Performance. Accuracy of Llama-3.2-3B, Gemma-4-12B, and Ministral-3-8B-Base-2512 on the two-step-right prediction task across five graph topologies. Each point reports accuracy on 300 test samples with 0, 1, 2, 4, 8, 16, or 32 few-shot demonstrations. Dotted gray lines indicate random guessing over the candidate nodes for each topology.
Figure 28: (a) Probe Performance. Probe AUROC. Gray lines show random feature probe baselines. (b) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. (c) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts. (d) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces. (e) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts. Results are from layer 28 of Llama-3.2-3B, layer 48 of Gemma-4-12B, and layer 34 of Ministral-3-8B-Base-2512 on Lattice Grid.
Figure 29: Intervention Results. Graph subspace erasure at iteration 20 in Llama-3.2-3B, Gemma-4-12B, and Ministral-3-8B-Base-2512 on Lattice Grid. Interventions target node subtokens in the final query’s walk at every Transformer block output. The random intervention uses an orthogonal subspace with matched rank and perturbation norm per node. Points show the mean difference in correct answer probability reduction between graph and random interventions, normalized by each query’s unperturbed probability, across 300 test samples. Probabilities are normalized over candidate answer sequence likelihoods. Positive values indicate greater degradation from graph erasure. Horizontal bars show 95% bootstrap confidence intervals obtained by resampling test samples 5000 times. Note that the wide confidence intervals for 0-shot Gemma-4-12B are because of near-zero baseline probabilities in the denominator.
Figure 30: (a) Probe Performance. Probe AUROC. Gray lines show random feature probe baselines. (b) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. (c) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts. (d) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces. (e) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts. Results are from Qwen3-4B-Base on Lattice Grid, where .
Figure 31: (a) Probe Performance. Probe AUROC. Gray lines show random feature probe baselines. (b) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. (c) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts. (d) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces. (e) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts. Results are from Qwen3-4B-Base on Lattice Grid, where .
Figure 32: Intervention Results. Graph subspace erasure at iteration 4 in Qwen3-4B-Base on Lattice Grid, where . Interventions target node subtokens in the final query’s walk at every Transformer block output. The random intervention uses an orthogonal subspace with matched rank and perturbation norm per node. Points show the mean difference in correct answer probability reduction between graph and random interventions, normalized by each query’s unperturbed probability, across 300 test samples. Probabilities are normalized over candidate answer sequence likelihoods. Positive values indicate greater degradation from graph erasure. Horizontal bars show 95% bootstrap confidence intervals obtained by resampling test samples 5000 times.
| Setting | Seed 1 | Seed 2 | Seed 3 | Seed 4 | Seed 5 | Mean Std |
|---|---|---|---|---|---|---|
| L4, 0-Shot | 0.6903 | 0.7072 | 0.7060 | 0.7048 | 0.7013 | |
| L4, 16-Shot | 0.6867 | 0.6824 | 0.7046 | 0.6763 | 0.6941 | |
| L20, 0-Shot | 0.7812 | 0.7761 | 0.7792 | 0.7784 | 0.7777 | |
| L20, 16-Shot | 0.7664 | 0.7573 | 0.7591 | 0.7688 | 0.7722 | |
| L36, 0-Shot | 0.7675 | 0.7508 | 0.7629 | 0.7591 | 0.7656 | |
| L36, 16-Shot | 0.7407 | 0.7359 | 0.7419 | 0.7406 | 0.7412 |
Table 1: Probe Performance. Results from Qwen3-4B-Base evaluated on 300 test samples on Lattice Grid. “L4, 0-Shot” denotes the probe performance at Layer 4 with 0-shot prompting.
| Setting | Seed 1 | Seed 2 | Seed 3 | Seed 4 | Seed 5 | Mean Std |
|---|---|---|---|---|---|---|
| L4, -Shot | 0.6406 | 0.6487 | 0.6525 | 0.6483 | 0.6425 | |
| L4, -Shot | 0.6857 | 0.6821 | 0.7036 | 0.6755 | 0.6942 | |
| L20, -Shot | 0.7191 | 0.7112 | 0.7028 | 0.7136 | 0.7061 | |
| L20, -Shot | 0.7669 | 0.7580 | 0.7592 | 0.7692 | 0.7731 | |
| L36, -Shot | 0.6329 | 0.6727 | 0.6551 | 0.6641 | 0.6538 | |
| L36, -Shot | 0.7428 | 0.7375 | 0.7439 | 0.7436 | 0.7415 |
Table 2: Cross Probe Performance. Results from Qwen3-4B-Base on Lattice Grid. “L4, -Shot” denotes the Layer 4 probe trained on 500 0-shot samples and evaluated on 300 16-shot samples.
| Setting | Seed 1 | Seed 2 | Seed 3 | Seed 4 | Seed 5 | Mean Std |
|---|---|---|---|---|---|---|
| L4, -Shot | 0.6295 | 0.8563 | 0.7353 | 0.7794 | 0.7285 | |
| L4, -Shot | 0.9804 | 0.9728 | 0.9681 | 0.9238 | 0.9928 | |
| L20, -Shot | 0.3329 | 0.4381 | 0.3081 | 0.3230 | 0.2904 | |
| L20, -Shot | 0.7970 | 0.7931 | 0.8709 | 0.8691 | 0.8247 | |
| L36, -Shot | 0.1776 | 0.2669 | 0.1812 | 0.1554 | 0.1601 | |
| L36, -Shot | 0.8654 | 0.8654 | 0.8993 | 0.8901 | 0.8861 |
Table 3: Representation Subspace Overlap. Results from Qwen3-4B-Base on Lattice Grid. “L4, -Shot” denotes the overlap between the rank-2 representation subspaces of probe layers trained on 500 samples.
| Setting | Seed 1 | Seed 2 | Seed 3 | Seed 4 | Seed 5 | Mean Std |
|---|---|---|---|---|---|---|
| 0-Shot | -0.401 | -0.649 | -0.380 | -0.359 | -0.648 | |
| 16-Shot | 1.010 | 1.293 | 1.069 | 1.246 | 1.635 |
Table 4: Mean Relative Degradation Difference after Interventions. Results from Qwen3-4B-Base on the Lattice Grid. Each row reports the mean difference in correct answer probability reduction between graph and random interventions after 20 iterations for the corresponding shot setting, normalized by the corresponding clean probability before interventions. Positive values indicate greater degradation under the graph intervention.
| Setting | Seed 1 | Seed 2 | Seed 3 | Seed 4 | Seed 5 | Mean Std |
|---|---|---|---|---|---|---|
| L4, 0-Shot | 0.5862 | 0.6337 | 0.5985 | 0.6015 | 0.5806 | |
| L4, 16-Shot | 0.5798 | 0.5897 | 0.5543 | 0.5758 | 0.5582 | |
| L20, 0-Shot | 0.7158 | 0.6900 | 0.6920 | 0.7095 | 0.7000 | |
| L20, 16-Shot | 0.6658 | 0.6783 | 0.6947 | 0.6891 | 0.6845 | |
| L36, 0-Shot | 0.6905 | 0.7208 | 0.7263 | 0.7059 | 0.7109 | |
| L36, 16-Shot | 0.6795 | 0.7056 | 0.6901 | 0.6915 | 0.6914 |
Table 5: Probe Performance after Graph Interventions. Results from Qwen3-4B-Base evaluated on 300 test samples on Lattice Grid after graph intervention for 20 iterations. “L4, 0-Shot” denotes the probe performance at Layer 4 with 0-shot prompting.