Latent Chain-of-Thought for Visual Reasoning

AI summary

Overview

Research area: Multimodal reasoning with Large Vision-Language Models (LVLMs), specifically how to train them to produce chain-of-thought (CoT) reasoning that generalizes to unseen tasks, and how to select the best reasoning trace at inference time.

Technical level: Advanced. The paper builds on Generative Flow Networks (GFlowNets), amortized variational inference, and Bayesian marginal-likelihood estimation; familiarity with reinforcement learning for language models (PPO, GRPO) is assumed.

Scope (one sentence): The paper reformulates visual CoT as posterior inference over latent reasoning traces and introduces a training algorithm (RGFN), a token-level reward approximation, and a Bayesian inference-scaling method (BiN) that together form a system called LaCoT.

What This Paper Is About

Existing approaches to giving LVLMs step-by-step reasoning ability—Supervised Fine-Tuning (SFT), PPO, and GRPO—optimize next-token distributions or scalar rewards and often fail to generalize to unseen reasoning tasks; PPO and GRPO also depend heavily on a reward model that can be biased, and can be exploited through reward hacking. The paper treats the reasoning chain as a latent variable Z drawn from a posterior P(Z|X,Y) given a question X and answer Y, and asks how to sample that posterior efficiently and use it to pick the best answer.

Key Contributions

1. RGFN (Reference-Guided GFlowNet fine-tuning): A training objective based on Sub-Trajectory Balance (SubTB) that adds a reference-guided exploration mechanism. Candidate rationales the model generates are compared against a reference rationale Z_ref and filtered by an indicator function before gradients are applied, which the authors say prevents catastrophic forgetting and removes the diversity restriction imposed by a KL penalty.

2. Token-level marginal reward approximation: Because Eq. (2) requires per-token rewards, which is infeasible for reasoning chains of thousands of tokens, the authors compute the true reward every λ steps and linearly interpolate the rewards of intermediate steps. They state a proposition (Proposition 1) showing the interpolation error is bounded by Mλ²/8, i.e., decays as O(λ²).

3. BiN (Bayesian inference over N latent rationales): An inference-scaling strategy that replaces Best-of-N (BoN) and Beam Search. It samples N latent rationales, samples answers, and ranks them by a length-normalized marginal likelihood, eliminating the need for a separate critic or reward model at inference time.

4. The LaCoT system: An instantiation of the above on Qwen2.5-VL 3B and 7B, evaluated on seven reasoning benchmarks (MathVista, MathVision, MathVerse, MMMU, MMMU-pro, MMVet, MME), with code released at https://github.com/heliossun/LaCoT.

Main Findings

• Headline improvements over base models: The paper reports that the 7B model improves 6.6% over its base model and outperforms GRPO by 10.6%; the 3B model surpasses its base model by 13.9% and outperforms larger models such as LLaVA-CoT-11B and LLaVA-OV-7B.

• Benchmark table (Table 1, test accuracy % and MME score): LaCoT-Qwen-7B scores 68.4 (MathVista), 24.9 (MathVision), 43.3 (MathVerse), 54.9 (MMMU), 35.3 (MMMU-pro), 74.2 (MMVet), and 2372 (MME), versus the Qwen2.5-VL-7B base model at 63.7, 25.4, 38.2, 50.0, 34.6, 70.5, and 2333. LaCoT-Qwen-3B scores 63.2, 20.7, 40.0, 48.8, 28.9, 69.6, and 2208, versus the Qwen2.5-VL-3B base at 60.3, 21.2, 26.1, 46.6, 22.4, 61.4, and 2134. The paper states LaCoT narrows the gap to GPT-4o to less than 3 points while using only 7 billion parameters.

• MathVerse gains: MathVerse-Vision-only improves the most, with the 3B model's accuracy jumping 14 points and outperforming all 7B models listed—which the authors attribute to better diagram comprehension and OCR robustness.

• MathVision remains hard: On MathVision, which the paper describes as real Olympiad diagrams that are more varied and often handwritten or low-resolution, a single misread propagates through long proof-style reasoning chains, leading to a performance drop.

• BiN beats BoN (Table 2): With the same policy model and no external reward model, BiN outperforms BoN at both sizes. For 3B: MathVerse 40.0 vs 21.2, MathVista 63.2 vs 57.1, MMMU 48.8 vs 44.7, MMVet 69.6 vs 67.1. For 7B: 39.7 vs 26.5, 68.4 vs 62.2, 54.9 vs 47.3, 74.2 vs 71.2. Both methods were evaluated at N ∈ {5, 10}, reporting the best score per method.

• Ablation of training algorithms (Table 3, Qwen2.5-VL-7B): RGFN reaches 68.4 / 43.3 / 54.9 on MathVista / MathVerse / MMMU, versus zero-shot 63.7 / 38.2 / 50.0, SFT 62.7 / 38.7 / 50.6, and GRPO 62.6 / 36.8 / 47.9. The authors attribute GRPO's poor result to inadequate reward-model guidance and limited exploration from the KL penalty.

• BiN generalizes to third-party models (Table 4): Applying BiN with N=5 and T=0.7 to Qwen2.5-VL models trained with ordinary SFT improves them on all benchmarks—7B (SFT): 62.7 / 38.7 / 50.6 to 64.4 / 38.9 / 51.6; 3B (SFT): 58.7 / 33.3 / 43.1 to 59.4 / 35.2 / 45.0.

• Diversity and likelihood: Sampling 5 rationale candidates with random temperature per instruction, the authors find LaCoT-Qwen-3B at T = 0.7 produces rationales with the highest log-likelihood and the highest semantic diversity, measured as average inter-sentence similarity against a reference set (Fig. 6).

