PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuheng Zha, Yilei Wang, Qiyue Gao, Junrong Chen, Yujia Wu, Zhengfeng Lai, Zhengzhong Liu, Eric P. Xing

arXiv:2609.34826v1cs.CV

TL;DR

VLMs often reason primarily through language, while spatial problems can require mentally simulating visual transformations. WM-VLM adds a world-model branch and staged training so a VLM generates intermediate visual states and reasons with them. Across Tetris-2D and Tetris-3D, it outperforms the fine-tuned backbone by 16.30–39.25 percentage points.

  • Problem

    VLMs may lack effective internal world models for predicting future visual states, and text-based intermediate reasoning may fail to preserve visual information needed for spatial problems.

    Source: §1 · p. 1, §4.2 · p. 7

  • Method

    WM-VLM equips a VLM with an internal world model and uses two-stage training to generate intermediate visual states before reasoning with its predictions.

    Source: §1 · p. 2

  • Results

    16.30–39.25 percentage points: WM-VLM outperforms the fine-tuned backbone across Tetris-2D and Tetris-3D.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results suggest that internal world models may enable VLMs to reason jointly in language and visual space.

    Source: §5 · p. 11

    Under the tested training setup, jointly learning visual-token prediction and reasoning from scratch remains near random chance.

    Source: §4.2 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

VLMs have persistent spatial-reasoning weaknesses, and text-only intermediate reasoning may fail to preserve visual information needed for spatial problems. WM-VLM adds an internal world model and staged training, and outperforms the fine-tuned backbone across the reported Tetris tasks.

1 Introduction

  • VLMs have persistent spatial-reasoning weaknesses, while text-only intermediate reasoning may fail to preserve visual information needed for spatial problems.The passage contrasts spatial reasoning with mathematical and chart-based tasks that can often be solved primarily through language.Source: §1 · p. 1
  • Prior approaches manipulate existing images, interleave generated states often in planning tasks, or delegate state prediction to external world models.These approaches leave open whether a VLM can learn to predict its own intermediate visual states to improve spatial reasoning.Source: §1 · p. 1
  • WM-VLM adds a shallow generation branch and trains first to generate intermediate visual states, then to reason with its predictions.Programmatically constructed Tetris-2D and Tetris-3D datasets provide explicit, verifiable intermediate states for separately evaluating state prediction and downstream reasoning.Source: §1 · p. 2
  • 16.30–39.25 percentage points: WM-VLM outperforms the fine-tuned backbone across Tetris-2D and Tetris-3D.The results also report that replacing or corrupting generated visual states substantially reduces performance.Source: §1 · p. 2

2 Related Work

Prior work explores visual intermediates for reasoning, including image operations, external world models, and generated visual representations. The paper situates its approach against open questions about whether models actually depend on those intermediates.

Visual Reasoning with VLMs

  • Visuospatial tasks involving transformation, localization, or state tracking may benefit more from visual intermediates or pixel-space operations than from purely textual reasoning.Prior VLM reasoning evaluations often focus on STEM, charts, or common-sense benchmarks.Source: §2 · p. 2

Interleaved Visual-Textual Reasoning

  • WM-VLM alternates textual mental actions with generated visual states before producing its final answer.This is the interleaved reasoning process illustrated in Figure 2.Source: Figure 2 · p. 3
  • Interleaved visual-textual methods use edited input images, external world models, generated reasoning images, or continuous and discrete visual latents.The cited analyses question whether models causally depend on visual intermediates, reporting that removing or corrupting latent tokens often has little effect.Source: §2 · p. 3

3 Method

The method pairs verifiable spatial-reasoning tasks with a VLM that alternates textual actions and predicted visual states. Training first teaches visual-state prediction, then teaches reasoning from the generated states.

3.1 Task Generation

  • Programmatically constructed 2D and 3D datasets provide interleaved action-state traces and verifiable multiple-choice questions requiring prediction of new visual states.The dataset design targets tasks with verifiable intermediate states as well as final answers.Source: §3.1 · p. 3

3.2 Interleaved Visual-Textual Reasoning with World Model

  • At each reasoning step, the policy produces a mental action, the world model predicts the resulting visual embedding, and the updated history conditions later reasoning and the final answer.Training uses the encoded intermediate image as the target embedding; inference generates the embedding directly.Source: §3.2 · p. 4

3.3 VLM with An Internal World Model

  • The light MoT design retains the pretrained understanding branch and aligns a shallow generation branch with consecutive understanding layers, routing noisy visual tokens to generation and text and clean visual tokens to understanding.The default generation branch uses four layers, and attention is global across tokens from both branches.Source: Figure 3 · p. 4, §3.3 · pp. 4–5

