WM-VLM: Probing Internal World Models for Interleaved Visual-Textual Reasoning

AI summary

Overview

Research area: Computer Vision / multimodal reasoning — specifically vision-language models (VLMs) and spatial intelligence.

Technical level: Advanced. The paper assumes familiarity with transformer architectures, mixture-of-experts/Mixture-of-Transformers routing, rectified flow / flow-matching generative objectives, and VLM fine-tuning conventions.

Scope (one sentence): The paper introduces WM-VLM, a Qwen2.5-VL-7B-Instruct backbone augmented with a shallow "generation branch" that predicts intermediate visual states, and tests on programmatically generated 2D and 3D mental-rotation datasets whether the model actually uses those generated states to answer spatial questions.

What This Paper Is About

Conventional VLMs reason in text, even for problems that are fundamentally spatial — like mentally rotating a shape. Text is a poor carrier for the visual information such rotations require. The authors ask whether a VLM can be given an internal world model: a component that predicts the next visual state after a "mental action," which the model then conditions on to continue reasoning. The goal is to test both whether the model can generate accurate intermediate visual states and whether it causally relies on them.

Key Contributions

1. WM-VLM architecture. A VLM with a lightweight Mixture-of-Transformers ("light MoT") design: a full understanding branch (vision encoder + language decoder, initialized from Qwen2.5-VL-7B-Instruct) plus a shallow generation branch of k layers (default k = 4) aligned with k consecutive layers of the understanding branch. Text and clean image tokens route to the understanding branch; noisy image tokens route to the generation branch; attention is global across both.

2. Two-stage training curriculum. Stage 1 freezes the entire understanding branch and trains only the generation branch with a rectified flow objective to predict the next visual state. Stage 2 keeps the vision encoder frozen and jointly trains the language decoder and generation branch with cross-entropy (Stage 1 uses α=1, β=0; Stage 2 uses α=0, β=1 in the combined loss).

3. Controlled, verifiable diagnostic datasets. Tetris-2D and Tetris-3D are programmatically constructed so that each sample has a question image, a trace of textual mental actions interleaved with visual states, and a verifiable multiple-choice answer. This isolates world modeling from action planning — unlike planning benchmarks where state prediction and planning are entangled.

4. Causal probing of the generated states. Beyond accuracy, the paper measures embedding cosine similarity and top-1 retrieval against ground-truth visual embeddings, and runs corruption ablations (all-zero, randomized, shuffled tokens) to test whether the visual states are genuinely used.

Main Findings

• Large accuracy gains over the SFT backbone. WM-VLM reaches 87.50 vs. 48.25 on Tetris-2D-ID (+39.25), 75.00 vs. 46.20 on Tetris-2D-OOD (+28.80), 91.00 vs. 71.75 on Tetris-3D-SC-ID (+19.25), 66.25 vs. 44.50 on Tetris-3D-C-ID (+21.75), and 71.80 vs. 55.50 on Tetris-3D-OOD (+16.30). The overall stated range is 16.30–39.25 percentage points.

• A sharp performance transition at ~25k Stage-1 steps. Figure 1 shows accuracy jumping suddenly around step 25k — the same point at which cosine similarity between generated and ground-truth visual embeddings rises faster. The interpretation is that the world model only helps once generation quality crosses a threshold. Qwen2.5-VL-7B-Instruct SFT tracks WM-VLM before that point but then plateaus.

• Baselines trail substantially. On Tetris-2D ID/OOD: Qwen2.5-VL-7B-Instruct 22.00/23.80, LatentUM 29.25/23.00, Mirage 47.75/43.00, ThinkMorph 22.50/21.60, WM-VLM 87.50/75.00. Fine-tuned BAGEL (ThinkMorph) emits no images on Tetris-3D at all, scoring 0.00 across all three Tetris-3D splits (marked with an asterisk in the paper).

• Corrupted visual states sharply degrade accuracy. On Tetris-2D-ID/OOD, replacing generated tokens with all-zero embeddings drops accuracy to 14.00/15.40 (−73.50/−59.60), randomized embeddings to 23.50/20.80 (−64.00/−54.20), and shuffling to 81.50/72.60 (−6.00/−2.40). Tetris-3D shows the same pattern, e.g. Tetris-3D-SC-ID falls from 91.00 to 15.50 with all-zero embeddings. The paper explains shuffling's smaller effect by noting that visual tokens attend bidirectionally in the generation branch and already carry positional information.

