PaperlayerView exact revision on arXiv

Source-linked AI summary

Variational Reasoning for Language Models

Xiangxin Zhou, Zichen Liu, Haonan Wang, Chao Du, Min Lin, Chongxuan Li, Liang Wang, Tianyu Pang

arXiv:2509.22637v2cs.CLcs.AIcs.LG

TL;DR

SFT can require costly curated reasoning traces, while RL training can be unstable and lose output diversity. The paper treats thinking traces as latent variables and develops multi-trace variational bounds with forward-KL posterior training. It reports consistent gains over strong baselines across diverse reasoning benchmarks, alongside an analysis of accuracy-weighted bias in existing methods.

  • Problem

    SFT and RL have shown empirical success for reasoning, but SFT can rely on costly curated traces and RL can be unstable or suffer output-diversity collapse.

    Source: §1 · p. 1

  • Method

    The framework treats thinking traces as latent variables, extends the ELBO to multiple traces, and trains a hint-conditioned variational posterior using forward KL.

    Source: §1 · p. 1, §2.1 · p. 2, §2.3 · p. 4

  • Results

    The method consistently improves over strong baselines across diverse reasoning benchmarks, with average accuracy over 8.5% higher than the strong Bespoke-Stratos-4B† baseline using the same training data.

    Source: §1 · p. 2, §4.2 · p. 7

  • Takeaways & Limitations

    The analysis places RFT and binary-reward RL, including GRPO, under a shared probabilistic view and identifies their implicit accuracy weighting as a bias toward easier questions.

    Source: §1 · p. 2, §5 · p. 9

    The forward-KL derivation assumes qϕ(y′) is a fixed prior distribution and does not involve gradients.

    Source: §2.3 · p. 4

Abstract

from arXiv · show

We introduce a variational reasoning framework for language models that treats thinking traces as latent variables and optimizes them through variational inference. Starting from the evidence lower bound (ELBO), we extend it to a multi-trace objective for tighter bounds and propose a forward-KL formulation that stabilizes the training of the variational posterior. We further show that rejection sampling finetuning and binary-reward RL, including GRPO, can be interpreted as local forward-KL objectives, where an implicit weighting by model accuracy naturally arises from the derivation and reveals a previously unnoticed bias toward easier questions. We empirically validate our method on the Qwen 2.5 and Qwen 3 model families across a wide range of reasoning tasks. Overall, our work provides a principled probabilistic perspective that unifies variational inference with RL-style methods and yields stable objectives for improving the reasoning ability of language models. Our code is available at https://github.com/sail-sg/variational-reasoning.

1 INTRODUCTION

Supervised finetuning and reinforcement learning have advanced reasoning, but each has limitations that motivate a more principled training objective. The paper treats thinking traces as latent variables, develops a variational framework, and reports consistent gains over strong baselines across diverse reasoning benchmarks.

1 INTRODUCTION

  • SFT can depend on costly curated traces and struggle to generalize or avoid catastrophic forgetting, while RL can be unstable, collapse output diversity, and reduce Pass@K on harder questions.These challenges motivate a more principled objective for training reasoning models.Source: §1 · p. 1
  • The framework treats thinking traces as latent variables, using variational inference to derive tractable lower bounds and a posterior that samples paths more likely to yield correct answers.The approach remains compatible with verifiable rewards.Source: §1 · p. 1
  • The method decomposes reasoning into a trace and answer, extends the ELBO to an IWAE-style multi-trace bound, and uses a forward-KL objective to stabilize posterior training.The resulting pipeline jointly improves the reasoning model and variational posterior.Source: §1 · p. 1
  • RFT and binary-reward RL, including GRPO, can be expressed as accuracy-weighted local forward-KL objectives, revealing a systematic bias toward easier questions.The authors describe this bias as previously unrecognized.Source: §1 · pp. 1–2
  • The method shows consistent improvements over strong baselines across diverse reasoning benchmarks, including MATH500, AIME24&25, OlympiadBench, LiveCodeBench, GPQA-Diamond, and MMLU-Pro.The experiments use Qwen2.5 and Qwen3 model families.Source: §1 · p. 2

2 VARIATIONAL REASONING

The framework models a reasoning output as a thinking trace followed by an answer, then optimizes the probability of correct answers through variational bounds. It extends the bound to multiple traces and uses a forward-KL objective to train a variational posterior conditioned on an answer hint.