3.4 Two-stage Training

  • Two-stage training first trains the generation branch to predict the next state while freezing the understanding branch, then trains the decoder and generation branch to use generated visual latents for reasoning.The generation branch uses a rectified-flow objective, while the understanding branch uses cross-entropy; the vision encoder remains frozen in Stage 2.Source: §3.4 · p. 5

4 Experiments

The experiments evaluate WM-VLM on programmatically constructed spatial-rotation tasks and compare it with fine-tuned baselines and other reasoning methods. They also test when generated visual states help and how training and architecture choices affect performance.

4.1.1 Data Construction

  • Programmatically generated Tetris-2D and Tetris-3D tasks test applying a demonstrated rotation to a new shape, with in-distribution and out-of-distribution test cases.The experiments also use ThinkMorph-SpatialNavigation, where the model predicts agent states and plans paths around ice holes.Source: §4.1.1 · pp. 5–6

4.1.2 Model and Training Details

  • Training freezes the understanding branch while the generation branch learns, then jointly trains both branches with generated visual tokens and a frozen vision encoder.The understanding branch starts from Qwen2.5-VL-7B-Instruct; the study also compares against its supervised fine-tuned baseline and LatentUM, Mirage, and ThinkMorph.Source: §4.1.2 · pp. 6–7

4.2 Results and Analysis

  • 39.25 points on Tetris-2D-ID and 28.8 on Tetris-2D-OOD; WM-VLM also gains 19.25, 21.75, and 16.3 points on the three Tetris-3D sets over Qwen2.5-VL-7B-Instruct SFT.The paper also reports that WM-VLM outperforms other visual-latent and interleaved visual-textual reasoning methods on Tetris-2D and Tetris-3D.Source: §4.2 · p. 7
  • Correct visual-token retrieval correlates positively with answer accuracy, while replacing generated tokens with zero or random embeddings substantially reduces accuracy.Shuffling has a smaller impact, which the authors relate to positional information and bidirectional attention among visual tokens.Source: Table 2 · p. 7, §4.2 · p. 8
  • Fewer visual tokens improve embedding MSE and cosine similarity, but ID accuracy rises with token count while OOD accuracy peaks at four tokens.The study varies token count by resizing the reasoning images.Source: Figure 4 · p. 8, §4.2 · p. 9
  • Under the tested setup, single-stage training from scratch stays near chance, while Stage 2 performs best with α = 0.1 and β = 0.9.That loss ratio gives cross-entropy the larger weight; increasing the flow-matching weight improves visual-token quality but does not necessarily improve downstream performance.Source: §4.2 · pp. 9–10, Figure 6 · p. 9, Figure 5 · p. 9
  • With a one-layer generation branch, placement at layers 23–27 performs best; with four layers, placement differences shrink, and matching the understanding-branch depth adds no significant gain.The authors attribute the smaller placement gap with four layers to added capacity mitigating information loss across understanding-branch locations.Source: §4.2 · pp. 10–11, Figure 7 · p. 10, Figure 8 · p. 10
  • At the same compute budget, reasoning with generated visual tokens outperforms reasoning with generated pixels.The pixel-generation model reconstructs pixels from visual tokens using an added MLP and pixel MSE loss.Source: §4.2 · p. 11

5 Conclusion

The conclusion presents WM-VLM as a VLM with an internal world model for visual spatial reasoning, evaluated using datasets with verifiable intermediate states. It reports improved spatial reasoning and reliance on generated visual tokens, with design ablations identifying several important choices.

5 Conclusion

  • WM-VLM improves spatial reasoning performance and relies on generated visual tokens to produce final answers.The conclusion frames internal world models as a possible route to reasoning jointly in language and visual space.Source: §5 · p. 11
  • The study constructs datasets with verifiable intermediate visual states to examine how world modeling affects interleaved visual-textual reasoning.Source: §5 · p. 11
  • Ablations identify visual-token count, cross-entropy and flow-matching loss balance, and two-stage training as important design choices.Source: §5 · p. 11

A Implementation Details

The implementation details describe modality-specific attention projections and outputs, then give the inference procedure for generating visual states through rectified-flow sampling.

A.1 VLM with an internal world model

  • The per-token query, key, and value projections use learnable weight matrices for the token’s modality.Source: §A.1 · p. 15
  • The section introduces an attention-output expression for token x_i and identifies O as the modality-specific learnable output-projection matrix.Source: §A.1 · p. 15