• Retrieval quality correlates with answer accuracy, but weakly. Tetris-2D-ID has cosine 0.9347 and top-1 retrieval 96.50 with φ = 0.298; Tetris-3D-OOD has cosine 0.7746 and top-1 retrieval 13.40 with φ = 0.129. Accuracy when retrieval is correct (Accuracy-✓) is much higher than when it is wrong — e.g. Tetris-3D-C-ID 92.77 vs. 59.31.

• Pixels are worse than continuous visual tokens. Adding an MLP to reconstruct pixels and training with pixel MSE yields 49.25/46.60 (all-white), 48.25/46.20 (pixel shuffle), and 47.25/46.20 (19×19 patches) on Tetris-2D ID/OOD — far below the 87.50/75.00 of token-based reasoning at the same compute budget.

• Single-stage joint training from scratch fails. With α=0.1 and β=0.9, the jointly trained model stays near random chance throughout training, while both Qwen2.5-VL-7B-Instruct SFT and WM-VLM exceed 40% accuracy within a comparable number of steps. The authors argue the staged curriculum is critical.

• Loss weighting matters. The best Stage-2 setting is α=0.1, β=0.9 (cross-entropy dominant), even though a larger flow-matching weight improves generated token quality. Better token quality does not automatically translate to better downstream accuracy.

• Token count has opposite effects on ID and OOD. Default is 361 visual tokens (19×19) for 512×512 reasoning images. Fewer tokens give lower MSE and higher cosine similarity, but ID accuracy rises with more tokens while OOD accuracy peaks at 4 (2×2) tokens.

• A shallow generation branch suffices. For a 1-layer generation branch, aligning with the top-most layers (reported as 23–27) performs best, with an unexplained spike at layer 3. With 4 layers, placement differences shrink. Using a full-depth generation branch (the original MoT setting) gives no significant gain.

• Random-guess reference is 25% (four-option multiple choice), shown as a dashed line in Figure 1.

• Stronger VLMs can solve these tasks textually via coordinates, but at high token cost — the paper raises this as an open question about whether continuous visual embeddings are better substrates than text for visual reasoning.

Methodology in Plain English

Data generation. The authors write programs that build shapes from a random number of atomic squares (2D) or cubes (3D), apply rotations to produce different views, and then pose a question: here are two views of one shape showing a rotation; here is a new shape; which view results from applying the same rotation? Because the generator knows the answer, every intermediate visual state and every final answer is checkable. Tetris-2D has 4k training samples with 400 ID and 500 OOD test cases. Tetris-3D has 16k training samples, with 400 easy in-domain (SC-ID), 400 hard in-domain (C-ID, seen cube count but unseen shapes), and 500 OOD cases (unseen shapes or counts). They also train on ThinkMorph-SpatialNavigation using the VSP maze navigation benchmark, where the model generates a visual state of the agent's position at each step and plans a move.

Model. Start from Qwen2.5-VL-7B-Instruct. Add a small parallel "imagination" stack of 4 transformer layers whose weights are copied from a matching block of the main model. When the model writes a mental action ("rotate 90° clockwise"), noisy image tokens are routed to this small stack, which denoises them into a visual embedding that takes the place of a real image embedding in the reasoning history.

Training. Stage 1: freeze the whole base model, teach the small stack to predict the next visual embedding, supervised against the vision encoder's embedding of the real intermediate image. Stage 2: unfreeze the language decoder and train both parts together, with the reasoning branch consuming the generated tokens rather than ground truth (an online setting), so the model learns to cope with its own imperfect predictions. Training runs on 8 H200 GPUs.

Probing. Because ground-truth intermediate images exist at test time, the authors can compare generated embeddings to true ones (cosine similarity, top-1 retrieval) and can deliberately corrupt the generated tokens to see whether accuracy collapses — a direct test of whether the visual states are load-bearing or decorative.

Why This Matters