2 VARIATIONAL REASONING

  • The reasoning model generates a thinking trace z and answer y from question x, with their joint probability decomposed into answer-given-trace and trace-given-question terms.The framework seeks to maximize the marginal probability of producing a correct answer, but computing it requires summing over possible traces.Source: §2 · p. 2, §2.1 · p. 2

2.1 EVIDENCE LOWER BOUND

  • The ELBO makes correct-answer likelihood tractable by minimizing the reverse KL from the hint-conditioned variational posterior to the true posterior over traces that produce correct answers.The hint y′ is drawn from the oracle set of correct answers and encourages traces more likely to yield correct answers.Source: §2.1 · p. 2

2.2 EXTENSION TO IWAE-STYLE LOWER BOUND

  • The IWAE-style objective uses K traces to form a lower bound that tightens as K increases; when |Yx| ≫ 1, the accuracy-based estimator has lower worst-case variance.The likelihood-based estimator averages answer likelihoods over the correct-answer set, while the accuracy-based estimator samples answers and checks membership in that set. Long trace likelihood ratios can have high variance; per-token geometric-mean normalization reduces variance but introduces bias.Source: §2.2 · p. 3

2.3 OPTIMIZING THE VARIATIONAL POSTERIOR VIA FORWARD KL DIVERGENCE

  • To train the variational posterior, the method replaces reverse KL with forward KL, viewed as weighted SFT on traces sampled from the reasoning model.The change addresses the posterior's difficulty using answer hints without collapsing into shortcut reasoning; the forward-KL objective is an approximation, not a lower bound. The derivation assumes qϕ(y′) is a fixed prior without gradients, and experiments use a single training round.Source: §2.3 · p. 4

3 CONNECTION TO OTHER METHODS

The paper interprets rejection-sampling finetuning and binary-reward RL as local forward-KL updates toward correct-answer-conditioned trace distributions. Their accuracy weighting favors easier questions, while the proposed variational objective treats questions more evenly.

3 CONNECTION TO OTHER METHODS

  • RFT is equivalent to maximizing a forward-KL objective toward the reference model’s correct-answer-conditioned trace distribution, weighted by reference-model accuracy.This weighting downplays hard questions with low reference accuracy and biases training toward easier ones; the paper contrasts it with its more even variational objective.Source: §3 · p. 5
  • For binary-reward RL, the local gradient for the thinking-trace policy is equivalent to minimizing forward KL toward the true correct-answer posterior, with updates weighted by model accuracy.The binary-reward RL objective itself maximizes the probability of a correct answer, whereas the variational framework maximizes its logarithm.Source: §3 · p. 5
  • GRPO retains the local forward-KL interpretation under 0–1 rewards, with a weight that increases monotonically with model accuracy and therefore emphasizes easier questions.The analysis omits GRPO’s min and clip operations for brevity.Source: §3 · p. 5

4 EXPERIMENTS

The experiments evaluate variational reasoning across Qwen model families and reasoning benchmarks, reporting stronger results than baselines and ablations of its training choices.

4 EXPERIMENTS

  • Experiments use Bespoke-Stratos-17k, compare against Bespoke-Stratos and other reasoning baselines, and evaluate math, coding, science, and general-knowledge benchmarks.GPQA-Diamond and MMLU-Pro are out-of-distribution relative to the math-and-code training data.Source: §4 · pp. 5–6

4.2 MAIN RESULTS

  • Variational reasoning surpasses all baselines in average accuracy, exceeding Bespoke-Stratos-4B† by over 8.5% and by over 14% in other domains.The reported Qwen3 results include substantial improvements over the base model in math and other domains.Source: Table 1 · p. 6, Table 2 · p. 6, §4.2 · p. 7
  • During Qwen3-Base training, the method has lower average loss and fewer gradient-norm spikes than Bespoke-Stratos, indicating greater training stability.The authors attribute the stability to adaptive weighting by trace quality and alignment with the reasoning policy.Source: Figure 1 · p. 7, §4.2 · p. 9

4.1 TRAINING DETAILS FOR VARIATIONAL REASONING

  • Training independently fits an initial reasoning model and variational posterior, then uses posterior-generated responses and weights to train the final model in 17K and 1K settings.The 17K setting retains the highest-weight generated response and original sample; the 1K setting uses all eight generated responses per sample.Source: §4.1 · p. 7

