Imaginative Perception Tokens Enhance Spatial Reasoning in Multimodal Language Models

Imaginative Perception Tokens: Teaching VLMs to Mentally Simulate Unseen Spatial Views

Current vision-language models (VLMs) excel at many tasks. They continue to struggle with spatial reasoning. This happens specifically when key information is not directly observable in the input. Most spatial questions require "imaginative perception" (the ability to simulate unobserved views). This includes tracing a trajectory through an occluded space or integrating partial views into a map.

The field has largely relied on two approaches. Some use intermediate visual representations like depth maps or bounding boxes. Others use textual chain-of-thought (CoT) (a method of generating step-by-step text reasoning). Both often fail. They either refine information already visible or force 3D geometry through the narrow bottleneck of natural language. This paper finds that forcing spatial reasoning through text can be detrimental. This suggests a fundamental modality mismatch. Instead, the authors propose training models to "imagine" missing spatial structure using specialized visual tokens.

The Problem

The status quo fails because spatial reasoning is often a constructive process. Standard VLMs are good at reading relations from an observed view. They break when asked to adopt a human-centered viewpoint or aggregate multiple views. As shown in, tasks like Perspective Taking (PET), Path Tracing (PT), and Multiview Counting (MVC) require building a mental representation.

This representation is not present in the source image.

Previous attempts used intermediate "visual thoughts." These typically focus on refining existing structure. They do not predict unobserved structure. Researchers also try to use textual CoT to bridge this gap. However, serializing viewpoint transformations and geometric constraints into text is error-prone. The authors demonstrate that forcing a model to describe a 3D transformation in words can degrade performance. This occurs compared to simple answer-only supervision.

How It Works

The authors introduce Imaginative Perception Tokens (IPT) to externalize spatial simulation. They implement this using BAGEL, a unified decoder-only transformer. This architecture treats image tokens as first-class sequence elements. This allows the model to interleave understanding and generation.

The mechanism works in two stages:

1. Perceptual Imagination: Given an observed context $C$ (images and a query), the model generates $\hat{I}{imag}$. These tokens represent the implied spatial structure. For example, in perspective-taking, this is a rendering of a new viewpoint. This is optimized using a Flow-Matching loss ($L$). This loss trains the model to transform Gaussian noise into the target latent representation (a compressed mathematical representation of an image).
2. Conditioned Reasoning: The model then produces the final answer $A$. This is conditioned on both the original observations and the new imaginative tokens.

The architecture uses a Mixture-of-Transformer-Experts (MoT) design. It employs separate transformer experts for multimodal understanding and generation. These experts share self-attention (a mechanism allowing tokens to interact). This ensures that "imagined" tokens are functionally tied to the underlying 3D geometry. As seen in, even when generated "thoughts" are visually imperfect, the training signal shapes the model's internal representations.

Numbers

The authors report significant gains by shifting from text-based reasoning to visual imagination. On the Multiview Counting (MVC) task, IPT improves accuracy by 3.4% over label-only supervision. This achieves 67.3% accuracy on the AI2-THOR benchmark.

The comparison to textual CoT is notable. In the MVC setting, IPT (67.3%) outperforms Text CoT (62.3%). On the Path Tracing (PT) task, "Mixed Training" (combining IPT with label-only data) yields the best results. It hits 66.7% on synthetic data and 58.6% on real-world benchmarks. This outperforms both pure label-only training and textual CoT.

These benefits persist even in "answer-only" inference mode. This means the model does not actually generate the image during inference. Even so, IPT supervision during training creates stronger internal spatial representations. These representations transfer to new environments. For example, on the Habitat benchmark for Perspective Taking, the IPT-trained model reaches 87.0% accuracy. This outperforms the label-only variant.

What's Missing

A practitioner should note several gaps in the reporting:

• Inference Latency and Cost: The paper discusses "answer-only" mode. However, it does not fully quantify the overhead of "imagination mode." This is the mode where the model performs iterative denoising to generate an image. High-resolution generation will impact throughput and tail latency.
• Hallucination Risks: The authors admit in and that generated thoughts can be spatially imprecise.

They may contain artifacts. If a system relies on these images for decision-making, this remains a reliability bottleneck. * Scaling Laws for IPT: The paper focuses on fine-tuning a 7B parameter model. It is not shown if the efficacy of IPT scales with model size. It is unclear if explicit visual supervision becomes redundant as models grow.

Should You Prototype This

Depends on your task.

Skip this if you are building a general-purpose VQA bot. Managing a unified generative/understanding architecture like BAGEL is complex. However, consider this if your roadmap includes embodied AI or robotics. These fields require "mental mapping." The takeaway that visual intermediates outperform textual CoT for geometry is a solid principle.

Code is reportedly available at the project page; see the paper for the canonical link. If you prototype, start with "answer-only" inference. You can reap the benefits of improved internal representations without the compute tax of real-time image generation.

Figures from the paper