A.2 Rectified Flow Sampling

  • At inference, the model integrates a predicted velocity field from noise toward a clean latent using K Euler steps, then appends the estimated visual state with T_j to guide later reasoning.The state is conditioned on c_j = (H_j, T_j).Source: §A.2 · p. 15

A.3 Training Hyperparameters

The training-hyperparameter section points to Table 8 and reports that all training tasks use 8 H200 GPUs.

A.3 Training Hyperparameters

  • Table 8 contains the training hyperparameters.Source: §A.3 · p. 15
  • All training tasks are conducted on 8 H200 GPUs.Source: §A.3 · p. 15
  • The reported 8-H200-GPU setup applies across all training tasks.Source: §A.3 · p. 15

B Data Construction

The data-construction section uses programs to generate the interleaved visual-textual reasoning dataset and adapts the 3D construction for 2D rotation.

B Data Construction

  • Programs generate the interleaved visual-textual reasoning dataset.Source: §B.1 · p. 15
  • The 2D rotation dataset is constructed in a similar way to the 3D dataset.Source: §B.1 · p. 15
  • The 2D rotation construction uses squares in place of cubes.Source: §B.1 · p. 15

B.2 Data Samples

The Tetris datasets pair mental rotation questions with textual actions and the resulting visual states, using different state-recording granularity in 2D and 3D.

B.2 Data Samples

  • Tetris-2D stores a visual state after the complete rotation sequence, while Tetris-3D stores one after each atomic rotation.The examples pair a question image and textual actions with visual states and a final textual answer.Source: §B.2 · p. 15, Figure 9 · p. 16
  • Tetris-3D construction enumerates connected shapes of 4–7 cubes and removes shapes equivalent under proper 3D rotations.The algorithm uses axes X, Y, and Z and rotation angles of 90°, 180°, or 270°.Source: §B.2 · p. 17
  • Mirror-related shapes stay in the same seen-versus-held-out split; training and ID evaluation then use disjoint configurations of seen shapes.ID evaluation uses unseen configurations of seen shapes, while OOD evaluation uses held-out shapes.Source: §B.2 · p. 17
  • For each selected configuration, the algorithm rotates a reference shape with a different cube count and pairs the correct pose with three distractors.The reference shapes, query, and four answer options are rendered together, with option positions randomized.Source: §B.2 · p. 17

B.3 Comparison between Different VLMs

The comparison evaluates VLMs with different reasoning inputs on 2D mental rotation. Stronger textual reasoners can solve by converting shapes to coordinates, but the authors describe this as a lossy, non-scalable visual reasoning approach.

B.3 Comparison between Different VLMs

  • The comparison tests different VLMs and reasoning types on 2D mental rotation, with results reported in Table 6.The table caption describes settings with question input, reasoning text, and intermediate reasoning images.Source: §B.3 · p. 17
  • Stronger VLMs can convert a shape into coordinates, rotate those coordinates through textual reasoning, and derive an answer.The passage describes this behavior for 2D or 3D shapes.Source: §B.3 · p. 17
  • The authors argue that coordinate-based textual reasoning is not a general, scalable visual reasoning approach because language is lossy.Source: §B.3 · p. 17

C Additional Experiment Results

Additional experiments examine branch initialization, generation depth, and visual planning; the paper also identifies interleaved embodied data as a future need for pretraining.

C Additional Experiment Results

  • Neither shared-weight nor random initialization shows a clear downstream advantage across the three tested generation-branch placement settings.Source: C Additional Experiment Results · p. 17, Figure 11 · p. 21
  • Increasing generation-branch depth from 8 to 24 generally improves visual-state reconstruction but does not consistently improve downstream accuracy.The result suggests better reconstruction alone need not translate into better reasoning performance.Source: Figure 10 · p. 19, C Additional Experiment Results · p. 21
  • WM-VLM shows competitive performance against baselines on visual planning, while BAGEL performs best there; on Tetris, BAGEL underperforms WM-VLM and can collapse during training.The authors suggest BAGEL's visual-planning advantage may reflect pretraining on large-scale interleaved data.Source: C Additional Experiment Results · pp. 17, 21

D Future Work

  • Pretraining with an internal world model requires interleaved visual-textual reasoning data, which existing datasets primarily source from static materials such as webpages.The authors propose embodied trajectories pairing actions with resulting observations as future work.Source: D Future Work · p. 21

Related papers

Loading 2609.34826v1…