Technical Summary: Imaginative Perception Tokens Enhance Spatial Reasoning in Multimodal Language Models

1. Problem Statement

Vision-Language Models (VLMs) currently excel at recognizing objects and attributes but struggle significantly with spatial reasoning, particularly when the required information is not directly observable in the input. Many spatial tasks require imaginative perception: simulating unseen viewpoints, tracing trajectories through occluded spaces, or integrating partial views into a coherent spatial map.

Existing approaches attempt to bridge this gap using intermediate visual representations (e.g., visual thoughts, depth maps, or bounding box tokens). However, these methods primarily refine structures already visible in the input or extract perceptual attributes. They fail to address the core challenge of constructing spatial representations that are not directly observed, which often requires predicting missing spatial structure implied by evidence rather than simply describing what is seen. Furthermore, forcing spatial computation through textual chain-of-thought (CoT) often results in a modality mismatch, as serializing geometric constraints and viewpoint transformations into language is error-prone and inefficient.

2. Methodology

2.1. Core Concept: Imaginative Perception Tokens (IPT)

The authors propose Imaginative Perception Tokens (IPT), a form of intermediate perceptual representation that externalizes what a VLM would perceive under an alternative spatial configuration while remaining consistent with the observed input. Unlike standard intermediates that describe visible structures, IPTs represent novel spatial views (e.g., a view from a new coordinate or a top-down map integrating multiple views) that must be inferred.

2.2. Tasks and Datasets

To study this capability, the authors formulate three tasks requiring imaginative perception, constructing datasets of approximately 20,000 examples each, spanning simulated (AI2-THOR, Habitat) and real-world settings:

1. Perspective Taking (PET): Given a first-person view and a target position, predict the scene from the new viewpoint (e.g., "After moving to 'X' and turning left, is the object on your left or right?"). The ground truth is a novel-viewpoint image.
2. Path Tracing (PT): Given a top-down map and forward views at waypoints, identify objects visible at a midpoint along a path. The ground truth is a sideview image (first-person rendering from the midpoint).
3. Multiview Counting (MVC): Given multiple first-person frames, determine the total count of a specific object by integrating views to resolve occlusions and duplicates. The ground truth is a top-down Bird's-Eye View (BEV) map.

Each dataset includes paired ground-truth intermediate imaginations and final answers, accompanied by human-filtered evaluation benchmarks.

2.3. Model Architecture and Training

The authors utilize BAGEL, a unified decoder-only transformer that natively supports interleaved multimodal understanding and generation.

• Architecture: BAGEL employs a Mixture-of-Transformer-Experts (MoT) design with separate experts for understanding (SigLIP2 ViT tokens) and generation (FLUX VAE latent tokens).
• Training Objective: The model is trained with a multi-task loss ($L_{total}$) comprising:
  1. Flow-Matching Loss ($L_{fm}$): The model learns to predict the velocity field to transform Gaussian noise into the target latent representation of the unobserved view (the IPT).
  2. Language Modeling Loss ($L_{lm}$): The model minimizes the negative log-likelihood of the final answer tokens, conditioned on the observed context and the ground-truth imaginative tokens.
• Inference Modes:
  • Imagination Mode: The model generates the intermediate image, re-encodes it, and uses it as context to predict the answer.
  • Answer-Only Mode: The model predicts the answer directly without generating an image at inference time, relying on the internal representations learned during training.

3. Key Results

3.1. Performance on Spatial Tasks

Using the BAGEL backbone, IPT supervision significantly improves spatial reasoning across all three tasks compared to baselines:

• Multiview Counting (MVC): IPT improves accuracy by 3.4% over label-only training (67.3% vs. 63.9%) and outperforms textual CoT (62.3%).
• Perspective Taking (PET): IPT achieves 96.8% accuracy on in-domain AI2-THOR data, outperforming label-only (97.5% is the max, but IPT is competitive) and significantly beating textual CoT (83.1%). On out-of-domain (Habitat) data, IPT reaches 87.0%, surpassing label-only (82.0%).
• Path Tracing (PT): Mixed training (IPT + label-only) achieves the best results, reaching 66.7% on synthetic data and 58.6% on real-world benchmarks, outperforming both label-only and textual CoT.

3.2. Comparison with Textual Chain-of-Thought

The paper finds that textual CoT is often detrimental on these tasks. For example, on PET, CoT drops performance to 83.1% compared to 96.8% for IPT. The authors attribute this to a modality mismatch: forcing spatial computation (viewpoint changes, occlusions) through language introduces noise and competes with answer prediction capacity. In contrast, IPTs represent these relationships directly in the visual modality.

3.3. Inference and Generalization

• Internal Representations: Crucially, the performance gains persist even when the model does not generate an image at inference time (Answer-Only mode). This suggests that training with IPT supervision helps models develop stronger internal spatial representations.
• Transfer Learning: Models trained with IPTs show improved transfer to external benchmarks (e.g., SAT, MessyTable) and other spatial tasks (ScanNet, MindCube), indicating that the learned spatial reasoning generalizes beyond the specific token targets.
• Resolution Ablation: Higher resolution latent tokens (Latent-64, 1024x1024) produce sharper, more spatially faithful imaginations, correlating with higher downstream accuracy.

4. Significance and Claims

The paper claims that Imaginative Perception Tokens provide a principled supervision signal for reasoning over unobserved structure. The key contributions and claims are:

1. Novel Intermediate Representation: IPTs are distinct from prior visual intermediates because they predict unobserved spatial configurations rather than refining visible ones.
2. Modality Alignment: Supervising models with intermediate visual predictions is more effective than textual CoT for spatial tasks, avoiding the bottlenecks of serializing geometric constraints into language.
3. Generalization: Training with imaginative supervision yields stronger spatial generalization and more interpretable intermediates aligned with the task's underlying geometry, even without explicit image generation during inference.
4. Benchmarking: The work introduces the first datasets explicitly designed to train and evaluate visually-grounded intermediate spatial reasoning, highlighting a "perspective gap" and the need for constructive spatial imagination in VLMs.

The authors conclude that while current unified models often trade understanding robustness for generation capability, integrating imaginative perception supervision allows models to achieve performance competitive with strong closed-source models on complex spatial reasoning tasks.