Divergence Minimization in LLM Alignment

• Divergence minimization in LLM alignment is a principle that minimizes statistical differences between safe and harmful behavior distributions using f-divergence metrics.
• It unifies methods such as RLHF, DPO, and offline preference optimization to balance alignment, diversity, and robustness in language models.
• Empirical implementations like KLDO and AOT showcase practical loss formulations and metrics that quantify and enhance model alignment.

Divergence minimization is a foundational principle in modern LLM alignment. The central concept is to cast alignment as the problem of minimizing a statistical divergence between distributions representing "aligned" (safe, preferred) and "unaligned" (harmful, less preferred) behaviors. This approach unifies and generalizes the majority of notable LLM alignment techniques, from supervised fine-tuning and reinforcement learning from human feedback (RLHF) to more recent offline preference optimization frameworks. By explicitly or implicitly maximizing the separation between the distributions of desirable and undesirable behaviors, divergence minimization not only guides model training but also yields natural metrics for analyzing and quantifying alignment success.

1. Formal Problem Setup: Alignment as Distributional Divergence

LLM alignment is formalized via two key joint distributions over prompt-response pairs: $\mathcal{D}^+$ (aligned, e.g., safe or preferred) and $\mathcal{D}^-$ (unaligned, e.g., harmful or less preferred). Denote the corresponding densities $p^+(x,y)$ and $p^-(x,y)$. Conditioning on $x$ yields conditional distributions $p^+(y|x)$ and $p^-(y|x)$. The objective is to train a model policy $\pi_\theta(y|x)$ such that $\pi_\theta$ closely matches $p^+$ and diverges from $p^-$.

Popular methods, including RLHF, Direct Preference Optimization (DPO), KTO, and BCO, instantiate this by constructing a scalar reward $r_\theta(x,y)$—typically the log-likelihood difference relative to a base policy—and optimizing to distinguish $\mathcal{D}^+$ from $\mathcal{D}^-$. These approaches can be interpreted as optimizing lower bounds or surrogates on standard $f$-divergences such as total variation, Jensen-Shannon, and Kullback-Leibler divergences (Haldar et al., 2 Feb 2025, Go et al., 2023, Han et al., 2024).

2. $f$-Divergence Minimization: Unified Theoretical Framework

Let $f:(0,\infty)\to\mathbb{R}$ be convex with $f(1)=0$. The $f$-divergence between distributions $P$ and $Q$ is

$D_f(P \,\|\, Q) = \sum_x Q(x)f(P(x)/Q(x)).$

This formulation encompasses forward KL ($f(t) = -\log t$), reverse KL ($f(t) = t\log t$), Jensen-Shannon, total variation, $\alpha$-divergences, squared Hellinger, and others (Go et al., 2023, Han et al., 2024).