• Scaling behavior (Fig. 7): Accuracy on LaCoT-Qwen-3B consistently increases with both the number of candidates N and temperature T. Increasing N from 1 to 5 significantly mitigates hallucination and improves accuracy on MMMU; the authors attribute this to reduced Monte-Carlo variance (standard error scaling as O(1/√N)), broader posterior coverage, smoothing of length-normalization fluctuations, and a larger candidate answer set.

Methodology in Plain English

The authors start from the observation that a reasoning chain is something we want the model to discover, not something it should simply memorize. So they treat the chain as a hidden variable and train a model to sample plausible chains given a question and answer.

To do this they use GFlowNets, a family of generative models that learn to sample objects with probability proportional to a reward rather than just maximizing reward. This is the key difference from PPO and GRPO: instead of hunting for the single highest-scoring trace, the model is trained to match a whole probability distribution over traces, which naturally keeps many different valid reasoning paths alive.

Three practical problems arise, and the paper solves each:

1. Rewards at every token are too expensive. Instead, they compute the true likelihood-based reward every λ steps (they use λ = 8) and linearly interpolate between those checkpoints. They prove the interpolation error is small when the segment is short.

2. Unconstrained exploration makes the model forget how to write sensible text. They generate m candidate rationales, score them, and discard any that do not beat a threshold based on a reference rationale Z_ref (in practice a CoT generated by a teacher model such as GPT-4o or Deepseek-R1). The threshold anneals: δ_s = τ_max − (τ_max − τ_min) × min(1, s/50), so early training allows more exploration and the bar tightens after 50 steps.

3. Picking the best answer at test time is costly. Instead of Best-of-N with a reward model, they sample N rationales and answers, compute each joint likelihood, normalize by sequence length, and return the answer with the highest estimated marginal likelihood. This removes the critic model entirely.

Implementation details: the reward/answer model π_Φ is a Qwen2.5-VL 3B or 7B fully fine-tuned for one epoch with ordinary token-prediction loss on a mixture of visual reasoning data from LLaVA-CoT and R1-Onevision, using a new special role token Analyzer so the model can selectively produce reasoning steps. The rationale sampler q_θ(Z|X) is initialized from π_Φ and trained with LoRA (r = 64, alpha = 128) on 3k resampled visual reasoning examples containing image, query, CoT, and answer.

Why This Matters

The paper targets a real weakness in current multimodal reasoning systems: models trained with reward maximization tend to produce one narrow style of reasoning, can be gamed by a biased reward model, and degrade on tasks that look different from their training data. Framing CoT as probabilistic inference offers a principled alternative, and the reported gains on seven benchmarks at the 3B and 7B scale suggest the approach is practical, not just theoretical.

Real-world applications:

• Educational and tutoring systems that need to show students a trustworthy, step-by-step solution to a diagram-based math or science problem, where multiple valid solution paths exist.
• Document and diagram understanding, such as charts, engineering drawings, or geometric figures, where the reported MathVerse-Vision-only gains point to better visual grounding and OCR robustness.
• Scientific and technical question answering, where multi-discipline college-level reasoning is required and the MMMU/MMMU-pro settings are a proxy.
• Assistive and accessibility tools that describe and reason about images, where hallucination control (demonstrated by increasing N) matters for user trust.

Industry relevance: the method is designed as a drop-in training and inference layer for any autoregressive LVLM. BiN in particular is shown to improve even models fine-tuned with plain SFT, meaning an existing deployed model could gain accuracy at inference time without retraining—relevant for companies serving vision-language products where retraining budgets and inference costs are both constrained.

Future Directions

• Scaling beyond 7B parameters. The paper explicitly notes that resource constraints limited experiments to models up to 7B, and states the expectation that conclusions hold for larger models.
• Exploration in complex latent spaces. The authors flag exploration as an open challenge, noting that factors such as sequence length and memory cost affect exploration time.
• Hallucination and internal knowledge. The paper states that despite improved inference performance, it does not address hallucination, which is closely related to internal knowledge.
• Knowledge distillation and synthetic data generation. The conclusion names these as the next applications of the system to investigate.

Target Audience

Researchers and engineers working on multimodal reasoning, reinforcement learning for language models, or probabilistic inference for generative models—particularly those who already understand PPO/GRPO-style training and want to see a variational-inference alternative. Practitioners deploying vision-language models who care about inference-time accuracy gains will find the BiN results (Tables 2 and 4) independently useful. Readers without a background in GFlowNets or variational inference will find the preliminaries section dense, though the high-level framing and the empirical tables are accessible on their own.

Authors’ abstract

Chain-of-thought (CoT) reasoning is critical for improving the interpretability and reliability of Large Vision-Language Models (LVLMs). However, existing training algorithms such as SFT, PPO, and GRPO may not generalize well across unseen reasoning tasks and heavily rely on a biased reward model. To address this challenge, we reformulate reasoning in LVLMs as posterior inference and propose a scalable training algorithm based on amortized variational inference. By leveraging diversity-seeking reinforcement learning algorithms, we introduce a novel sparse reward function for token-level learning signals that encourage diverse, high-likelihood latent CoT, overcoming deterministic sampling limitations and avoiding reward hacking. Additionally, we implement a Bayesian inference-scaling strategy that replaces costly Best-of-N and Beam Search with a marginal likelihood to efficiently rank optimal rationales and answers. We empirically demonstrate that the proposed method enhances the state-of-the-art LVLMs on seven reasoning benchmarks, in terms of effectiveness, generalization, and interpretability.