Latent path exploration improves reasoning in language models
Latent Path Exploration via Potential-Guided Subtrajectory Balance
Computation and LanguageMachine Learning
Summary
Language models can solve problems by thinking through steps inside their 'mind' without saying each step out loud. The authors found a way to better explore different possible thinking paths using a technique that scores each step based on how well it might lead to the right answer. Their approach helps the model try many paths and pick the better ones, especially in math and logic puzzles. This leads to more correct answers, even when tested on harder or unfamiliar problems.
What this means in practice
- •For natural language processing developers: Enable language models to better explore intermediate reasoning steps to improve accuracy on complex arithmetic and logical tasks.
- •For ai system engineers: Incorporate guided latent path exploration to enhance robustness of reasoning modules under domain shifts or out-of-distribution challenges.
- •For automated tutoring platforms: Use improved latent reasoning techniques to provide more accurate step-by-step problem-solving guidance in educational applications.$Commercial implications: This paper enables smarter AI tutors that reason more effectively, potentially increasing their market value and adoption.
Authors
Xiaoling Zeng, Qi Zhang
Abstract
Continuous latent reasoning enables language models to perform intermediate computation without generating explicit reasoning traces. Advances in stochastic latent reasoning allow multiple paths to be sampled in continuous space, but stochasticity alone does not ensure effective exploration. A key challenge is guiding stochastic transitions between intermediate latent states toward useful reasoning paths. We introduce SubLAT, which couples a relative answer potential with Subtrajectory Balance to learn stochastic latent exploration. The relative potential evaluates the answer support of sampled states against replacement states, with a correctness signal at the terminal state. Subtrajectory Balance relates potential differences to forward and backward transition densities across subtrajectories of different lengths. The model is jointly trained with answer supervision, CoT supervision, and representation distillation. Experiments on arithmetic and deductive reasoning show increasing solution coverage with larger sampling budgets across in-domain and out-of-domain benchmarks, as well as improvements on several structural OOD benchmarks. Trajectory analyses reveal progressive divergence among sampled latent paths and higher mean relative potentials for correct trajectories across arithmetic benchmarks.