Beyond KL Divergence: Policy Optimization with Flexible Bregman Divergences for LLM Reasoning

2026-02-04 · Rui Yuan, Mykola Khandoga, Vinay Kumar Sankarapu arxiv

Policy optimization methods like Group Relative Policy Optimization (GRPO) and its variants have achieved strong results on mathematical reasoning and code generation tasks. Despite extensive exploration of reward processing strategies and training dynamics, all existing group-based methods exclusively use KL divergence for policy regularization, leaving the choice of divergence function unexplored. We introduce Group-Based Mirror Policy Optimization (GBMPO), a framework that extends group-based policy optimization to flexible Bregman divergences, including hand-designed alternatives (L2 in probability space) and learned neural mirror maps. On GSM8K mathematical reasoning, hand-designed ProbL2-GRPO achieves 86.7% accuracy, improving +5.5 points over the Dr. GRPO baseline. On MBPP code generation, neural mirror maps reach 60.1-60.8% pass@1, with random initialization already capturing most of the benefit. While evolutionary strategies meta-learning provides marginal accuracy improvements, its primary value lies in variance reduction ($\pm$0.2 versus $\pm$0.6) and efficiency gains (15% shorter responses on MBPP), suggesting that random initialization of neural mirror maps is sufficient for most practical applications. These results establish divergence choice as a critical, previously unexplored design dimension in group-based policy optimization for LLM reasoning.

📄 PDF Abstract BibTeX arXiv:2602.04380

Code (0)

등록된 구현이 없습니다.

Tasks

Mathematical Reasoning Code Generation

Results from the Paper

Rank	Task	Dataset	Model	Metrics
#101	Code Generation	MBPP	Group-Based	Accuracy: 5.5

Bregman Gradient Policy Optimization

2021-06-23 · ICLR 2022 4 · Feihu Huang, Shangqian Gao, Heng Huang

In the paper, we design a novel Bregman gradient policy optimization framework for reinforcement learning based on Bregman divergences and momentum techniques. Specifically, we propose a Bregman gradient policy optimizat…

reinforcement-learning Reinforcement Learning Reinforcement Learning (RL)

Divergence-Augmented Policy Optimization

2025-01-25 · NeurIPS 2019 12 · Qing Wang, Yingru Li, Jiechao Xiong, Tong Zhang

In deep reinforcement learning, policy optimization methods need to deal with issues such as function approximation and the reuse of off-policy data. Standard policy gradient methods do not handle off-policy data well, l…

Atari Games Deep Reinforcement Learning Policy Gradient Methods reinforcement-learning+1

On the Convergence of Policy in Unregularized Policy Mirror Descent

2022-05-17 · Dachao Lin, Zhihua Zhang

In this short note, we give the convergence analysis of the policy in the recent famous policy mirror descent (PMD). We mainly consider the unregularized setting following [11] with generalized Bregman divergence. The di…

BSO: Safety Alignment Is Density Ratio Matching

2026-05-12 · Tien-Phat Nguyen, Truong Nguyen, Thin Nguyen, Duy Minh Ho Nguyen 외 arxiv

Aligning language models for both helpfulness and safety typically requires complex pipelines-separate reward and cost models, online reinforcement learning, and primal-dual updates. Recent direct preference optimization…

Reinforcement Learning

On the connection between Bregman divergence and value in regularized Markov decision processes

2022-10-21 · Brendan O'Donoghue

In this short note we derive a relationship between the Bregman divergence from the current policy to the optimal policy and the suboptimality of the current value function in a regularized Markov decision process. This …