visual-reasoning-revisitation

name	visual-reasoning-revisitation
title	Don't Look Only Once: Towards Multimodal Interactive Reasoning with Selective Visual Revisitation
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2505.18842
keywords	["Multimodal Reasoning","Visual Grounding","Active Referencing","Vision-Language"]
description	Enable multimodal models to dynamically revisit and re-ground reasoning steps in images using point-and-copy mechanisms for better long-horizon reasoning.

Re-ground Visual Reasoning Through Selective Image Revisitation

Multimodal language models encode images once into key-value caches, then reason purely in text. This approach works for simple questions but fails for complex reasoning: as reasoning chains lengthen, models progressively lose focus on relevant visual regions. Humans don't work this way—they revisit visual evidence repeatedly while thinking.

The solution is dynamic visual referencing: enable models to selectively revisit the image during reasoning, pointing to relevant regions and updating their understanding. This "active visual referencing" grounds intermediate reasoning steps back in the image, preventing drift and improving accuracy on multi-step visual reasoning tasks.

Core Concept

The key insight is that reasoning about images should be interactive, not one-shot. Long reasoning chains require multiple passes over the image:

• Dynamic referencing: Model selects when to revisit the image during reasoning
• Point-and-copy mechanism: Select regions (via coordinates/points) to attend to
• Re-grounding: Update visual context at specific reasoning steps
• Attention preservation: Maintain focus on relevant regions as reasoning progresses
• Selective revisitation: Only revisit when reasoning confidence drops or new information is needed

This prevents catastrophic forgetting of visual details during text-only reasoning phases.

Architecture Overview

• Visual encoder: Standard vision backbone (CLIP, DINO, etc.)
• Spatial attention mechanism: Ability to focus on regions identified by coordinates
• Dynamic referencing controller: Decides when to revisit image and which regions
• Point selector module: Model outputs coordinates to re-attend to
• Hybrid reasoning: Alternates between visual and text-only reasoning steps
• KV cache management: Maintains fresh visual context for referenced regions

Implementation

Implement selective visual revisitation by adding a pointing mechanism to multimodal models:

# Dynamic visual referencing for multimodal reasoning
import torch
import torch.nn as nn
from einops import rearrange

