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

WM-VLM. Probing Internal World Models for Interleaved Visual-Textual Reasoning. By Yuheng Zha and seven others, from UC San Diego, MBZUAI, and one other institution. Published on the twenty-eighth of September. Its code repository is public and real: about twenty-five source files.

People often solve spatial problems by mentally simulating visual transformations, turning an object over in the mind's eye. Conventional vision-language models, by contrast, reason primarily through language. The question investigated here is whether a VLM can solve spatial problems by reasoning with both text and visual states it generates itself, and the evidence suggests that it can.

To test it, the authors build WM-VLM, which equips a pretrained VLM with a lightweight world model branch for generating intermediate visual states. Training has two stages. The model first learns to generate the next visual state, and then learns to use that state for reasoning. The tasks are spatial reasoning problems constructed programmatically, so their intermediate visual states are verifiable. That makes two separate measurements possible: how well the model generates visual states, and how much it relies on them to answer the question.

On mental rotation tasks in two and in three dimensions, WM-VLM consistently outperforms the supervised fine-tuned backbone, with gains of up to thirty-nine point two five percentage points. The ablations are what tie the gain to the visual states. Removing the generated states, or corrupting them, sharply reduces performance, which suggests the gains depend on them.

The scope is specific. The tasks are programmatically constructed spatial problems, and the reported gains are on two-dimensional and three-dimensional mental rotation. The comparison is against the model's own supervised fine-tuned backbone. The ablations suggest the dependence on generated states rather than proving it, and the broader conclusion is framed as a promising path rather than a settled one.

The direction the authors draw from this is that internal world models offer a promising path toward vision-language models that reason in both language and visual space. Here, in a verifiable setting, a model that pictures the next state before answering does better than one that only talks its way through.