Mitigating Visual Forgetting via Take-along Visual Conditioning for Multi-modal Long CoT Reasoning

TL;DR

The paper identifies a visual-forgetting problem in multimodal long CoT reasoning, where models progressively neglect visual inputs as reasoning unfolds. It presents Take-along Visual Conditioning (TVC), a two-stage approach with Dynamic Visual Reaffirmation during training and Periodic Visual Calibration at inference, plus visual token compression to maintain grounding. Through a data-centric long-chain reasoning pipeline and iterative distillation with reject sampling, TVC achieves substantial gains across multiple mathematical reasoning benchmarks and demonstrates robustness through ablations and case studies. The work significantly improves multimodal reasoning fidelity and offers a practical path toward more reliable visual grounding in complex tasks.

Abstract

Recent advancements in Large Language Models (LLMs) have demonstrated enhanced reasoning capabilities, evolving from Chain-of-Thought (CoT) prompting to advanced, product-oriented solutions like OpenAI o1. During our re-implementation of this model, we noticed that in multimodal tasks requiring visual input (e.g., geometry problems), Multimodal LLMs (MLLMs) struggle to maintain focus on the visual information, in other words, MLLMs suffer from a gradual decline in attention to visual information as reasoning progresses, causing text-over-relied outputs. To investigate this, we ablate image inputs during long-chain reasoning. Concretely, we truncate the reasoning process midway, then re-complete the reasoning process with the input image removed. We observe only a ~2% accuracy drop on MathVista's test-hard subset, revealing the model's textual outputs dominate the following reasoning process. Motivated by this, we propose Take-along Visual Conditioning (TVC), a strategy that shifts image input to critical reasoning stages and compresses redundant visual tokens via dynamic pruning. This methodology helps the model retain attention to the visual components throughout the reasoning. Our approach achieves state-of-the-art performance on average across five mathematical reasoning benchmarks (+3.4 points vs previous sota), demonstrating the effectiveness of TVC in enhancing multimodal reasoning systems.

Figures (8)

• Figure 1: The visual forgetting phenomenon by removing the image at different reasoning stages. It shows that by the midpoint of the reasoning process, the model becomes less dependent on the image, causing text-over-relied outputs.
• Figure 2: Illustration of layer-level and token-level attention weights. (a) The layer-level attention weights of image tokens across different response token positions. (b) The token-level attention weights at the middle layer. It shows that the model's attention to the image gradually decreases during the reasoning process.
• Figure 3: Overview of TVC System Design. We enable the model to have take-along visual conditioning capabilities through two stages: training and inference.
• Figure 4: Data Generation Pipeline of TVC. We use iterative distillation to collect long-chain reasoning data, followed by a comprehensive response filtering process to ensure high-quality reasoning.
• Figure 5: Ablations on the amount of training data. TVC benefits from data scaling, continually improving the reasoning capabilities.