Aligning Language Models with Preferences through f-divergence Minimization

• Alignment can be viewed as approximating a desired target distribution, while GDC uses explicit energy-based targets and forward KL and RLHF corresponds to reverse KL minimization.This distinction separates describing aligned behavior from approximating it.Source: §1 · pp. 1–2
• Divergence choice matters under model misspecification because different objectives can yield different optima and optimization landscapes.Prior work largely coupled target-distribution form with the approximation algorithm.Source: §1 · p. 2
• f-DPG unifies DPG and RL with KL penalties by approximating any evaluable target energy-based model with any f-divergence, including KL, reverse KL, TV, and JS.The framework provides a universal gradient-descent formula, variance-reduction baselines, and support for conditional target distributions.Source: §1 · p. 2
• The study evaluates KL-DPG, RKL-DPG, TV-DPG, and JS-DPG across thirteen language-model alignment tasks, three target-distribution forms, four objectives, and eight metrics.Tasks include lexical control, social-bias reduction, factual summarization, and code compilability.Source: §1 · p. 2
• JS-DPG often balances alignment and diversity best, significantly improves over prior work and KL-DPG, and remains superior to suboptimal objectives from 127M to 1.5B parameters.Other divergences can outperform KL-DPG even when quality is evaluated using forward KL, while no objective is universally best.Source: Figure 1 · p. 1, §1 · p. 2

2. Background

• Alignment approaches differ in how they construct a target distribution and how they approximate it to produce samples.The first axis represents human preferences as probabilities over texts; the second enables sampling from that distribution.Source: §2 · p. 2

2.1. Defining a Target Distribution

• A target distribution assigns probabilities to texts according to preference satisfaction, often via a nonnegative energy-based function P(x), but it can score rather than directly generate samples.The paper focuses on targets used in GDC and RLHF (Ziegler et al., 2019; Stiennon et al., 2020; Ouyang et al., 2022; Menick et al., 2022; Bai et al., 2022a).Source: §2.1 · p. 2
• The paper covers binary constraints, scalar rewards, and distributional feature-moment constraints as three forms of preference targets.Binary targets minimally deviate from a pretrained LM under KL(pGDC bin||a), RLHF uses reward maximization with KL penalties and equivalently reverse-KL minimization toward pRLKL, and distributional targets match desired feature moments while minimally deviating under KL(pGDC dist||a).Source: §2.1 · p. 3

2.2. Approximating the target distribution

• Approximating a target distribution can use inference-time decoding or sampling, or train a parametric model πθ; this paper focuses on minimizing D(πθ||p), including DPG’s KL(p||πθ) objective.Inference-time methods include guided decoding, rejection sampling, QRS, and MCMC, while Khalifa et al. (2021) uses Distributional Policy Gradients (DPG; Parshakova et al., 2019) to minimize KL(p||πθ), equivalently CE(p, πθ).Source: §2.2 · p. 3

3.1. f-divergences

• The f-divergence family includes forward KL, reverse KL, Jensen-Shannon divergence, and Total Variation, with each measure specified by a convex generator f.Under the stated conditions, f-divergences are nonnegative and equal zero when the distributions match; strict convexity at 1 makes equality diagnostic.Source: §3.1 · pp. 3–4

3.2. Distributional alignment with f-divergences

• f-DPG minimizes the f-divergence between a target distribution p and a parametrized generative model πθ, using samples from πθ when direct sampling from p is unavailable.The model can both sample texts and evaluate their probabilities, enabling the optimization framework.Source: §3.2 · p. 4
• When p belongs to the model family, every chosen divergence has minimum zero; under model misspecification, the divergence determines which approximating distribution is preferred.This dependence motivates treating divergence choice as part of the alignment method rather than merely an optimization detail.Source: §3.2 · p. 4
• Theorem 1 expresses the f-divergence gradient as a policy-gradient-style expectation under πθ, defining a parameter-dependent pseudo-reward and the f-DPG algorithm.The theorem applies when either the target support is contained in every model support or the model support is independent of θ; standard softmax language models have full support.Source: §3.2 · p. 4

3.3. Adding a baseline

