From Noisy Traces to Stable Gradients: Bias-Variance Optimized Preference Optimization for Aligning Large Reasoning Models
Mingkang Zhu, Xi Chen, Bei Yu, Hengshuang Zhao, Jiaya Jia
TL;DR
This work tackles the instability in aligning Large Reasoning Models (LRMs) that stems from trace-based gradient variance when reasoning traces are sampled during training. It introduces Bias–Variance Optimized Preference Optimization (BVPO), which forms a gradient mix $g_c = \alpha g_t + (1-\alpha) g_e$ between a high-variance trace-based gradient $g_t$ and a low-variance empty-trace gradient $g_e$, and derives the MSE-optimal mixing weight to minimize bias and variance. The authors prove that BVPO reduces conditional variance from trace sampling and provides convergence guarantees for SGD with the optimal mixing, linking statistical optimality to algorithmic performance. Empirically, BVPO improves alignment on AlpacaEval 2 and Arena-Hard and enhances reasoning on six math benchmarks, without sacrificing base-model reasoning capabilities. Overall, BVPO identifies trace-sampling variance as a key bottleneck and demonstrates that bias–variance trade-off optimization yields stable training and stronger performance in LRMs.
Abstract
Large reasoning models (LRMs) generate intermediate reasoning traces before producing final answers, yielding strong gains on multi-step and mathematical tasks. Yet aligning LRMs with human preferences, a crucial prerequisite for model deployment, remains underexplored. The statistically correct objective for preference alignment requires marginalizing over reasoning traces, but this computation is intractable in practice. A common workaround optimizes a single sampled trajectory, which introduces substantial gradient variance from stochastic trace sampling. To address this challenge, we frame preference optimization for LRMs through the lens of the bias--variance trade-off and propose Bias--Variance Optimized Preference Optimization (BVPO), a simple, drop-in method that mixes two gradient estimators: a high-variance trace-based estimator and a low-variance empty-trace estimator obtained by disabling reasoning trace generation. Our theory shows that BVPO strictly reduces trace-induced variance for any nontrivial mixture, provides a closed-form choice of the mixing weight that minimizes mean-squared error relative to the true marginal gradient, and under standard smoothness and step-size conditions, tightens classical convergence bounds for stochastic gradient descent. Empirically, BVPO improves alignment over the best baseline by up to 7.8 points on AlpacaEval~2 and 6.8 points on Arena-Hard. Despite being trained only on general conversational data, BVPO also boosts reasoning performance for base models by up to 4.0 points on the average of six math reasoning benchmarks. These results identify variance from trace sampling as a key bottleneck and demonstrate that directly optimizing the bias--variance trade-off yields more stable training and stronger overall performance.