4.2 MAIN RESULTS

  • As Pass@K increases, the method's advantage grows on complex benchmarks but diminishes on simpler tasks and multiple-choice questions.The comparison covers Qwen3-4B/8B-Base models.Source: Figure 2 · p. 8, §4.2 · p. 8
  • The method significantly outperforms General-Reasoner on GPQA-Diamond and MMLU-Pro, which are out-of-distribution for its math-and-code training data.Those benchmarks are in-domain for General-Reasoner-4B.Source: §4.2 · p. 8
  • Performance remains consistent across prompt templates; accuracy-based and geometric-mean estimators perform similarly across four scales, with a slight accuracy-based advantage on math.Source: §4.2 · p. 8

4.3 ABLATION STUDIES

  • Increasing trace count K can improve performance at added computational cost; removing y′ conditioning reduces performance, while accuracy and geometric-mean estimators outperform naive likelihood.The K sweep runs from 1 to 32; the estimator comparison uses the 1K data setting.Source: Table 4 · p. 8, Figure 3 · p. 9, §4.3 · p. 9

5 CONCLUSION

The conclusion presents the variational reasoning framework as a stable training objective with gains over strong baselines, and the appendix derives its objectives and analyzes reward-shaping behavior. It identifies extending training beyond one round and enriching the answer-hint posterior as future directions.

5 CONCLUSION

  • The framework reports consistent gains over strong baselines on diverse reasoning tasks and clarifies biases in existing SFT/RFT and RL methods.The authors identify extending training beyond a single round and exploring richer answer-hint posterior designs as future directions.Source: §5 · p. 9

A.1 DERIVATION FOR EQ. (4)

  • The appendix rewrites the ELBO objective induced by Eq. (2).It also identifies Pθ(Yx|x) as the true posterior distribution.Source: §A.1 · p. 15, §A.2 · p. 15

A.2 DERIVATION FOR EQ. (6)

  • The appendix derives gradients for model parameters and the variational posterior from the ELBO and IWAE-style lower bound.It also derives the forward-KL gradient for the variational posterior.Source: §A.2 · p. 15, §A.4 · p. 16

A.3 PROOF OF THEOREM 1

  • The accuracy-based estimator’s variance is independent of the number of correct answers and their probability distribution; the stated worst-case condition almost always holds when |Yx| ≫ 1.The estimator is defined as expected correctness over answers sampled from the model.Source: §A.3 · pp. 15–16

A.5 CONNECTION TO MORE GENERAL RL REWARD SHAPING

  • With α > β > γ, reward shaping can favor reward-hacking modes on hard questions that are easy to format correctly.When β = γ, the optimization reduces to the no-format-reward case.Source: §A.5 · p. 17

A.6 SPECIAL CASES IN EQ. (6)

  • The analysis states that a unique-correct-answer setting is a special case of Eq. (6) and relates its estimator to normalized VeriFree.The passage labels this relationship as a special case.Source: §A.6 · p. 17

A.7 MORE DERIVATIONS FOR EQ. (11)

The appendix notes that the ideal solution requires perfect answer correctness independent of the thinking trace, then derives a low-variance behavior policy from the variational posterior. It also connects the framework to Reinforcement Learning Teachers, reviews related reasoning approaches, and describes a sentence-level training adjustment.

A.7 MORE DERIVATIONS FOR EQ. (11)

  • The optimal solution is usually unattainable because it requires 100% correct answers regardless of the thinking trace.This is the stated condition for achieving the optimum.Source: §A.7 · p. 17

A.8 APPLYING VARIATIONAL POSTERIOR FOR RL

  • Using the trained variational posterior as the behavior policy approximately minimizes the variance of the binary-reward RL objective estimator.The passage says this policy is compatible with actor-critic frameworks using advantage estimation and baselines.Source: §A.8 · p. 18

A.9 CONNECTION TO REINFORCEMENT LEARNING TEACHERS

  • The paper gives a variational-inference justification for Reinforcement Learning Teachers and adds a tighter IWAE-style bound and an accuracy-based estimator.RLTs train reasoning models as teachers for student distillation using GRPO and dense rewards [Cetin et al., 2025].Source: §A.9 · p. 19