class SelectiveVisualRevisitation(nn.Module):
    """
    Enable models to dynamically revisit and re-ground visual information.
    """
    def __init__(self, hidden_dim=768, image_size=336, num_visual_tokens=256):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.image_size = image_size
        self.num_visual_tokens = num_visual_tokens

        # Spatial attention: focus on specific regions
        self.spatial_attention = nn.MultiheadAttention(
            hidden_dim, num_heads=8, batch_first=True
        )

        # Point predictor: model outputs (x, y) to attend to
        self.point_predictor = nn.Sequential(
            nn.Linear(hidden_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 2)  # Output (x, y) coordinates
        )

        # Region encoder: extract context around pointed region
        self.region_encoder = nn.Linear(hidden_dim, hidden_dim)

        # Revisitation decision module
        self.revisit_decision = nn.Sequential(
            nn.Linear(hidden_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward():
        
        batch_size, seq_len, hidden_dim = text_state.shape

        
        revisit_prob = .revisit_decision(text_state[:, -])  

        
        predicted_points = .point_predictor(text_state[:, -])  
        predicted_points = torch.sigmoid(predicted_points)  
        predicted_points = predicted_points * .image_size

        
        
        distances = torch.cdist(
            predicted_points.unsqueeze(),  
            image_spatial_coords  
        )  

        
        spatial_attention_weights = torch.exp(-distances / (.image_size / ))
        spatial_attention_weights = spatial_attention_weights / (spatial_attention_weights.(dim=-, keepdim=) + )

        
        attended_visual = torch.einsum(
            ,
            spatial_attention_weights,
            visual_features.unsqueeze()
        )  

        
        reground_vector = .region_encoder(attended_visual.squeeze())

        
        updated_state = text_state.clone()
        
        updated_state[:, -] = updated_state[:, -] + (reground_vector * revisit_prob)

         updated_state, {
            : revisit_prob.mean().item(),
            : predicted_points.cpu().numpy(),
            : spatial_attention_weights.squeeze().cpu().detach()
        }

Implement a wrapper that enables dynamic revisitation during multimodal reasoning:

class MultimodalReasonerWithRevisitation:
    """
    Multimodal model that can revisit image during reasoning chain.
    """
    def __init__(self, base_model, revisitation_module):
        self.base_model = base_model
        self.revisitation = revisitation_module

    def reason_with_revisitation(self, image, question, max_reasoning_steps=10,
                                revisit_threshold=0.3):
        """
        Generate reasoning with dynamic visual revisitation.
        """
        # Encode image once
        image_features = self.base_model.encode_image(image)
        image_coords = self._get_patch_coordinates(image)

        # Initial text encoding
        text_encoding = self.base_model.encode_text(question)

        reasoning_trace = []
        current_state = text_encoding

        for step in range(max_reasoning_steps):
            # Generate next reasoning token
            next_token, logits = self.base_model.generate_token(current_state)

            reasoning_trace.append({
                'step': step,
                'token': next_token,
                'revisited': False
            })

            # Check if we should revisit image
            revisit_info = self.revisitation(current_state, image_features, image_coords)
            if revisit_info['revisit_prob'] > revisit_threshold:
                # Re-ground reasoning in image
                current_state = revisit_info['updated_state']
                reasoning_trace[-][] = 
                reasoning_trace[-][] = revisit_info[]

            
            current_state = .base_model.update_state(current_state, next_token)

            
             next_token == .base_model.eos_token_id:
                

         {
            : reasoning_trace,
            : (  t  reasoning_trace  t[]),
            : (reasoning_trace)
        }

     ():
        
        batch_size = image.shape[]
        num_patches_h = image.shape[] //   
        num_patches_w = image.shape[] // 

        
        h_coords = torch.linspace(, image.shape[], num_patches_h)
        w_coords = torch.linspace(, image.shape[], num_patches_w)
        coords = torch.stack(torch.meshgrid(h_coords, w_coords, indexing=), dim=-)

         coords.unsqueeze().expand(batch_size, -, -, -)

Implement a training procedure that teaches models to use revisitation effectively:

def train_with_revisitation_supervision(model, train_data, num_epochs=10):
    """
    Train multimodal model to learn when and where to revisit images.
    Uses supervision signal: which regions are relevant for each reasoning step.
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        for batch in train_data:
            images = batch['images']
            questions = batch['questions']
            reasoning_steps = batch['reasoning_steps']  # Multi-step reasoning
            visual_grounding = batch['visual_grounding']  # Which regions matter per step

            # Forward pass with revisitation
            image_features = model.encode_image(images)
            image_coords = model._get_patch_coordinates(images)

            losses = []

            for step_idx, step_tokens in enumerate(reasoning_steps):
                # Run one step of reasoning
                state = model.encode_text(questions)

                # Get revisitation decision
                revisit_logits = model.revisitation.revisit_decision(state)

                # Get spatial attention
                spatial_attention = model.revisitation(state, image_features, image_coords)

                # Supervision: which regions should be attended?
                target_regions = visual_grounding[step_idx]

                # Loss 1: spatial attention alignment
                target_attention = create_attention_mask(target_regions, image_features.shape)
                spatial_loss = F.kl_div(
                    F.log_softmax(spatial_attention, dim=-1),
                    target_attention,
                    reduction='batchmean'
                )

                # Loss 2: revisitation decision
                should_revisit = ((target_regions) > ).()
                revisit_loss = F.binary_cross_entropy(revisit_logits, should_revisit.unsqueeze())

                total_loss = spatial_loss +  * revisit_loss
                losses.append(total_loss)

            total_loss = torch.stack(losses).mean()
            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

             epoch %  == :
                ()

Practical Guidance

Parameter	Typical Range	Notes
Revisit threshold	0.3 - 0.7	Probability needed to trigger revisitation
Point prediction scale	image_size	Normalize coordinates appropriately
Spatial attention sigma	image_size/8 to image_size/4	Controls region size around point
Max revisits per sequence	3 - 8	Balance accuracy with compute
Patch size	14 - 32 pixels	Smaller = higher resolution, higher compute

When to use selective visual revisitation:

• Multi-step visual reasoning tasks (VQA, visual understanding)
• Long reasoning chains over images
• Problems requiring attention to multiple image regions
• Accuracy is more important than single-pass latency
• Models struggle with mid-chain reasoning drift

When NOT to use:

• Single-question image tasks (captioning, classification)
• Simple visual understanding (no complex reasoning needed)
• Latency is critical (adds compute per revisit)
• Image complexity is low (doesn't benefit from re-grounding)
• Models already encode sufficient visual context in cache

Common pitfalls:

• Revisit threshold too low (excessive re-encoding, no benefit)
• Point predictor not well-calibrated (predicting outside image)
• Not supervising which regions to attend during training
• Not measuring whether revisitation actually helps on the task
• Using too-coarse spatial patches (losing region specificity)
• Not balancing revisitation frequency (must be selective)

Reference

Don't Look Only Once: Towards Multimodal Interactive Reasoning with Selective Visual Revisitation https://arxiv.org/abs/2505.18842