Source-linked AI summary
Latent Chain-of-Thought for Visual Reasoning
Guohao Sun, Hang Hua, Jian Wang, Jiebo Luo, Sohail Dianat, Majid Rabbani, Raghuveer Rao, Zhiqiang Tao
TL;DR
Existing LVLM training methods can struggle to generalize across reasoning questions, motivating latent-rationale posterior inference. LaCoT uses amortized variational training and Bayesian inference scaling; its 7B model improves over its base model and GRPO, while BiN improves SFT performance across benchmarks.
Problem
SFT, PPO, and GRPO may struggle to generalize across diverse reasoning questions because they do not explicitly capture dependencies across trajectories.
Source: §1 · p. 1Method
LaCoT treats rationales as latent variables, trains an amortized variational sampler, and uses BiN to rank answers by length-normalized marginal likelihood.
Source: §1 · p. 2, §3 · p. 3, §3.3 · p. 6Results
BiN consistently improves Qwen2.5-VL 3B and 7B SFT performance on all benchmarks; LaCoT 7B improves 6.6% over its base model and outperforms GRPO by 10.6%.
Source: §1 · p. 2, §4.4 · p. 9Takeaways & Limitations
BiN provides a scalable, probabilistically justified strategy for ranking rationales and answers without a critic model.
Source: §1 · p. 2Experiments apply the methods to models up to 7B parameters, and hallucinations related to internal knowledge remain unaddressed.
Source: §6 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
The introduction presents visual chain-of-thought as important for LVLM interpretability and reliability, while identifying generalization, exploration, and reward-hacking weaknesses in existing training. It proposes latent-rationale inference with LaCoT and Bayesian inference scaling.
1 Introduction
- SFT, PPO, and GRPO struggle to generalize across diverse reasoning questions, with constrained exploration and possible reward hacking among the stated limitations.SFT relies on teacher-forced traces, while PPO and GRPO use KL penalties that keep exploration close to the SFT baseline.Source: §1 · p. 1
- LaCoT treats rationales as latent variables sampled from a question-answer posterior and uses amortized variational inference to encourage diverse, high-likelihood trajectories.The posterior is intractable to sample directly; AVI optimizes an ELBO for token-level learning.Source: §1 · p. 2
- RGFN combines reference-guided policy exploration with token-level reward approximation to support diverse sampling and efficient mini-batch exploration.The reference-guided method is intended to address catastrophic forgetting and the diversity constraint from KL penalties.Source: §1 · p. 2
- BiN ranks rationales and answers using length-normalized marginal likelihood, replacing costly Best-of-N or Beam Search and avoiding a critic model.The introduction describes this as a Bayesian, probabilistically justified inference-time search strategy.Source: §1 · p. 2
- LaCoT’s 7B model improves 6.6% over its base model and outperforms GRPO by 10.6%; its 3B model improves 13.9% over its base and surpasses larger models.The larger-model comparisons named are LLaVA-CoT-11B and LLaVA-OV-7B.Source: §1 · p. 2
2 Preliminaries
The preliminaries describe GFlowNets as amortized variational methods for sampling structured trajectories according to reward, then outline trajectory-flow matching and its visual-reasoning comparison with PPO.
2 Preliminaries
- GFlowNets are amortized variational methods for sampling structured objects in proportion to an unnormalized reward, trained by minimizing divergence between learned and target trajectory distributions.The passage names Sub-Trajectory Balance as a hierarchical variational objective.Source: §2 · p. 2
- For causal token generation, each state has one parent, so the backward policy is 1 and the forward policy is modeled by qθ.The loss enforces flow consistency across complete rationales and penalizes discrepancies between flows along truncated trajectories.Source: §2 · p. 2
- The trajectory-flow relation compares a state’s flow, scaled by intervening transition probabilities, with the later state’s flow; terminal flow is defined by the trajectory reward.The terminal state is usually an eos token in an LLM.Source: §2 · p. 3
3 Amortizing Variational Inference for Latent Visual CoT
The section formulates visual reasoning as inference over latent CoT rationales, then develops methods for approximating token rewards, guiding exploration, and ranking answers by marginal likelihood.
3 Amortizing Variational Inference for Latent Visual CoT
- Given question-answer observations, the model learns to sample latent rationales in proportion to their unnormalized posterior likelihood.The section also identifies token-level rewards for long reasoning chains and efficient, diverse exploration as challenges.Source: §3 · p. 3
3.1 Token-level Marginal Reward Approximation
- Linear interpolation between rewards computed every λ steps approximates intermediate token rewards, enabling mini-batch exploration for diverse sampling and broader posterior coverage.Exact rewards for all states are computationally expensive for long rationale sequences.Source: §3.1 · p. 4, Figure 2 · p. 4
3.2 Reference-Guided GFlowNet Fine-tuning
- Unconstrained exploration can cause catastrophic forgetting; reference-guided exploration filters low-reward trajectories and gradually tightens its acceptance threshold.The threshold allows more exploration early and becomes stricter over 50 steps.Source: §3.2 · pp. 4–5, Figure 3 · p. 5
3.3 Bayesian Inference over Latent Rationales
- BiN samples latent rationales, generates answers, estimates length-normalized marginal likelihoods, and selects the answer with the highest estimate.It uses the learned policy to approximate the marginal likelihood rather than sampling all rationales from the posterior.Source: §3.3 · pp. 5–6
4 Empirical results
The experiments report LaCoT’s benchmark performance, rationale diversity, and inference-time scaling, alongside ablations of RGFN and BiN. They also describe the implementation used to train the reward model and rationale sampler.
4.1 Implementation
- The reward model πΦ is fine-tuned for one epoch on structured rationale examples, while the rationale sampler is initialized from it and trained with LoRA.The sampler uses 3,000 resampled visual-reasoning examples and teacher-generated reference rationales.Source: §4.1 · p. 6
4.2 Multi-modal Reasoning
- MathVerse Vision-only accuracy jumps 14 points at 3B, outperforming all 7B models; LaCoT also narrows the GPT-4o gap to less than 3 points using 7B parameters.The authors report consistent improvements on MathVista and MMMU and identify LaCoT as the best open-source LVLM in Table 1.Source: §4.2 · p. 7, Table 1 · p. 7
- LaCoT-Qwen-3B at T = 0.7 samples rationales with the highest log-likelihood and diversity among the tested models.The comparison uses five rationale candidates per visual instruction and measures semantic diversity by average inter-sentence similarity to the reference set.Source: Figure 6 · p. 7, §4.2 · p. 8
4.3 Inference-time Scaling
- BiN consistently outperforms Best-of-N on visual reasoning benchmarks when both use LaCoT-Qwen as the shared policy model.The comparison evaluates N ∈ {5, 10} and reports each method’s best score without an external reward model.Source: §4.3 · p. 8, Table 2 · p. 8
4.4 Ablation Studies
- RGFN performs better on complex examples than the listed baselines, with reference-guided exploration covering diverse trajectories rather than collapsing to one reward mode.The authors attribute GRPO’s poor performance partly to inadequate external-reward guidance and limited exploration from the KL penalty.Source: §4.4 · p. 8, Table 3 · p. 8
- Test accuracy consistently rises with higher candidate count N and temperature T; increasing N also stabilizes answer rankings and broadens posterior coverage.The authors further describe length-normalization smoothing and a larger candidate-answer set as effects of increasing N.Source: §4.4 · p. 8, Figure 7 · p. 8
- On MMMU, increasing N from 1 to 5 significantly mitigates hallucinated or misleading reasoning and improves answer accuracy.At N = 1, BiN may produce incorrect reasoning steps that lead to lower accuracy.Source: Figure 7 · p. 8, §4.4 · p. 9
- With N = 5 and T = 0.7, BiN consistently improves Qwen2.5-VL 3B and 7B SFT performance across all benchmarks.The paper selects this setting for a strong performance boost with relatively shorter inference time.Source: §4.4 · p. 9
5 Related Work
Related work covers learning-based multimodal CoT, reinforcement-learning methods for reasoning, and inference-time search. The paper positions its approach as amortized variational inference for diverse outputs and probabilistic scaling.
5 Related Work
- Learning-based multimodal CoT trains LVLMs on rationale-augmented data; cited approaches add multimodal consistency, majority voting, or structured rationales [40] [56] [46] [25].The section presents these methods as embedding reasoning trajectories into models through supervised learning.Source: §5 · p. 9
- RL methods including PPO [34] and GRPO [12] advance reasoning, but prior approaches are described as depending on human-annotated data and producing policies with limited diversity.The paper proposes amortized variational inference for LVLM training as an alternative intended to generate diverse outputs.Source: §5 · p. 9
- Plan-based inference scaling uses search strategies such as Best-of-N and beam search, while the paper proposes aligning marginal likelihood with reward to avoid explicit reward computation.The related-work passage notes that dedicated reward-model forward passes make explicit trajectory evaluation computationally expensive.Source: §5 · p. 9
6 Conclusion
The conclusion presents LaCoT as a system for sampling latent visual reasoning chains and generating answers from them. It reports more diverse chains than GRPO and notes limits in scale, exploration, and hallucination coverage.
6 Conclusion
- LaCoT samples more diverse and comprehensive reasoning chains than the GRPO model.Figure 8 presents this qualitative comparison.Source: Figure 8 · p. 10
- LaCoT combines a rationale sampler for diverse, plausible visual reasoning trajectories with an answer generator enhanced by those chains.The conclusion also describes token-level reward approximation, RGFN training, and Bayesian inference scaling with statistically robust selection criteria.Source: §6 · p. 10
- Experiments are limited to models up to 7B parameters; complex-latent exploration remains open, and hallucination is not addressed.The authors say their methods can apply to autoregressive models of various sizes, but that is an expectation rather than an evaluated result here.Source: §6 · p. 10
A Proof of Proposition 1
The proof bounds the error from interpolating rewards between skipped steps. Linear reward growth gives exact interpolation, while bounded curvature yields quadratic error and flow consistency for sufficiently small segments.
A Proof of Proposition 1
- Under linear reward growth across the segment, interpolation reproduces every intermediate reward exactly.The proof states that the interpolation incurs zero error for every intermediate step.Source: A Proof of Proposition 1 · p. 14
- With a bounded discrete second derivative, the interpolation error decays as O(λ^2), and sufficiently small λ keeps it arbitrarily close to zero.The bound follows from the classical linear-interpolation error bound.Source: A Proof of Proposition 1 · p. 14
- For sufficiently small segments, the forward and backward product flows match up to a multiplicative factor that vanishes quadratically with λ.The proof bounds relative perturbations using the positive minimum reward Rmin.Source: A Proof of Proposition 1 · pp. 14–15
B.1 Qualitative results
The appendix reports qualitative comparisons of reasoning chains and examines how interpolation reward skipping affects supervision granularity and training time.
B.1 Qualitative results
- LaCoT-Qwen-7B provides a more accurate reasoning chain leading to the correct answer, while SFT and GRPO samples show incorrect visual CoT.Figures F9–F11 show LaCoT sampling more straightforward and accurate chains.Source: §B.1 · p. 15
B.2 Study of interpolation reward
- A smaller interpolation skip λ gives finer-grained reward supervision but increases training time.The study varies skipped steps in the policy-optimization reward approximation.Source: §B.2 · p. 15
B.3 Efficiency analysis
LaCoT-Qwen-7B is reported to balance stronger reasoning performance with modestly higher inference time. Mini-batching reduces the forward-pass cost of generating multiple rationales.
B.3 Efficiency analysis
- With batch size k=5, generating N rationales takes N/k forward passes.The method uses mini-batching to contain the additional inference cost of sampling multiple rationales.Source: §B.3 · p. 15
- LaCoT-Qwen-7B achieves stronger performance even with modest increases in inference time.The comparison reports per-sample reasoning-plus-answering time and performance on MathVista and MathVerse.Source: §B.3 · pp. 15–16
- Compared with other multi-rationale baselines, LaCoT is reported to offer a favorable balance between computational cost and reasoning reliability.The authors connect this balance with improved rationale trustworthiness and final-answer accuracy.Source: §B.3 · p. 16
B.4 Experiments compute resources
The appendix reports the compute setup and training duration, describes LaCoT’s exploration sampling, and identifies tables for inference-time and training-hyperparameter details.
B.4 Experiments compute resources
- Training used an 8*80GB GPU node; supervised fine-tuning took around 30 hours on 250k samples, while GRPO and RGFN fine-tuning took 120 hours on 3k samples.DeepSpeed Zero-3 and gradient checkpointing were used to reduce optimization memory costs.Source: §B.4 · p. 16
- Figure F11 presents qualitative visual-reasoning results and highlights important reasoning steps.Source: Figure F11 · p. 17
B.5 Hyperparameter
- LaCoT randomly samples mini-batch Zs for each (X, Y) during training as exploration; Table T7 lists training hyperparameters, and Table T6 studies inference time with multiple rational sampling.Source: §B.5 · p. 17, Table T7 · p. 18, Table T6 · p. 18