• Subtracting a constant baseline from the pseudo-reward leaves f-DPG gradient estimates unbiased, and experiments use an estimate of the average pseudo-reward unless specified otherwise.The baseline is justified by the zero expectation of the score-function gradient and is typically chosen as the current average reward.Source: §3.3 · p. 4

3.4. Recovering Some Existing Methods

• GDC is recovered by choosing forward KL, KL(p||πθ), while RL with KL penalties is recovered by choosing reverse KL, KL(πθ||p), within the f-DPG framework.The forward-KL case corresponds to f(t) = −log t, whereas reverse KL corresponds to f(t) = t log t; the latter recovers the usual KL-penalty gradient up to a constant factor β.Source: §3.4 · p. 5

3.5. Estimating Z

• When the target is proportional to a nonnegative function P(x), f-DPG can estimate its normalizing constant Z by importance sampling from current and previously sampled model distributions.Each importance-sampling estimate is unbiased, while averaging estimates across different πθ can improve precision.Source: §3.5 · p. 5

3.6. Conditional Target Distributions

• For conditional tasks, f-DPG generalizes to conditional models πθ(x|c) by combining a context distribution τ(c) with a context-specific target distribution pc.The conditional formulation extends the DPG approach to machine translation, summarization, and dialogue.Source: §3.6 · p. 5

4. Experiments

• The experiments compare forward KL, reverse KL, Total Variation, and Jensen-Shannon f-DPG using divergence, alignment, language-model deviation, entropy, and pseudo-reward-variance metrics across diverse tasks.The study uses GPT-2 small with 117M parameters initially, with implementation details and hyperparameters provided separately.Source: §4 · pp. 5–6

4.1. Alignment with Scalar Preferences

• JS-DPG performs best overall among KL-DPG, RKL-DPG, TV-DPG, and JS-DPG on the scalar sentiment-preference task, consistently minimizing the tracked divergences.RKL-DPG achieves the strongest reverse-KL and alignment scores but sacrifices diversity and departs substantially from the original language model.Source: §4.1 · p. 6, Figure 2 · p. 6

4.2. Alignment with Lexical Constraints

• On lexical constraints, all f-DPG variants reduce target divergence and improve desired-word production, while every alternative outperforms KL-DPG even on forward KL.KL-DPG’s comparatively poor performance is hypothesized to result from high pseudo-reward variance; RKL-DPG also produces lower normalized entropy without significantly reducing sentence-level diversity.Source: §4.2 · pp. 6–7, Figure 3 · p. 7

4.3. Alignment with Distributional Constraints

• Distributional constraints reduce regard-score bias from 1:0.677 to 1:0.801, while RKL-DPG best matches the gender pointwise constraint but appears less diverse.TV-DPG converges suboptimally on regard rebalancing because hard-thresholded pseudo-rewards have high variance when the original distribution is already near the target.Source: §4.3 · p. 7, Figure 4 · p. 8

4.4. Alignment with Conditional Constraints

• For factual summarization, f-DPG increases consistent named entities and indirectly improves Rouge quality without training on ground-truth summaries, while JS-DPG converges better to the target than KL-DPG.The conditional constraint requires generated named entities to be contained in the source and include at least four entities.Source: §4.4 · pp. 7–8, Figure 5 · p. 8

4.5. Scaling Trends of f-DPG

• As model size grows from GPT-2 small (117M) to xl (1.5B), alignment steadily improves, but differences among divergence objectives and their general ordering persist.The smooth scaling trend suggests that increasing model size alone does not close the gap between optimal and suboptimal divergence objectives.Source: §4.5 · p. 8, Figure 6 · p. 9

4.6. Ablation Study

• With matched-capacity target and policy families, KL-DPG still lags in divergence and has high pseudo-reward variance, whereas RKL-DPG’s initial entropy loss can recover.Baseline techniques improve lexical-constraint training, especially for RKL-DPG; large batches reduce KL-DPG variance but do not eliminate its inferior performance, and importance-sampling estimates converge to the true normalization value.Source: §4.6 · pp. 8–9

5. Discussion and Conclusion