Impact on research. Recent work (the paper cites Viveiros et al., 2026) has questioned whether interleaved visual-textual reasoning models causally depend on their visual intermediates at all, pointing to latent bypass, representation collapse at inference, and an oracle-generation gap. This paper directly addresses that critique by building a setting where visual states are verifiably necessary, and by reporting corruption ablations that show dependence rather than assuming it. The finding that a 4-layer generation branch matches a full-depth one, and that staged training is essential while joint training fails, are practical design signals for anyone building multimodal reasoning systems.

Real-world applications:

• Robotics and embodied agents — an agent that can imagine the visual consequence of a candidate action before committing, without querying an external simulator.
• Assistive and industrial spatial tasks — packing, assembly guidance, and navigation where a user needs to know what a scene looks like after a transformation.
• Design and CAD tooling — predicting intermediate states of a 3D object under rotation or manipulation.
• Diagrammatic and STEM tutoring — systems that reason about geometry and physical transformations rather than only describing them in words.

Industry relevance. The approach bolts a small generative branch onto an existing pretrained VLM rather than training a unified model from scratch, and the results suggest only 4 extra layers are needed. That is attractive for teams that already have a VLM in production and want to add spatial imagination without a full retrain. The baseline comparisons against LatentUM, Mirage, and ThinkMorph/BAGEL provide a benchmark landscape for this design space.

Future Directions

• Caveats and limitations are not fully reported in the available content. Effect sizes vary widely: the φ correlations between retrieval correctness and answer correctness are modest (0.104 to 0.399), and on Tetris-2D-OOD and Tetris-3D-OOD, Accuracy-✗ remains high (70.21 and 69.52), implying the model often answers correctly without retrieving the right visual state. Reconciling this with the strong corruption-ablation results is an open question.

• Why does joint training fail, and how general is the curriculum requirement? The paper reports the failure under one setting (α=0.1, β=0.9) but does not establish whether other schedules or initializations could make single-stage training work.

• Layer placement and the ~25k-step transition are unexplained. The spike at layer 3, the preference for the top-most layer block, and the sudden jump in accuracy at 25k steps are documented but not mechanistically accounted for.

• Does this transfer beyond synthetic Tetris? The tasks are programmatically generated, controlled, and use a fixed answer format. Whether internal world models help on natural-image spatial reasoning, real embodied control, or open-ended interleaved reasoning is untested. The paper also raises the question of whether continuous visual embeddings are a fundamentally better substrate than coordinates-and-text reasoning, which a strong textual model can already exploit.

• Results on ThinkMorph-SpatialNavigation / VSP-Nav are referenced (Table 7) but the numbers are not present in the provided content, so the navigation setting's outcome cannot be summarized here.

Target Audience

Researchers and engineers working on multimodal LLMs, visual reasoning, and world models — particularly those interested in interleaved visual-textual chain-of-thought, latent visual reasoning, and mixture-of-transformers architectures. Also relevant to embodied-AI and robotics practitioners considering internal prediction versus external simulation. The paper is not beginner-friendly: it presumes comfort with flow matching, transformer routing, and VLM fine-tuning. Readers looking for a critical empirical treatment of whether latent visual tokens are actually used will find the ablation design and the retrieval/corruption analysis the most distinctive part.

Authors’ abstract

Humans often solve spatial problems by mentally simulating visual transformations. In contrast, conventional vision-language models (VLMs) reason primarily through language. We investigate whether VLMs can solve spatial problems by reasoning with both text and generated visual states. To this end, we introduce WM-VLM, which equips a pretrained VLM with a lightweight world model branch for generating intermediate visual states. Our two-stage training first teaches the model to generate the next visual state and then to use that state for reasoning. We programmatically construct spatial reasoning tasks with verifiable intermediate visual states. These tasks allow us to evaluate how well the model generates visual states and how much it relies on them to answer the question. On 2D and 3D mental rotation tasks, WM-VLM consistently outperforms the supervised fine-tuned backbone, with gains of up to 39.25 percentage points. Ablations suggest that these gains depend on the generated visual states, as removing or corrupting them sharply reduces performance. Together, these results suggest that internal world models offer a promising path toward VLMs that reason in both language and visual space. Code is publicly available at: https://github.com/yuh-zha/WM-VLM