Do multimodal models imagine electric sheep?

TL;DR

A Qwen3.5 9B VLM fine-tuned to solve twelve open-loop spatial puzzles (tangram, jigsaw, sokoban, Bloxorz, rush hour, 3D mental rotation, etc.) by behavior-cloning the full action sequence from a single initial image develops emergent mental imagery — a separately trained visual decoder can reconstruct intermediate world states from the LLM’s hidden activations at action boundaries, despite the model receiving no visual supervision past the first frame. Sharpening that latent imagery with an auxiliary visual-prediction loss, or surfacing it as 16-64 discrete visual tokens interleaved with action tokens in the chain of thought, lifts the average solve rate from 83% (action-only) to 89.5%, with the largest gains on reasoning-heavy tasks (jigsaw, 3D mental rotation). Visual CoT beats a matched JSON scene-graph textual CoT by ~10 points, so the win is the visual modality itself, not test-time compute.

Key claims

Off-the-shelf Qwen3.5 9B solves 13% of the twelve puzzles; action-only open-loop behavior cloning lifts the average solve rate to 83% despite the model never observing intermediate states during 20+-step rollouts [§4.2, Fig. 5].
Mental imagery emerges without visual supervision: a four-layer MAE-style decoder trained on the frozen behavior-cloned LLM’s last-action-token hidden state can reconstruct meaningful intermediate geometry (tangram grid layout, character shape) [§4.2, Fig. 2].
Mental-image fidelity correlates with solve rate across tasks; the strongest gains from removing the stop-gradient on the visual probe (“visual supervision”) occur on initially-low-solve-rate puzzles like 3D mental rotation and 3D shape matching [§4.2, Fig. 6].
Visual chain-of-thought — interleaving 16 or 64 discrete VQ-VAE tokens of the post-action state between action blocks — achieves 89.5% / 89.0% average solve rate, vs 83% for action-only and ~86% for the auxiliary visual loss [§4.3, Fig. 5].
A matched textual CoT baseline that emits structured JSON scene graphs per step (same Qwen3.5 9B, same behavior cloning) reaches only 78.7% — visual CoT beats it by ~10 points and the gap is largest on 3D shape matching, 3D mental rotation, jigsaw, and Bloxorz [§4.3].
16 visual tokens per state (32×32 grid) and 64 tokens per state (64×64 grid) have similar averages but complementary per-game profiles: 16-token is better on long-horizon games (sokoban, tangram, Bloxorz), 64-token is better on resolution-hungry tasks (jigsaw, 3D shape matching, 3D mental rotation) [§4.3].
Auxiliary visual supervision improves average solve rate by 3.4 points over action-only with no test-time overhead (the visual head is dropped at inference) [§4.3].
Fine-tuning a unified understanding-and-generation model (BAGEL) as the base did not reach reasonable performance under the same recipe [§4.3, appendix].

Method

The setup is open-loop behavior cloning: given an initial puzzle image and a natural-language instruction, the VLM must emit the full action sequence a_0 … a_{T-1} (plus optional final answer for certificate variants) from the initial image alone, with no observation of intermediate states. Twelve puzzles — tangram, jigsaw, sokoban, Bloxorz, rush hour, anagram, plus 2D/3D shape matching and 2D/3D mental rotation certificate variants — are implemented as Rust game engines with expert solvers (BFS, Algorithm X, closed-form) that emit ground-truth action sequences. Base model is Qwen3.5 9B; each variant is fine-tuned for 72 hours on one node of 8 B200 GPUs.

Three training variants share the action-prediction loss but differ in how internal visual state is shaped: (a) action-only — pure behavior cloning, no visual signal past the first frame; (b) visual supervision — an auxiliary head reads hidden states at action boundaries and predicts FSQ codes of the next state, gradients flow back into the LLM, head is dropped at inference; (c) visual tokens — the model’s output vocabulary is extended with LlamaGen VQ-VAE tokens and it autoregresses interleaved <action>…</action><image>…</image> blocks, with the image block being either 16 or 64 tokens. To probe (a) for emergent imagery without changing its training, the authors freeze the action-only LLM and train a four-layer MAE-style decoder that takes only the hidden state of the last token of each action and reconstructs the post-action visual state.

Results

Average solve rate over twelve puzzles, 200 held-out instances each, greedy decoding via vLLM [§4.3, Fig. 5]:

Stock Qwen3.5 9B: 13%
Textual CoT (JSON scene graph): 78.7%
Action-only behavior cloning: 83%
Visual supervision (auxiliary loss, head dropped at inference): ~86%
Visual tokens, 64 per state: 89.0%
Visual tokens, 16 per state: 89.5%

The visual-CoT lead over textual-CoT (~~+10 points) localizes to tasks where state is hard to verbalize: 3D shape matching, 3D mental rotation, jigsaw, Bloxorz. The visual-supervision-vs-action-only gain (~~+3.4 points) clusters on the same reasoning-heavy puzzles and on anagram. Probing the action-only model with the frozen visual decoder shows recognizable post-action geometry on tangram and jigsaw rollouts (Fig. 2), and the per-task correlation between reconstruction quality and solve rate is positive on every puzzle except anagram, where the base model has near-zero solve rate but artifactually good reconstructions due to sparse foreground pixels [§4.2, Appendix D.3]. Total compute for a full sweep is 9,000 GPU hours.

Why it’s interesting

This is the cleanest filed evidence that a behavior-cloned VLM builds an internal world model of the visual state as a byproduct of learning to plan, and that this latent imagery can be either (i) probed and visualized post-hoc with a frozen decoder, (ii) sharpened with an auxiliary loss at zero inference cost, or (iii) surfaced as a compact 16-token-per-step visual CoT. For Thinking with Modalities this is a sharper instance of substrate (a) — tagged latent skills inside an autoregressive CoT — than Video-Thinker: Sparking "Thinking with Videos" via Reinforcement Learning (Video-Thinker), because here the inserted tokens are the visual state rather than tags that invoke perception ops, and the matched textual-CoT control rules out test-time compute as the explanation. For VLM Perception Failures it is the positive counterpart to Hidden in plain sight: VLMs overlook their visual representations (Fu et al.’s finding that visual features are intact at every LLM layer but overridden at the readout): here the visual features are not just intact, they are predictive of an unseen future state, and a tiny auxiliary loss is enough to keep the readout honest. The compact-token result (16 tokens per state is competitive with 64) is the practically useful number — minimal sequence-length overhead, no architectural surgery.