Variational Proximal Policy Optimization

Abstract: Reinforcement Learning from Human Feedback via Proximal Policy Optimization often suffers from policy mode collapse, brittle exploration loops, and distribution drift. This paper introduces Variational Proximal Policy Optimization ((\textsc{VP}_2\textsc{O})), a particle-based variational inference framework that maps policy optimization to Stein Variational Gradient Descent within a Mixture-of-Experts architecture. By leveraging functional kernels over localized expert prototypes alongside an expert orthogonalization loss, (\textsc{VP}_2\textsc{O}) introduces a geometry-based proximal-control mechanism that can reduce reliance on fixed clipping or KL schedules. Our results on a 33B/4B sparse Mixture-of-Experts model show several improvements across complex reasoning benchmarks, establishing a (+\mathbf{179}) ELO gain on Codeforces and a (\mathbf{32\%}) reduction in token count on AIME mathematical reasoning tasks.

• The paper introduces VP²O, reframing RLHF as a variational inference problem using SVGD for enhanced policy improvement.
• It replaces heuristic clipping and KL penalties with geometric trust regions over Mixture-of-Experts policy particles.
• Experimental results show improved performance, stability, and token efficiency across mathematical reasoning, code generation, and instruction tasks.

Variational Proximal Policy Optimization: A Variational Inference Perspective on RLHF with Mixture-of-Experts

Motivation and Problem Formulation

Reinforcement Learning from Human Feedback (RLHF), typically instantiated with Proximal Policy Optimization (PPO), remains the canonical choice for model alignment in LLMs. However, standard PPO- and GRPO-based RLHF algorithms exhibit several persistent limitations: policy mode collapse, brittle exploration, and susceptibility to reward model overfitting and distributional drift. These constraints are further exacerbated by the ad-hoc nature of trust region enforcement (clipping, KL penalties), which often yield unstable optimization dynamics and reduced behavioral diversity.

This work introduces Variational Proximal Policy Optimization (VP₂O), which recasts KL-regularized RLHF as a particle-based variational inference problem, leveraging Stein Variational Gradient Descent (SVGD) over Mixture-of-Experts (MoE) policy manifolds. VP₂O establishes a direct geometric correspondence between expert diversity and regularized policy improvement, offering structural proximal control that reduces dependence on heuristic exploration bonuses and clipping schedules.

Variational Inference and Stein Gradient Formulation

The central insight in VP₂O is the reframing of PPO as minimizing $D_{\text{KL}}(\pi_\theta\|p^*)$, where $p^*$ is the Boltzmann policy induced by the KL-regularized reward. Instead of relying on importance sampling and scalar trust region heuristics, the method transports an ensemble of policy particles toward high-reward support via SVGD. For policy parameter distribution $q$ with target $p^*$, SVGD defines a functional transport field optimizing Stein discrepancy, with attraction (gradient) and repulsion (covariance) terms acting on a population of experts.

The updates are realized in the output space using low-dimensional expert prototypes (principal eigenvectors of expert output projections), enabling efficient and robust kernel geometries, and capturing policy similarity and diversity at the expert granularity. The Stein field is decomposed according to the MoE’s routing structure, with co-activated experts exchanging information (attraction) and inactive experts repelled in prototype space.

Mixture-of-Experts as Policy Particle Manifold

VP₂O leverages Transformer-based MoE models, treating each expert as a distinct policy particle. Top-$K$ routing mechanisms induce conditional computation regimes, promoting specialization and coverage. The flexible SVGD update moves expert parameters in functional output space rather than high-dimensional parameter space, enhancing diversity and mitigating collapse. Experts are regularized to be orthogonal via a specialized loss, and routing diversification losses further promote functional disentanglement and utilization.

Notably, the Stein transport field is constructed with behavioral/directional similarity kernels among prototypes, parameterized with a temperature scaling relative to the number of experts. This encapsulates both the diversity-promoting repulsion and reward-driven attraction of the SVGD dynamics.

Geometric Trust Regions and On-Policy Diagnostics

A key distinction from PPO/GRPO is that VP₂O replaces fixed clipping and scalar KL bounds with geometric controls in prototype barycenter space. Updates are accepted only if the learner’s barycenter (in the current data batch) remains sufficiently close to the actor’s anchor snapshot, providing a data- and routing-driven trust region. Further, synchronization of actor and learner rollouts is governed by on-policy behavior budgets using explicit KL drift and Effective Sample Size (ESS) criteria, rather than periodic or fixed refresh intervals.

This controls policy drift while avoiding the variance inflation and gradient distortion characteristic of token-wise importance sampling and excessive clipping.

Experimental Evidence and Numerical Benchmarks

The method is demonstrated on 33B/4B MoE architectures with 20 experts per FFN layer, using GRPO (the strongest available PPO-variant) as the baseline. VP₂O achieves consistent improvements across a suite of challenging benchmarks:

• Mathematical Reasoning: On AIME 2024/2025, VP₂O improves by +2.6/+2.8 pp at 8K context, converges faster (by ~2000 steps) at 16K, and consistently surpasses the baseline.
• Scientific and Knowledge Reasoning: On GPQA Diamond and MMLU-Pro, gains are more prominent at 16K context, achieving +1.8 and +1.1 pp, respectively.
• Code Generation: VP₂O yields a maximal +179 ELO and +3.6 Pass@1 lift on Codeforces at 16K, confirming efficacy in tasks with multi-modal solution structure.
• Instruction Following: Gains on IFBench and IFEval are robust and consistent (+4-5 pp at 8K, increasing further at 16K).

Additionally, VP₂O demonstrates solution token efficiency: at 16K, solution length for AIME drops by 59 tokens (−11%), and for IFEval by 59 tokens (−22%) at equal or improved accuracy. For Codeforces, solution tokens decrease by 135 at 16K with a significant performance boost. The method does not exhibit late-stage degradation, unlike baseline PPO variants.

Practical and Theoretical Implications

VP₂O provides a principled methodology for RLHF grounded in functional variational inference rather than heuristic trust region engineering. The approach furnishes diversity control, exploration, and specialization, which address core failure modes of PPO-style RLHF—mode collapse, reward overoptimization, and insufficient coverage. The MoE-Stein hybridization enables scalable and efficient particle approximation of the optimal RLHF policy, with geometric controls that naturally adapt to routing and expert utilization statistics.

Computationally, the primary overhead arises from prototype updates and kernel diagnostics (5–8% wall-clock), which is modest compared to the gains in stability and sample efficiency. However, extension to regimes with larger expert counts or different backbone architectures remains an open problem.

Future Directions

Future lines of research involve:

• Scaling the method to larger MoE architectures ($N \gg 20$) and diverse Transformer families.
• Incorporating richer diagnostics of expert overlap, routing entropy, and drift statistics.
• Extending the framework to off-policy RLHF and reward model uncertainty settings.
• Studying the integration of nonparametric and adaptive kernel constructions for SVGD in ultra-high dimensional settings.
• Systematic ablation of orthogonalization and routing regularization loss impacts.

Conclusion

VP₂O establishes a systematic bridge between RLHF policy optimization and variational inference, with SVGD-based ensemble dynamics instantiated in MoE models. The approach achieves both empirical gains and improved training stability by replacing ad-hoc trust region heuristics with kernelized geometric controls and diversity-aware particle optimization. The implications are directly relevant for scaling RLHF to longer sequence lengths and more complex tasks. While the experiments validate VP₂O in a moderately large MoE stack, further work is necessary to examine generalization to broader architectures and data domains, as well as to integrate uncertainty quantification and off-policy extensions (2606.08032).