Papers
Topics
Authors
Recent
Search
2000 character limit reached

BSO: Safety Alignment Is Density Ratio Matching

Published 12 May 2026 in cs.LG and cs.AI | (2605.12339v1)

Abstract: Aligning LLMs 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 approaches simplify training but incorporate safety through ad-hoc modifications such as multi-stage procedures or heuristic margin terms, lacking a principled derivation. We show that the likelihood ratio of the optimal safe policy admits a closed-form decomposition that reduces safety alignment to a density ratio matching problem. Minimizing Bregman divergences between the data and model ratios yields Bregman Safety Optimization (BSO), a family of single-stage loss functions, each induced by a convex generator, that provably recover the optimal safe policy. BSO is both general and simple: it requires no auxiliary models, introduces only one hyperparameter beyond standard preference optimization, and recovers existing safety-aware methods as special cases. Experiments across safety alignment benchmarks show that BSO consistently improves the safety-helpfulness trade-off.

Summary

  • The paper presents a novel BSO framework that reframes safety alignment for LLMs as a density ratio matching problem using Bregman divergence minimization.
  • It unifies prior heuristic methods and eliminates the need for auxiliary reward models by offering a single-stage, theoretically grounded optimization approach.
  • Empirical results show improved safety and helpfulness trade-offs especially in smaller models, with clear guidelines for hyperparameter tuning.

Safety Alignment as Density Ratio Matching: The BSO Framework

Motivation and Context

The problem of aligning LLMs simultaneously for helpfulness and safety is fundamentally a constrained optimization challenge. Maximizing helpfulness under a safety constraint typically demands complex architectures involving reward and cost models, online reinforcement learning (RL), and primal-dual optimization. These pipelines introduce significant computational overheads and instabilities. Recent direct preference optimization (DPO) approaches simplify learning for helpfulness, but safety is often incorporated through heuristic procedures without principled theoretical guarantees.

The central contribution of "BSO: Safety Alignment Is Density Ratio Matching" (2605.12339) is a unified, theoretically grounded approach to safety-aware preference optimization. The framework reinterprets safety alignment as a density ratio matching problem, leveraging Bregman divergence minimization to directly optimize for the optimal safe policy. This principled perspective both subsumes existing heuristic approaches as special cases and yields new algorithms with provable optimality guarantees.

Formalization of Safety Alignment

The paper begins by clarifying the entangled nature of helpfulness and safety constraints in LLM alignment. It highlights that preference datasets, collected for optimizing utility, often reward unsafe behavior when the preferred response is itself harmful. Consequently, solutions must formally separate optimizing helpfulness from enforcing safety.

The conventional solution models the problem as constrained optimization:

maxπ E(x,y)π[r(x,y)Cs(x,y)]βKL(π(x)πref(x)) s.t.s(x,y)0(x,y)\begin{align*} &\max_{\pi} \ \mathbb{E}_{(x,y)\sim \pi} [r(x, y) - C s(x, y)] - \beta \mathrm{KL}(\pi(\cdot | x) \| \pi_{\text{ref}}(\cdot | x)) \ &s.t. \quad s(x, y) \leq 0 \quad \forall (x, y) \end{align*}

where r(x,y)r(x, y) is the helpfulness reward, s(x,y)s(x, y) is the safety indicator, CC is the penalty coefficient, and πref\pi_{\text{ref}} is the reference policy.

Prior work such as SafeRLHF [Dai et al., 2023] operationalizes this schema via auxiliary reward and cost models, Lagrangian multipliers, and online RL. Direct preference optimization (DPO) [Rafailov et al., 2023] eliminates reward modeling and RL for helpfulness alignment, but its adaptations for safety are typically ad hoc.

Density Ratio Matching Principle

The crucial insight of BSO is that the log-likelihood ratio of the optimal safe policy admits a tractable decomposition involving the reference policy, helpfulness preference, and a safety correction term. This allows reframing safety alignment as a density ratio matching task.

The optimal safe policy's likelihood ratio is shown to satisfy:

