Skip to content

Latent Chain-of-Thought for Visual Reasoning

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LaCoT reformulates chain-of-thought reasoning in Large Vision-Language Models (LVLMs) as posterior inference over latent rationale variables, training a policy via amortized variational inference (GFlowNets) rather than SFT/PPO/GRPO. A token-level reward approximation and reference-guided exploration enable diverse, high-quality rationale sampling. At inference, a Bayesian marginal-likelihood ranking (BiN) replaces reward-model-dependent Best-of-N search.

Problem

SFT encourages parroting reference traces; PPO and GRPO constrain exploration via KL penalties, limiting rationale diversity and causing reward hacking. None explicitly model the full posterior over reasoning chains, leading to poor generalization on unseen visual reasoning tasks.

Method

Visual CoT is cast as posterior inference: given observation (X, Y), sample latent rationale Z ~ P(Z|X,Y). The policy model q_θ(Z|X) is trained with RGFN (Reference-Guided GFlowNet Fine-tuning), which combines: 1. Interpolated Sub-Trajectory Balance (ISubTB): sparse token-level rewards computed every λ=8 tokens with linear interpolation between checkpoints, replacing the prohibitively expensive exact per-token reward. 2. Reference-guided exploration: m candidate rationales are sampled; those whose joint log-likelihood falls below an annealed threshold δ_s relative to a teacher-generated reference Z_ref are filtered before backprop, preventing catastrophic forgetting without KL penalties.

At inference, BiN samples N rationale–answer pairs {Z_i, Y_i}, computes length-normalized joint log-likelihoods π_Φ(Z_i Y_i | X) / |Z_i Y_i|, and selects the answer with the highest estimated marginal likelihood—no external critic needed. The full system (LaCoT) pairs a LoRA-finetuned rationale sampler q_θ with an SFT-initialized answering model π_Φ (Qwen2.5-VL 3B/7B backbone).

Key Contributions

  • RGFN training objective: GFlowNet-based amortized VI with reference-guided candidate filtering, replacing KL-penalized RL for LVLM CoT.
  • ISubTB reward approximation: piecewise linear interpolation for token-level rewards reduces compute without significant accuracy loss (Proposition 1).
  • BiN inference scaling: critic-free, Bayesian marginal-likelihood ranking over N sampled rationale–answer pairs.
  • LaCoT system: end-to-end integration of a latent rationale sampler and an answering model, demonstrating improved generalization over SFT and GRPO baselines.

Results

  • LaCoT-Qwen-7B vs. Qwen2.5-VL-7B base: +6.6% average across 7 benchmarks; vs. R1-Onevision (GRPO baseline): +10.6% on MathVista (68.4 vs. 64.1).
  • LaCoT-Qwen-3B vs. Qwen2.5-VL-3B base: +13.9% improvement; outperforms LLaVA-CoT-11B and LLaVA-OV-7B on most benchmarks despite 4× smaller size.
  • MathVerse Vision-Only (3B): 40.0% vs. 26.1% base (+14 points); surpasses all 7B models on this split.
  • MMMU val (7B): 54.9% vs. 47.9% base; 3B reaches 48.8%, matching LLaVA-CoT-11B (48.8%).
  • BiN vs. BoN (Table 2, 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.
  • Ablation (7B): Zero-shot 63.7/38.2/50.0 → SFT 62.7/38.7/50.6 → GRPO 62.6/36.8/47.9 → RGFN 68.4/43.3/54.9 (MathVista/MathVerse/MMMU).
  • BiN scales with N: MathVerse accuracy rises from 32.7 (N=1) to 40.0 (N=10) at T=0.7 for 3B model.

Limitations

  • Experiments restricted to models ≤7B parameters due to resource constraints; scalability to larger models is assumed but untested.
  • On-policy exploration remains challenging for long reasoning chains (∼1k tokens) due to high memory cost and exploration time.
  • Does not directly address hallucination; performance on MathVision (complex Olympiad diagrams) degrades relative to other benchmarks, highlighting sensitivity to visual grounding quality.
  • Requires SFT-pretrained Ï€_Φ and teacher-generated reference rationales (GPT-4o / DeepSeek-R1), introducing dependence on proprietary data pipelines.

Relevance to Vision-Language Models

LaCoT directly addresses a core bottleneck in LVLM reasoning: existing RLHF-style fine-tuning collapses rationale diversity and risks reward hacking, whereas this work applies principled amortized VI (GFlowNets) to the multimodal setting. The BiN inference strategy is model-agnostic and immediately applicable to any reasoning LVLM, making it a practical upgrade over Best-of-N without an external reward model. For researchers tracking VLMs, this represents a methodologically distinct alternative to the dominant GRPO/PPO paradigm, with strong empirical gains on standard multimodal math benchmarks (MathVista, MathVerse, MMMU) while improving interpretability through probabilistic rationale ranking.