• JS-DPG consistently lies on the alignment–diversity Pareto frontier, balancing RKL-DPG’s high-alignment, low-diversity behavior against KL-DPG’s low-alignment, high-diversity behavior.These findings hold across the reported alignment tasks and experiments.Source: Figure 7 · p. 9, §5 · p. 9
• Pseudo-reward shapes explain the trade-offs: forward KL emphasizes coverage, reverse KL can induce collapse, Total Variation is outlier-robust but high-variance near matching distributions, and JS is smooth and robust.The paper identifies JS as a reasonable default because it avoids relying heavily on a single direction.Source: §5 · p. 9, Figure 8 · p. 9
• The proposed f-DPG framework approximates evaluable target distributions with any f-divergence, unifies earlier alignment approaches, and shows that objective choice remains important as model size increases.Larger models improve alignment scores but do not inherently eliminate gaps between divergence objectives.Source: §5 · p. 9

A.1. Equivalent definitions for f-divergences

• The symmetrical f-divergence definition uses a convex generator with f(1)=0, incorporates boundary masses through f(0) and f*(0), and satisfies D_f(p,q)=D_f*(q,p).The perspective transform is f*(t)=t f(1/t), with f**=f.Source: §A.1 · p. 16
• The appendix establishes equivalence between the paper’s two f-divergence definitions using f'(∞)=f*(0) and conventions for zero-mass boundary terms.The conventions include 0f(0)=0 and 0f*(0)=0, including infinite endpoint values.Source: §A.1 · p. 16

A.2. Illustrations of a few f-divergences

• The formal illustrations recover forward and reverse KL through the generator and its perspective transform, while Total Variation is symmetric under swapping its arguments.For Total Variation, both orientations can be represented as the corresponding f-divergence.Source: §A.2 · p. 17

A.4. About non-differentiability of f

• Although Total Variation’s generator is nondifferentiable, subgradients are typically unnecessary because nondifferentiability occurs on a measure-zero set, making Eq. 8 well-defined almost surely over θ.The relevant exceptional set is countable within the parameter space.Source: §A.4 · p. 18

A.3. Proof of Theorem 1

• Theorem 1 applies when either Supp(p) ⊂ Supp(πθ) for every θ or Supp(πθ) is independent of θ; under either condition, ∇θp(πθ = 0) = ∇θ0 = 0.Source: §A.3 · pp. 17–18
• The proof handles zero-probability terms using f′(∞) .= f*(0), a multivariate chain-rule variant, and an equivalent notation when p(x) = 0 and πθ(x) > 0.The identity limt→∞ f(t)/t = f*(0) motivates the extended-derivative notation.Source: §A.3 · p. 18
• The same gradient result applies to the opposite-order divergence through Df(p||πθ) = Df*(πθ||p), with the derivative taken on f’s perspective transform.The support assumptions remain unchanged under this reformulation.Source: §A.3 · p. 18

A.6. Baseline: alternative derivation

• Adding c(t − 1) to an f-divergence generator leaves the divergence unchanged, enabling baseline construction by replacing f with an equivalent generator g.For Df(πθ||p), choosing g(x) = f(x) − B(x − 1) preserves both the divergence and its gradient while introducing a baseline B.Source: §A.6 · pp. 18–19
• The alternative baseline derivation is paired with an iterative estimator that samples from πθ, tracks the normalization constant Z with a moving average, and updates πθ using Theorem 1.The procedure starts from an unnormalized target P(·) and outputs the updated model πθ.Source: §A.6 · p. 19

B. Extended Related Work

• Prior RL-for-language-model work includes policy-gradient, actor-critic, and heuristic- or task-reward methods, while KL-control penalizes deviation from the original policy and is optimized with algorithms such as PPO or Q-learning.This line of work has supported aligned language models including LaMDA, InstructGPT, Sparrow, and CAI.Source: B. Extended Related Work · p. 19
• Related f-divergence research spans GANs, variational losses, divergence-specific convergence analysis, forward-KL limitations, and RL or imitation-learning formulations connecting divergence minimization with trajectory distributions.These studies include Goodfellow et al. (2020), Nowozin et al. (2016), Arjovsky et al. (2017), Theis et al. (2016), and Ke et al. (2021); Ghasemipour et al. (2020).Source: B. Extended Related Work · pp. 19–20
• Compared with Wang et al. (2018), Theorem 1 is more general because Df(πθ||p) requires no absolute continuity in either direction, covering cases where p(x) = 0 and πθ(x) > 0.The authors identify this case as crucial for language-model alignment.Source: B. Extended Related Work · p. 20