B RELATED WORK

  • The related-work discussion situates the method among latent-trace approaches and distinguishes its forward-KL posterior training, tighter bound, and accuracy-based estimator.It also connects the framework to RFT and GRPO.Source: B RELATED WORK · pp. 19–20

C TRAINING DETAILS

  • Dr. SFT sums token losses and divides by the precomputed average response length to align training with sentence-level reasoning and answer variables.The same modified objective is used to train the paper’s models and Bespoke-Stratos-4B/8B†.Source: C TRAINING DETAILS · p. 20

C.1 TRAINING THE INITIAL REASONING MODEL AND VARIATIONAL POSTERIOR

The training details specify separate configurations for the initial reasoning model and variational posterior. Posterior samples are weighted using model and posterior likelihoods, with alternative estimators used to assess correctness likelihood.

C.1 TRAINING THE INITIAL REASONING MODEL AND VARIATIONAL POSTERIOR

  • The initial reasoning model is trained for 3 epochs with batch size 96, while the variational posterior is trained for 10 epochs with batch size 16.Both use AdamW; the initial model uses learning rate 1.0e-5 and the posterior follows its other hyperparameters.Source: §C.1 · p. 20

C.2 DETAILS OF SAMPLING FROM VARIATIONAL POSTERIOR

  • For each question, the trained posterior samples 8 reasoning traces, and importance weights are computed from likelihoods under the initial model and posterior.The likelihoods are evaluated in forward mode without backpropagation.Source: §C.2 · p. 20
  • The “-Acc” estimator averages correctness across 8 verifier-checked answers per sampled trace; “-GML” uses the geometric mean of token probabilities to mitigate length bias without an external verifier.The naive estimator is labeled “-L”.Source: §C.2 · pp. 20–21

C.3 DETAILS OF TRAINING FINAL REASONING MODEL

Final-model training differs by data setting: the 17k setting selects one high-weight trace per question and mixes synthetic data with the original set, whereas the 1k setting uses all weighted traces without that mixture.

C.3 DETAILS OF TRAINING FINAL REASONING MODEL

  • In the 17k setting, training selects the highest-weight of 8 posterior-sampled traces per question and mixes the resulting synthetic data with the original dataset.With the accuracy-based estimator, the selected trace is paired with a randomly chosen verified answer.Source: §C.3 · p. 21
  • For 17k training, batch size increases to 192 so the optimization-step count matches baseline models despite the doubled mixed-dataset size.For the 1k method setting, batch size is adjusted proportionally to match baseline optimization steps while other hyperparameters remain unchanged.Source: §C.3 · p. 21
  • In the 1k setting, the method uses all 8 traces weighted by eρk and does not combine them with the original dataset.The baseline uses batch size 16 and cutoff length 32768.Source: §C.3 · p. 21

D DETAILS OF EVALUATION

Evaluation uses sampled responses with reproducibility-focused inference settings, while practical prompt templates structure the reasoning and answer in distinct sections.

D DETAILS OF EVALUATION

  • Responses are sampled at temperature 0.7 with top_p=1.0 and a 38,912-token budget; inference uses float32 and vLLM with model-size-specific tensor parallelism.The vLLM versions differ by model family: 0.7.0 for Qwen2.5 and 0.8.4 for Qwen3.Source: D DETAILS OF EVALUATION · p. 21

E PROMPT TEMPLATES

  • The practical templates elicit a solution plus a detailed reasoning trace, with formats labeling that trace as Explanation or Thought and placing the sections in different orders.The templates also vary in whether they reconstruct reasoning from a provided correct answer or ask the model to solve the question.Source: E PROMPT TEMPLATES · pp. 22–23
  • Table 6 reports Qwen2.5-7B-Instruct models trained on Bespoke-Stratos-17k, with columns for GPQA-D, LCB-E, LCB-M, LCB-H, MMLU-Pro, and averages.The supplied table text gives its scope and column labels, but not cell values.Source: Table 6 · p. 24

F.1 EXTENDED MAIN RESULTS

Extended results report that variational-reasoning variants outperform baselines on average accuracy and show length-related correlations. The section also compares against officially reported results and outlines further ablations.