Alignment objectives are typically structured as minimizing some $D_f(\pi_\theta \,\|\, \pi^*)$ or $D_f(\pi^* \,\|\, \pi_\theta)$, where $\pi^*$ is an explicit or implicit target policy derived from human feedback, demonstration data, or reward models. The gradient of $D_f(\pi_\theta \,\|\, p)$ admits a general policy gradient estimator: $$\nabla_\theta D_f(\pi_\theta \,\|\, p) = \mathbb{E}_{x \sim \pi_\theta} \left[ f'( \pi_\theta(x)/p(x) ) \cdot \nabla_\theta \log \pi_\theta(x) \right].$$ Different choices of $f$ correspond to different trade-offs between alignment, diversity, robustness, and mode collapse (Go et al., 2023, Han et al., 2024).

Offline preference optimization methods such as $f$-PO generalize these ideas and demonstrate that DPO, EXO, and related algorithms are special cases where $f$ is respectively $-\log$, $u\log u$, etc. (Han et al., 2024). Empirically, $\alpha$-divergence interpolates between forward and reverse KL, allowing practitioners to control mode coverage and specialization.

3. Instantiations: KLDO, AOT, Distributional Alignment

KLDO (KL-Divergence Optimizer) is a direct application of the Donsker–Varadhan variational representation of KL divergence: $$L_{\mathrm{KLDO}}(\theta) = -\, \mathbb{E}_{\mathcal{D}^+}[ r_\theta(x, y) ] + \log \mathbb{E}_{\mathcal{D}^-}[ e^{ r_\theta(x, y) } ],$$ which at its optimum yields $$-L_{\mathrm{KLDO}}(\theta^*) = \mathrm{KL}( \mathcal{D}^+ \| \mathcal{D}^- )$$ (Haldar et al., 2 Feb 2025). This direct optimization of the KL divergence shows superior separation over TV and JS-based approaches in the presence of well-separated distributions.

Alignment via Optimal Transport (AOT) frames the task as enforcing first-order stochastic dominance (FSD) of the reward distributions from chosen over rejected responses. FSD constraints are relaxed to a one-dimensional optimal transport problem with a convex cost function $h$: $$L_{\mathrm{AOT}}(\theta) = \frac{1}{n} \sum_{i=1}^n h( u_\theta^{(i)} - v_\theta^{(i)} ),$$ where $u_\theta^{(i)}$ and $v_\theta^{(i)}$ are the sorted reward scores of positive and negative samples. AOT produces consistent empirical gains and admits closed-form (sorting-based) solutions for the OT coupling (Melnyk et al., 2024).

Distributional alignment methods for LLM-as-a-Judge align the model’s output distribution with the human empirical distribution using a KL objective augmented by cross-entropy regularization and adversarial perturbations to improve robustness: $$L_{\mathrm{KL}}(\theta) = \frac{1}{|D|} \sum_{x \in D} D_{\mathrm{KL}}( p(x) \| q_\theta(x) )$$ with adversarial minimax extensions to guard against overfitting when sample sizes are limited (Chen et al., 18 May 2025).

4. Data Regimes and Their Impact on Divergence Separation

The choice of data regime has pronounced effects on divergence-based alignment. Compliance–refusal (CR) datasets create maximally separated clusters in latent space, as every prompt–response pair is forced to have $R(x, y) \ne 1$ (where $R$ is the likelihood ratio $p^+(y|x)/p^-(y|x)$). This separation enables a Bayes classifier on prompt embeddings or reward scores to perfectly recover a prompt’s latent safety label $z_x$. For standard preference data, in contrast, the separation is weakened since some pairs correspond to two compliant (both aligned) responses, reducing the information available for divergence maximization. Empirically, CR data leads to larger divergence values, better embedding separation, and lower adversarial success rates in jailbreaking evaluations (Haldar et al., 2 Feb 2025).

In contrast, demonstration-based methods (AfD/Inverse-RLignment) alternate between mass-covering (forward KL, SFT-like) and mode-seeking (reverse KL, discriminator-based) divergence objectives depending on the diversity of the expert data (Sun et al., 2024).

5. Comparison of Divergence Families and Their Effects

Selection of the divergence family exerts a strong influence on the alignment-diversity-robustness trade-off:

• Reverse KL ($\mathrm{KL}(\pi_\theta \| p^*)$): Highly mode-seeking; produces sharp alignment but risks mode collapse. Basis for RLHF and DPO.
• Forward KL ($\mathrm{KL}(p^* \| \pi_\theta)$): Mass-covering; expands response diversity but may dilute alignment and induce unstable gradients. Basis for SFT and GDC.
• Jensen–Shannon: Yields more stable policy updates, balancing alignment and diversity, and often forms the Pareto frontier (Go et al., 2023, Han et al., 2024).
• $\alpha$-divergence: Interpolates smoothly between forward and reverse KL, allowing tuning for domain requirements (Han et al., 2024).
• TV, Hellinger, Jeffrey’s, and others: Offer additional options with distinct convergence and generalization characteristics.

In ablation studies, Jensen–Shannon and $\alpha$-divergence consistently outperform solely forward or reverse KL in producing simultaneously high win rates and reasonable entropy across preference datasets. Larger models magnify divergence-induced distinctions, with orderings by alignment and diversity persisting despite scaling (Go et al., 2023, Han et al., 2024).

6. Quantifying Alignment: Statistical and Representation-Based Metrics

Divergence-minimization not only aligns policy distributions but also structurally separates aligned and unaligned behaviors in latent space. A canonical method is to map each prompt $x$ to a representation $h(x)$ in the model’s last hidden layer, fit Gaussians to clusters of safe and unsafe embeddings, and compute the Bhattacharyya distance: $$D_B = \frac{1}{8} (\mu_s - \mu_u)^\top \Sigma^{-1} (\mu_s - \mu_u) + \frac{1}{2} \log \frac{|\Sigma|}{ \sqrt{ |\Sigma_s| \cdot |\Sigma_u| } },$$ with $\mu_s, \mu_u$ (means) and $\Sigma_s, \Sigma_u$ (covariances) for safe/unsafe respectively, and $\Sigma = (\Sigma_s + \Sigma_u )/2$. Larger $D_B$ correlates with lower adversarial (ASR) rates, demonstrating its utility as a safety indicator (Haldar et al., 2 Feb 2025). Other clustering metrics, such as the Silhouette Score, are also used but are less directly linked to the divergence-maximization principle.

Distributional alignment methods further quantify performance through KL to empirical human label distributions, top-1 prediction accuracy, and robustness to perturbations of the label distribution. Empirical studies show that removing any divergence-based term materially degrades alignment quality (Chen et al., 18 May 2025).

7. Practical Implementation and Algorithmic Considerations

Training recipes for divergence-based alignment instantiate SGD or Adam/LoRA updates over the chosen divergence surrogate, with careful construction of positive and negative batches, often involving adversarial sampling strategies or buffer approximations for variance reduction. For instance, KLDO relies on moving averages or negative sample buffers to stabilize estimation of $\mathbb{E}_{\mathcal{D}^-}[e^r]$ (Haldar et al., 2 Feb 2025). AOT leverages deterministic sorting of reward scores thanks to its one-dimensional OT structure, which permits closed-form loss computation (Melnyk et al., 2024).

Adversarial and hybrid objectives (e.g., KL plus cross-entropy with adversarial perturbations) have proven necessary for improving robustness to annotation noise and finite-sample effects (Chen et al., 18 May 2025). Dual-formulation analyses enable sample complexity guarantees, with AOT achieving parametric $O(1/\sqrt{n})$ generalization rates under natural assumptions (Melnyk et al., 2024).

An overview of example loss formulations:

Method	Divergence	Core Loss Expression
DPO/RLHF	Reverse KL/TV	$$-\mathbb{E}_{\text{paired}} \log \sigma [ r_\theta(y^+) - r_\theta(y^-) ]$$
KLDO	KL	$$-\mathbb{E}_{\mathcal{D}^+} [r_\theta] + \log \mathbb{E}_{\mathcal{D}^-} [e^{r_\theta}]$$
AOT	OT/FSD	$$\frac{1}{n} \sum_{i=1}^n h(u_\theta^{(i)}-v_\theta^{(i)})$$
f-PO	General $f$	$$\mathbb{E}_{(x,y^+,y^-)} [ (1-\epsilon)f(u_1)+\epsilon f(u_2) ]$$
LLM-as-J	KL (+CE, Adv)	$$\mathrm{KL}(p(x)\Vert q_\theta(x)) + (1-\alpha)\mathrm{CE}(q_\theta,r(x))$$

The practitioner’s choice of divergence, surrogate, and data regime determines the alignment-diversity trade-off, sample efficiency, and robustness properties.

---

Collectively, these results establish divergence minimization as a rigorous, principled foundation for virtually the entire spectrum of LLM alignment methods, from RLHF to modern offline and distributional preference optimization. The efficacy, robustness, and quantifiability of alignment are directly rooted in the mathematical properties of divergences and their optimization. The science and engineering of LLM alignment are now predominantly cast in these distributional and statistical terms.