C. Implementation Details

• Models used PyTorch, HuggingFace Transformers, and Adam on Nvidia V100 GPUs; hyperparameters were searched primarily for KL-DPG, then tuned once for all f-DPG methods under comparable settings.The longest training run took approximately 2 days, and hyperparameter details were reported in Tab. 2.Source: C. Implementation Details · p. 20

D.1. Generation Quality

• The reported generation-quality metrics are entropy, Self-BLEU-5, Distinct-1, and perplexity, respectively measuring diversity, distributional repetition, within-sample diversity, and fluency.Perplexity is computed with a separate Distil-GPT-2 model to avoid inflated estimates.Source: Table 3 · p. 21, §D.1 · p. 21
• RKL-DPG has notably lower entropy and higher Self-BLEU than other f-DPGs, indicating lower distributional diversity, while achieving better perplexity-based individual-sample quality.This suggests RKL-DPG captures a subset of the target distribution.Source: §D.1 · p. 21, Table 3 · p. 21
• For lexical-constraint generations, the authors found no significant difference in diversity among generated sentences.These results are aggregated in Table 4 using reward, Self-BLEU-5, Distinct-1, and perplexity metrics.Source: §D.1 · p. 21, Table 4 · p. 21

E. f-DPG on Conditional Target Distributions

• Conditional control seeks a model p_c(x) that satisfies a binary constraint for every sample while staying minimally KL-divergent from the pretrained conditional model a(x|c).The target has the form p_c(x) = 1/Z_c a(x|c)b(x,c).Source: E. f-DPG on Conditional Target… · p. 21
• The conditional target P is difficult to sample because it represents potentially infinitely many unconditional models, each lacking an autoregressive factorization.Korbak et al. (2022a) address this by fitting a generative model π_θ through expected KL(p_c||π_θ), equivalently expected cross-entropy.Source: E. f-DPG on Conditional Target… · pp. 21–22
• The conditional formulation is presented as an extension of the unconditional f-DPG objective, with gradients defined for optimization across contexts.The gradient is described as a conditional extension of Eq. 5.Source: E. f-DPG on Conditional Target… · p. 22
• f-DPG extends the objective to expected f-divergence D_f(π_θ(·|c)||p_c(·)) over contexts, including expected KL(p_c||π_θ(·|c)).The extension uses the gradient implied by Theorem 1.Source: E. f-DPG on Conditional Target… · p. 22

E.1. Additional Conditional Preferences Experiments and Details

• Summarization quality is evaluated with precision-source, recall-target, and Rouge, covering entity faithfulness, target-entity recovery, and unigram overlap.Precision-source is defined over named entities in the source, recall-target over named entities in the target summary, and Rouge over source–summary unigram overlap.Source: §E.1 · p. 22
• For code generation, evaluation uses average PEP8 error count and the fraction of concatenated prompt–function samples that compile.The Python experiment uses GPT-Neo-125 and marks a function as uncompilable when the Python interpreter raises an exception.Source: §E.1 · p. 22
• JS-DPG increases the fraction of compilable Python functions while reducing average PEP8 violations, converging to p better than KL-DPG.The code-generation experiment conditions on Python function signatures and evaluates compilability and PEP8 error count.Source: §E.1 · p. 22, Figure 10 · p. 23

F. Optimal Reward Model for a Decision Maker with a Categorical Distribution

• Under categorical choice feedback, the cross-entropy-optimal reward model reproduces the decision-maker’s choice probabilities, and the construction rϕ is optimal when the internal scoring function ϕ is known.The choice function returns a one-hot preferred-sample vector, while the reward model is trained through a parametrized discrete choice model.Source: F. Optimal Reward Model for a Decision… · pp. 23–24

G. Additional Figures

• Additional figures evaluate sentiment-preference metrics, constrained single-word distributions, IMDB alignment and diversity, and alignment scores across model sizes.The reported metrics include f-divergences, expected internal scores, KL to the initial model, diversity and entropy measures, pseudo-reward summaries, and model-size comparisons.Source: Figure 11 · p. 24, Figure 12 · p. 25, Figure 13 · p. 26
• The constrained-generation evaluations note that reverse KL is undefined when the target distribution assigns zero probability to some points.This qualification applies to the target distribution induced by the GDC framework in the single-word constraint experiments.Source: Figure 12 · p. 25