F.1 EXTENDED MAIN RESULTS

  • All variational-reasoning variants outperform all baselines in average accuracy, and the two variational-posterior prompt templates perform similarly.The authors attribute template robustness to fine-tuning the posterior rather than relying on prompt engineering alone.Source: §F.1 · p. 24
  • Thinking-trace and answer lengths correlate strongly with their respective likelihood measures, indicating length biases and motivating accuracy- or geometric-mean-based estimators over naive likelihood.Figure 4 plots thinking-token length against the trace likelihood ratio and answer-token length against answer log-likelihood.Source: §F.1 · p. 24, Figure 4 · p. 25
  • The General-Reasoner-4B results differ from Ma et al. (2025), alongside differences in evaluation framework, sampling temperature, and token budget.This comparison uses SkyThought, temperature 0.7, and a 38,912-token budget, unlike the reported simple-evals, mostly greedy decoding, and 8,192-token budget.Source: §F.1 · p. 24
  • On benchmarks shared with Ma et al. (2025), the method still outperforms General-Reasoner-4B‡; Table 7 identifies the comparison as using their officially reported accuracy.The table covers Qwen3-4B-Base models, with General-Reasoner-4B as an exception to the Bespoke-Stratos-17k training-data setup.Source: §F.1 · p. 24, Table 7 · p. 25

F.2 EXTENDED ABLATION STUDIES

  • The extended ablations examine training-data sources and usage, overlap between posterior and reasoning-model training, Dr. SFT versus naive SFT, and training cutoff length.Source: §F.2 · p. 25

F.3 EFFECTS OF DIFFERENT FINAL DATA SOURCES AND WAYS OF DATA USAGE

Ablations compare trace-selection and data-mixing choices, data overlap, and two SFT objectives. Weighted multi-trace training without mixing performs best in the 1k setting, while overlap and objective variants show comparable accuracy.

F.3 EFFECTS OF DIFFERENT FINAL DATA SOURCES AND WAYS OF DATA USAGE

  • In the 1k setting, weighted multiple traces without original-data mixing perform best; mixing helps single-best trace selection but slightly degrades weighted multi-trace performance.The main 17k setup instead selects the highest-weight trace and mixes it with original data for efficiency; the ablation compares this strategy with weighted multiple traces.Source: §F.3 · pp. 25–26, Table 8 · p. 26

F.4 EFFECTS OF DATA OVERLAP

  • Overlap and non-overlap variants have similar average accuracy, and both outperform the Bespoke-Stratos-7B-1K† baseline.For the non-overlap setting, the posterior and initial reasoning model are trained on 16k samples, then the final model is trained on a separate 1k subset using weighted multiple traces.Source: §F.4 · p. 26, Table 9 · p. 27

F.5 COMPARING DR. SFT WITH NAIVE SFT

  • Dr. SFT and naive SFT differ by less than 2% in average accuracy, and both outperform the baseline; the authors attribute the primary improvement to variational reasoning.Dr. SFT sums valid-token losses normalized by a constant, rather than averaging valid-token losses across the batch.Source: §F.5 · pp. 26–27

F.6 ABLATION STUDY ON THE EFFECTS OF CUTOFF LENGTH WHEN TRAINING

The cutoff-length ablation compares 16k and 32k training cutoffs and finds similar performance and inference lengths, supporting the more efficient 16k setting.

F.6 ABLATION STUDY ON THE EFFECTS OF CUTOFF LENGTH WHEN TRAINING

  • The ablation compares cutoff_len=16384 (-Len16k) with cutoff_len=32768 (-Len32k) in the 17k data setting.Source: §F.6 · p. 27, Table 11 · p. 28
  • The 16k and 32k cutoff variants show similar performance and similar inference generation lengths, while both produce longer traces than Qwen2.5-7B-Instruct.The supplied Table 11 caption identifies the cutoff-length ablation, and Table 12 is described as reporting average completion token lengths.Source: §F.6 · p. 27, Table 11 · p. 28, Table 12 · p. 28
  • Increasing the cutoff beyond 16k has minimal effect on output, so the authors use 16k for better training efficiency without sacrificing performance.Source: §F.6 · p. 27

G EXAMPLES OF GENERATED RESPONSES

The examples identify an equiangular parallelogram as a rectangle and show responses from the variational posterior and a final model. The geometry response checks the answer against the biconditional and the alternatives.

G EXAMPLES OF GENERATED RESPONSES

Related papers

Loading 2509.22637v2…