πsafe(ywx)πsafe(ylx)=πref(ywx)πref(ylx)Phelp(yw<ylx)Phelp(yw>ylx)exp(C(swsl)/β)\frac{\pi_{\text{safe}}(y_w|x)}{\pi_{\text{safe}}(y_l|x)} = \frac{\pi_{\text{ref}}(y_w|x)}{\pi_{\text{ref}}(y_l|x)} \frac{P_{\text{help}}(y_w<y_l|x)}{P_{\text{help}}(y_w>y_l|x)} \exp(-C(s_w-s_l)/\beta)

where (x,yw,yl,sw,sl)(x, y_w, y_l, s_w, s_l) denotes the (prompt, winner, loser, winner safety, loser safety) tuple, PhelpP_{\text{help}} is the Bradley-Terry model over helpfulness, and CC is the safety penalty. This decomposition implies that enforcing the matching between model and data density ratios is sufficient to recover the optimal safe policy.

Bregman Safety Optimization (BSO)

BSO formalizes safety alignment as minimizing the Bregman divergence between the empirical data ratio and the model-induced ratio. Each choice of convex generator hh yields a tractable, theoretically guaranteed single-stage training loss:

r(x,y)r(x, y)0

where r(x,y)r(x, y)1 is the safety-aware ratio defined above. Under sufficient model capacity, minimizing this loss provably recovers the optimal safe policy.

This family of losses unifies prior safety-aware preference optimization methods, including SafeDPO [Kim et al., 2026], as special cases (e.g., logistic generator).

Practical Considerations and Algorithm Design

Implementing large safety penalties r(x,y)r(x, y)2 directly in the loss introduces numerical instabilities. The paper develops deterministic data transformations—label swapping and pair removal—that realize the qualitative effects of large r(x,y)r(x, y)3 without destabilizing gradient scaling. Concordant pairs (where both responses are equally safe) are weighted moderately, while safe-winner pairs receive amplified gradient contributions.

The generator analysis demonstrates the suitability of scaled Basu's power divergence (SBA), which enables controlled amplification of safe-winner pairs without disrupting the baseline gradient scale. The hyperparameters (safety penalty r(x,y)r(x, y)4 and amplification r(x,y)r(x, y)5) are shown through ablation to require moderate values for optimal joint performance.

Empirical Evaluation

Experiments are conducted on PKU-SafeRLHF-30K for Qwen2.5-0.5B and Llama3.2-3B backbones, against strong baselines (SFT, SafeRLHF, SACPO, SafeDPO):

  • BSO consistently improves both helpfulness and safety, shifting the Pareto frontier upward and right (higher helpfulness and harmlessness) compared to prior art. Gains are pronounced in smaller models that are more sensitive to safety weighting.
  • Ablations reveal that overly large safety penalties r(x,y)r(x, y)6 degrade reward without proportional safety improvements. Similarly, excessive amplification r(x,y)r(x, y)7 destabilizes optimization and harms the safety-helpfulness balance.
  • LLM-based evaluations corroborate the main results, confirming BSO's superiority in balancing safety and helpfulness.
  • The single-stage, auxiliary-model-free nature of BSO enables efficient training and more reliable convergence.

Theoretical and Practical Implications

BSO provides a unified theoretical framework for safety-aware preference optimization based on density ratio matching and Bregman divergence minimization. By subsuming prior methods under this principle, it enables systematic comparative analysis, principled algorithm design, and transparent hyperparameter tuning. The implications are multifold:

  • Algorithmic Simplicity: BSO eliminates the need for reward/cost models and multi-stage training, facilitating deployment and reproducibility.
  • Provable Optimality: The approach guarantees recovery of the optimal safe policy under stated assumptions and sufficient capacity.
  • Generalization Potential: Extensions to richer safety annotations, nuanced harmfulness taxonomies, and broader alignment settings are theoretically straightforward.
  • Trade-off Management: Practitioners gain access to a tunable safety-helpfulness trade-off, moving beyond heuristic margins.

Conclusion

"BSO: Safety Alignment Is Density Ratio Matching" fundamentally advances safety alignment in LLMs by casting the problem as density ratio matching, operationalized through Bregman Safety Optimization. This unified framework yields single-stage objectives with provable guarantees, improves empirical performance across benchmarks, and enables principled tuning of safety and useful behavior. Future work may extend these concepts to non-binary safety signals, interactive settings, and holistic alignment architectures. The BSO paradigm holds promise for simplifying and strengthening safety mechanisms in LLM training pipelines while maintaining practical utility.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.