H.1. Matching Other Language Model within Parameter Family

• The experiment reverses fine-tuning by training a 117M-parameter GPT-2 model fine-tuned on IMDB to match the original GPT-2 architecture and distribution.The parameter family is intentionally correctly specified to isolate loss-function behavior from model-capacity limitations.Source: §H.1 · p. 28
• RKL-DPG initially loses substantial entropy but ultimately converges to p(x), because the correctly specified family allows recovery and coverage of the remaining distribution.The initial entropy drop is attributed to mode-seeking behavior.Source: §H.1 · p. 28
• The resulting models recover much of the original GPT-2’s quality when evaluated zero-shot on summarization.The summarization evaluation follows Radford et al. (2019).Source: §H.1 · p. 28
• TKL-DPG, TV-DPG, and JS-DPG rapidly approach GPT-2, whereas KL-DPG progresses slowly and regresses in TV distance, consistent with high-variance pseudo-reward gradients.These metrics were averaged over three independent seeds; Fig. 17 reports divergences, diversity, perplexity, entropy, and pseudo-reward statistics.Source: Figure 17 · p. 28, §H.1 · p. 28

Motivation

• The experiment tests whether matching GPT-2’s distribution transfers its general fluency and multitask capabilities beyond unconditional generation.This relies on GPT-2’s previously demonstrated ingrained multi-task capabilities.Source: §H.2 · p. 29

Evaluation setup

• The evaluation generates 100 tokens after “TL;DR:” with top-k sampling (k = 2), using the first three sentences and average ROUGE-1,2,L for comparison.Ground-truth summaries are used only to compute evaluation metrics, not during training or sampling.Source: §H.2 · p. 29

Results

• f-DPG successfully recovers the initial model’s lost summarization ability while matching GPT-2’s distribution without using CNN/Daily Mail data or ROUGE during training.The result is shown in Figure 18 through the evolution of average ROUGE-1,2,L scores across training epochs.Source: §H.2 · p. 29, Figure 18 · p. 29

Effect of baseline

• Baseline techniques improve every f-DPG method, with RKL-DPG benefiting most because averaging mitigates its large negative pseudo-rewards.The ablation uses a scalar preference rewarding outputs containing “amazing.”Source: §H.3 · p. 29, Figure 19 · p. 30

Effect of batchsize

• Large batches are necessary to address KL-DPG’s high variance, while increasing batch size does not change the study’s conclusions.This supports improved sample efficiency and performance when applying f-DPG within the GDC framework.Source: §H.3 · p. 29, Figure 20 · p. 31

Z estimation

• Partition-function estimates converge to the true value across f-DPG models, with no significant learned-model difference from using the true value.Estimating Z is needed when the unnormalized target distribution’s partition function is unknown; for the binary-feature constraint, bootstrap estimation uses the initial Bernoulli frequency.Source: §H.3 · p. 31

I. Samples

• The Z-estimation ablation tracks the estimated Z against its true value over time.Source: Figure 21 · p. 32
• Convergence curves are compared across f-DPG models using the true Z value.Source: §I · p. 32
• The comparison includes KL-DPG, TV-DPG, and JS-DPG variants.Source: §I · p. 33

RKL-DPG

• Generation samples compare KL-DPG, TV-DPG, and JS-DPG for sentiment preference, with the corresponding examples presented in Table 5.The passages provide sample outputs for each method but no quantitative comparison or stated winner.Source: RKL-DPG · p. 34, Table 5
• For the “amazing” preference, the section again shows KL-DPG, TV-DPG, and JS-DPG outputs, summarized visually in Table 6.The supplied samples illustrate generated text, without reporting a metric or qualitative ranking.Source: Table 6 · p. 34, RKL-DPG · pp. 34–35
• Table 7 presents generation samples for a female 50% science 100% preference across KL-DPG, TV-DPG, and JS-DPG.The passages contain method-specific sample outputs but do not state which method performs best.Source: Table 7 · p. 35, RKL-DPG · p. 35
• Tables 8–10 present generation samples for summarization.Source: Table 8 · p. 36, Table 9 · p. 37, Table 10 · p. 38