AMVL trains a multimodal LLM to reason in a continuous latent space by adding a reverse KL divergence term that keeps the answer-aware training Posterior (variational) from drifting into regions the answer-blind inference Prior (variational) cannot reach, closing the train/test gap that plagues variational latent reasoning.
Suppose you’re building a visual QA agent that inspects diagrams, screenshots, or medical images. The dominant recipe is Chain-of-Thought: force the model to verbalize each intermediate step as text tokens. That serialization is lossy for spatial and perceptual content, and it’s where hallucinations creep in. A newer line of work (LVR, Monet, Mull-Tokens) skips the text and reasons through a handful of continuous latent vectors instead. Those methods work, but each one hand-designs what the latent should encode via reconstruction or alignment losses. This paper argues that supervision should come from the task itself, through variational inference, and identifies why the naive version of that idea fails.
The setup is a standard conditional Variational Autoencoder wrapped around an MLLM. You insert k placeholder tokens between the prompt and the answer. Two lightweight heads read the hidden states at those positions and output Gaussian parameters for a prior (which sees only the input) and a posterior (which also sees the ground-truth answer). At training time you sample from the posterior; at inference time you sample from the prior.
The problem the paper names is answer leakage. The posterior can cheat: it encodes information from the answer directly into the latent, gets low reconstruction loss, and the standard Evidence Lower Bound then pulls the prior toward this cheating posterior. At test time the prior has to hallucinate that leaked information. The authors prove formally that one-sided ELBO training bakes the average answer-dependent shift into the prior as a residual bias.
The fix, AMVL, adds a second KL divergence term in the opposite direction, with stop-gradients so each term updates only one side:
# per training step
z_post = sample(q_phi(z | x, y)) # posterior sample
L_ntp = -log p_theta(y | x, z_post) # next-token loss
L_fwd = KL(stopgrad(q_phi) || p_theta) # updates prior only
L_rev = KL(stopgrad(p_theta) || q_phi) # updates posterior only
loss = L_ntp + beta * L_fwd + gamma * L_rev
The forward term pulls the prior toward the posterior. The reverse term pulls the posterior back toward regions the prior can actually cover. Both KLs are optimized on their second argument, which gives them a mass-covering (rather than mode-seeking) effect, so neither distribution is allowed to collapse away from the other. The reverse weight is annealed in later than the forward weight so the prior gets a head start before it starts constraining the posterior.
The usual variational recipe treats the posterior as a fixed teacher and asks the prior to imitate it. This paper shows the opposite. When the teacher has privileged information (the answer), imitation alone transfers that privilege as a phantom signal the student can never reproduce at test time; the teacher itself must be regularized to stay inside the student’s reachable region. The clean evidence is the latent-spread analysis: forward-only training leaves a wide gap between prior and posterior geometry, while adding the reverse term shrinks that gap and lifts downstream accuracy.
The load-bearing result is the objective ablation on three perception benchmarks. Forward-KL alone actually hurts the base model on V* (40.84 vs 81.15 for next-token-only), because it drags the prior toward a leaky posterior. Reverse-KL alone under-performs the baseline slightly. Only the combined objective, with forward warmed up before reverse kicks in, beats plain next-token training on all three (84.29 / 72.12 / 68.50). Order matters too: reverse-first scheduling drops V* to 80.63.
Secondary evidence, all built on Qwen2.5-VL-7B:
•
On the BLINK perception suite, AMVL lifts the average by +10.83 over the base model, with a +32.00 jump on the Jigsaw task and +12.67 on IQ Test.
•
On V*, HRBench 4K and 8K, the average gain over Qwen2.5-VL-7B is +5.57, beating both discrete-reasoning baselines (Vision-R1, PAPO, PixelReasoner, DeepEyes) and prior latent-reasoning methods (Latent Visual Reasoning (LVR), Mull-Tokens, Monet).
•
Out-of-distribution on VisualPuzzles, AMVL scores 33.90 overall and 52.50 on deductive reasoning, beating a 72B baseline.
•
Latent-spread analysis (Appendix H) shows the paired prior-posterior L2 distance drops from 16.95 (no KL) to 5.43 (AMVL), with cosine similarity rising from roughly zero to 0.89.
•
Ablations pin the sweet spot at k=8 latent slots, d=512; more slots hurt because continuous vectors are information-dense and extra capacity dilutes the signal.
Reach for this when you’re building a vision-language system where the intermediate reasoning is inherently spatial or perceptual (diagram QA, chart reading, visual search, jigsaw-like puzzles) and you’ve noticed that forcing a text chain-of-thought loses grounding. AMVL is a drop-in modification to any Qwen2.5-VL-style MLLM: add k placeholder tokens, a small variational head, and two KL terms with stop-gradients. The vision encoder stays frozen; only the language backbone and the new heads train.
The paper does not link to a public code release in the provided text. Training was done on 16 A100s for about 20 hours on a mix of Visual-CoT, ReFocus, CogCoM, and Zebra-CoT. If you want to replicate, the recipe (k=8, d=512, forward KL warmup over 2000 steps to weight 1.0, reverse KL delayed 1000 steps then ramped to 0.5) is fully specified in Appendix E.
When a training-time signal has access to information the inference-time signal doesn’t, don’t just teach the deployed model to imitate it. Constrain the privileged teacher to stay inside the reachable world of the student. Every conditional variational setup with an asymmetric posterior has this pathology latent in it; a second KL in the opposite direction is a cheap and principled fix.
•
All experiments are at 7B. The paper explicitly flags scaling to 70B+ as future work, so it’s unknown whether the mismatch gets worse, better, or qualitatively different at frontier scale.
•
Gains are largest on perception-heavy and spatial tasks (Jigsaw, V*, deductive puzzles). On tasks where text CoT already grounds well, the delta is smaller and the extra machinery may not pay for itself.
•
The theoretical results (Propositions C.1-C.4) rely on diagonal-Gaussian latents and a local linear-response assumption around a stationary point. Real posteriors that leak through variance or higher-order statistics, rather than through the mean, are outside the analysis, and the authors acknowledge this restriction.