Direct Preference Optimization LLM Alignment

# Direct Preference Optimization & LLM Alignment

## Introduction & Motivation

Aligning large language models with human preferences, so that they produce outputs judged helpful, honest, and harmless rather than merely fluent, has historically been accomplished through Reinforcement Learning from Human Feedback (RLHF): a pipeline that trains a separate reward model on human preference comparisons, then optimizes the language model against that reward model using an RL algorithm such as Proximal Policy Optimization (PPO). This pipeline works well in practice but is notoriously complex to implement correctly, requires training and hosting multiple large models simultaneously (the policy, the reward model, a reference model, and often a separate value function), is sensitive to a wide range of RL-specific hyperparameters, and can suffer from instability during training.

Direct Preference Optimization (DPO), introduced by Rafailov and colleagues in 2023, showed that the entire RLHF pipeline could be replaced by a single, simple classification-style loss computed directly on preference data, with no reward model, no RL rollouts, and no value function, while provably optimizing the same objective that the RLHF pipeline targets. This result triggered a wave of follow-on preference-optimization methods, including Identity Preference Optimization (IPO), Kahneman-Tversky Optimization (KTO), and Odds Ratio Preference Optimization (ORPO), each adjusting the underlying loss to address specific weaknesses of the original DPO formulation, collectively forming what is now often called the family of direct alignment algorithms.

The practical significance of this shift is substantial: direct alignment methods dramatically simplify the engineering effort required to fine-tune a language model on human or AI-generated preference data, reducing the infrastructure from a four-model RL pipeline to a supervised-learning-style loop resembling ordinary fine-tuning, which has made preference alignment accessible to teams without dedicated RL infrastructure and has become the default alignment technique for the vast majority of openly released instruction-tuned language models.

## Core Concepts & Theory

The starting point for DPO is the same as for RLHF: a dataset of preference comparisons, each consisting of a prompt and a pair of candidate responses, one marked as preferred (the "chosen" response) and one marked as dispreferred (the "rejected" response), typically collected by presenting both responses to a human annotator or, increasingly, to a stronger language model acting as an automated judge. The RLHF pipeline uses this data to train a reward model that assigns a scalar score to any prompt-response pair, then uses RL to find a policy that maximizes expected reward while staying close, as measured by KL divergence, to a fixed reference policy (typically the supervised-fine-tuned model prior to preference training).

The key theoretical insight behind DPO is that, for the specific KL-regularized reward-maximization objective used in RLHF, the optimal policy has a closed-form relationship to the reward function: the reward function can be expressed entirely in terms of the optimal policy and the reference policy, meaning that a reward model need never be explicitly instantiated at all. Substituting this closed-form relationship into the standard Bradley-Terry preference model, which expresses the probability that one response is preferred over another as a function of the difference in their underlying reward, yields a loss function expressed purely in terms of the policy being trained and a fixed reference policy, with the reward model canceled out algebraically.

This substitution transforms preference alignment from a two-stage RL problem into a single supervised classification problem: given a chosen and rejected response for a given prompt, the DPO loss increases the policy's relative log-probability of the chosen response compared to the rejected response, relative to how the reference policy scores the same pair, with the overall strength of this push controlled by a temperature-like parameter beta that plays the same role as the KL penalty coefficient in RLHF.

## Mathematical Formulation

The RLHF objective that both PPO-based RLHF and DPO ultimately optimize is a KL-regularized reward maximization problem, where pi_theta is the policy being trained, pi_ref is the fixed reference policy, r(x, y) is the reward function, and beta controls the strength of the penalty keeping the trained policy close to the reference:

$$ \max_{\pi_ heta} \; \mathbb{E}_{x \sim \mathcal{D}, \, y \sim \pi_ heta(y \mid x)} \left[ r(x, y) ight] - \beta \, \mathbb{D}_{\mathrm{KL}}\!\left[ \pi_ heta(y \mid x) \, \| \, \pi_{\mathrm{ref}}(y \mid x) ight] $$

The closed-form optimal policy for this objective, for a given reward function r, can be shown to take the following form, where Z(x) is an intractable partition function that normalizes the distribution over all possible responses y:

$$ \pi^*(y \mid x) = \frac{1}{Z(x)} \, \pi_{\mathrm{ref}}(y \mid x) \, \exp\!\left(\frac{1}{\beta} r(x, y) ight) $$

Rearranging this expression algebraically solves for the reward function in terms of the optimal policy, the reference policy, and the (intractable) partition function:

$$ r(x, y) = \beta \log \frac{\pi^*(y \mid x)}{\pi_{\mathrm{ref}}(y \mid x)} + \beta \log Z(x) $$

Substituting this reward expression into the Bradley-Terry preference model, in which the probability that response y_w (the chosen or "winning" response) is preferred over y_l (the rejected or "losing" response) is a sigmoid of the reward difference, causes the intractable partition function Z(x) to cancel out entirely, since it depends only on x and appears identically in both reward terms, yielding the final DPO loss as a simple binary cross-entropy over policy log-probability ratios:

$$ \mathcal{L}_{\mathrm{DPO}}(\pi_ heta; \pi_{\mathrm{ref}}) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \log \sigma\!\left( \beta \log \frac{\pi_ heta(y_w \mid x)}{\pi_{\mathrm{ref}}(y_w \mid x)} - \beta \log \frac{\pi_ heta(y_l \mid x)}{\pi_{\mathrm{ref}}(y_l \mid x)} ight) ight] $$

The implicit reward that DPO training induces for any response, without ever training an explicit reward model, is simply the beta-scaled log-ratio of the trained policy's probability to the reference policy's probability for that response:

$$ \hat{r}_ heta(x, y) = \beta \log \frac{\pi_ heta(y \mid x)}{\pi_{\mathrm{ref}}(y \mid x)} $$

## Advanced Theory & Extensions

Identity Preference Optimization (IPO) addresses a specific weakness of DPO: because the DPO loss is a sigmoid of an unbounded log-ratio difference, and because real preference datasets are rarely perfectly separable (the same prompt may sometimes receive contradictory preference labels, or preferences may be genuinely close to a coin flip), DPO's loss can be driven arbitrarily close to zero by pushing the log-ratio difference to positive infinity, causing the policy to overfit the preference data and drift arbitrarily far from the reference policy in extreme cases. IPO replaces the sigmoid-based loss with a bounded squared-error objective on the log-ratio difference, directly regularizing how far the implicit reward gap is allowed to grow, which empirically produces more stable training on noisy or non-separable preference data.

Kahneman-Tversky Optimization (KTO) removes the requirement for paired preference data (a chosen and rejected response for the same prompt) entirely, instead requiring only a binary desirable/undesirable label on individual (prompt, response) examples, which is a substantially weaker and cheaper form of supervision to collect since it does not require presenting two responses side by side. KTO's loss is derived from prospect theory, the behavioral economics framework describing how humans evaluate gains and losses asymmetrically relative to a reference point, and empirically performs comparably to DPO despite using less structured feedback, making it attractive when only single-response quality judgments (such as thumbs up/thumbs down feedback from production traffic) are available.

Odds Ratio Preference Optimization (ORPO) takes a different approach, eliminating the reference model altogether by combining a standard supervised fine-tuning loss on the chosen response with an odds-ratio-based penalty term that suppresses the probability of the rejected response, allowing alignment and instruction-tuning to be performed in a single training stage without needing to separately produce and store a frozen reference model's log-probabilities. Self-Play Preference Optimization and related iterative approaches extend direct alignment methods into a loop where the current policy generates its own candidate responses, which are then re-ranked (by a judge model or reward model) to produce fresh preference pairs for the next round of DPO-style training, effectively closing the loop between generation and preference-based improvement without requiring PPO-style online rollouts.

A recognized failure mode across the DPO family is that, because the loss only ever compares the relative probability of chosen versus rejected responses, it can inadvertently decrease the absolute probability of the chosen response even while correctly increasing the margin between chosen and rejected, a phenomenon sometimes called likelihood displacement; several proposed fixes add an explicit supervised fine-tuning term or a floor constraint on the chosen response's likelihood to counteract this effect.

## Computational Considerations

DPO's computational profile is substantially lighter than PPO-based RLHF: a training step requires only two forward-and-backward passes through the policy being trained (one for the chosen response, one for the rejected response) plus two forward-only passes through the frozen reference model to obtain reference log-probabilities, none of which require the generation (sampling) step that PPO-based RLHF needs at every training iteration. Because sampling from a language model is typically far slower than a single forward pass (autoregressive generation requires one forward pass per generated token, whereas computing log-probabilities of an already-given sequence requires only one forward pass total), removing the online generation step is the single largest source of DPO's computational advantage over RLHF.

The reference model's log-probabilities for both the chosen and rejected responses depend only on the fixed reference policy and the fixed preference dataset, so they can be precomputed once, in a single pass over the entire training set before training begins, and cached to disk, entirely eliminating the need to hold the reference model in memory during the main training loop; this optimization roughly halves the peak memory and compute footprint of DPO training compared to keeping both the policy and reference model resident simultaneously.

Because DPO's loss is computed per preference pair rather than per token, and because preference pairs from the same prompt can have substantially different response lengths, care is needed in batch construction: naively batching chosen and rejected responses of very different lengths together wastes computation on padding tokens, and length imbalances between chosen and rejected responses can introduce a systematic length bias in the learned implicit reward, since longer sequences accumulate more total log-probability mass purely as an artifact of having more tokens, a known DPO pathology addressed by length-normalizing the log-probability sums used in the loss.

## Practical Implementation Strategies

A practical DPO training pipeline begins with a model that has already undergone supervised fine-tuning (SFT) on instruction-following data, since DPO is designed to refine an already-reasonable policy's behavior at the margin rather than to teach a base language model how to follow instructions from scratch; running DPO directly on a base (non-instruction-tuned) model tends to produce poor results. The reference model is almost always initialized as an exact copy of this SFT model's weights and then frozen for the duration of DPO training.

The beta hyperparameter requires careful tuning: values that are too small allow the policy to drift very far from the reference model per unit of preference-loss reduction, risking degenerate outputs or catastrophic forgetting of general capabilities, while values that are too large overly constrain the policy and produce only marginal alignment improvements; typical values used in practice range from roughly 0.1 to 0.5, and practitioners commonly sweep this parameter alongside learning rate on a held-out validation split of preference data.

Preference data quality and diversity matter enormously in practice: preference pairs where the chosen and rejected responses differ only trivially (for example, differing by a single word with no substantive quality difference) provide a weak or noisy training signal, while pairs with a large but well-justified quality gap provide much more useful signal; many practical pipelines filter or up-weight preference pairs based on the confidence or margin of the underlying preference judgment (whether from human annotators or an AI judge) to concentrate training signal on the most informative comparisons. It is also standard practice to interleave a small amount of the original SFT loss into DPO training, or to monitor SFT-loss regression on held-out data during DPO training, as a safeguard against the likelihood-displacement pathology degrading the model's general fluency and factuality.

## Benchmark Datasets & Evaluation

Anthropic's Helpful and Harmless (HH-RLHF) dataset and OpenAI's summarization preference dataset were among the earliest large-scale human preference datasets used to train and evaluate both RLHF and, later, DPO-style methods, each consisting of tens to hundreds of thousands of pairwise comparisons between model-generated responses. UltraFeedback, a more recent large-scale dataset, uses GPT-4 as an automated judge to score responses from a diverse pool of models along multiple quality axes (helpfulness, honesty, instruction-following, and harmlessness), and has become one of the most widely used preference datasets for training open-source DPO models due to its scale and the elimination of costly human annotation.

Direct alignment methods are typically evaluated using automated pairwise win-rate benchmarks such as AlpacaEval and MT-Bench, in which a strong judge model (commonly GPT-4 or a comparably capable model) compares the outputs of the DPO-trained model against a fixed reference model (often the original SFT checkpoint or a well-known strong baseline like GPT-3.5) across a diverse set of prompts, reporting the percentage of prompts on which the trained model's response is judged superior. Because these automated judges have their own known biases (favoring longer responses, favoring responses stylistically similar to their own outputs, and exhibiting positional bias depending on which response is presented first), careful benchmark design randomizes response order and often applies length-controlled win-rate corrections to avoid rewarding models that have merely learned to produce longer, more verbose outputs without genuine quality improvement.

Beyond win-rate benchmarks, alignment quality is also assessed through targeted safety and honesty evaluations (measuring refusal rates on harmful requests, susceptibility to jailbreak prompts, and calibration of expressed uncertainty), as well as through regression testing on standard capability benchmarks (such as general knowledge and reasoning tasks) to confirm that preference alignment has not degraded the model's underlying capabilities, a failure mode sometimes referred to as an "alignment tax."

## Key Challenges & Limitations

DPO and its variants are fundamentally limited by the quality and coverage of the offline preference dataset used to train them: because there is no online exploration or generation step during training (unlike PPO-based RLHF, which continually samples fresh responses from the current policy), DPO can only ever refine the policy's behavior on the distribution of responses represented in the static preference dataset, and cannot discover or reinforce genuinely novel high-quality response strategies that never appeared in the training data, a limitation increasingly addressed by iterative and online variants of direct preference optimization.

Reward hacking and specification gaming, well-known failure modes in RLHF, also manifest in DPO training, typically as an increase in response verbosity (since human and AI preference judgments often correlate, whether intentionally or not, with length and apparent thoroughness) without a genuine improvement in helpfulness, requiring explicit length-controlled evaluation and, in some pipelines, length-penalized variants of the DPO loss to counteract this drift.

The theoretical equivalence between DPO and RLHF derived in the original paper relies on several assumptions, including that the preference data is generated according to the Bradley-Terry model, that the policy has sufficient capacity to represent the optimal solution, and that the reference model used during training matches the one used to originally generate the (implicitly assumed) preference data; violations of these assumptions in practice, particularly the use of off-policy preference data collected from a different model's outputs than the one being trained, can weaken the theoretical guarantees and produce sub-optimal alignment.

## Hyperparameter Tuning

Beyond the beta regularization coefficient, effective DPO training requires tuning the learning rate, which is typically set substantially lower than the learning rate used for the initial SFT stage (often by an order of magnitude, in the range of 1e-7 to 1e-6 for large models), reflecting the fact that DPO is intended to make a comparatively small refinement to an already-competent policy rather than a large-scale weight update; excessively high learning rates during DPO training are strongly associated with rapid degeneration into repetitive or incoherent outputs.

The number of training epochs over the preference dataset is another sensitive hyperparameter: because DPO's loss can, in principle, be driven arbitrarily low by overfitting to the specific preference pairs seen during training, most practical pipelines train for only a small number of epochs (often just one to three passes over the preference dataset) and monitor a held-out validation split of preference pairs for early stopping, since continued training well past the point of validation-loss improvement reliably produces degraded output quality even as the training loss continues to decrease.

Batch size and gradient accumulation settings interact with the effective strength of the DPO gradient signal in practice, since preference losses computed over larger batches average over more comparisons and therefore produce smoother, lower-variance gradient estimates; in production pipelines, effective batch sizes (accounting for gradient accumulation across many GPUs) commonly range from the tens to low hundreds of preference pairs per optimizer step, following patterns similar to those used in standard large-batch fine-tuning of large language models.

## Real-World Applications & Case Studies

Direct preference optimization has become the standard final training stage for the majority of widely used open-weight instruction-tuned language models released since 2023, frequently applied after an initial SFT stage and, in some pipelines, layered on top of or in place of a separate RLHF stage, as teams have found that DPO achieves comparable or superior alignment quality with substantially less engineering overhead and infrastructure investment than maintaining a full PPO-based RLHF pipeline.

Beyond general-purpose chat assistants, direct preference optimization methods have been applied to domain-specific alignment tasks including code generation (aligning models to prefer functionally correct, efficient, and well-documented code over superficially plausible but subtly incorrect code), mathematical reasoning (using preference pairs derived from whether a generated solution reaches the correct final answer, sometimes combined with process-level preferences over intermediate reasoning steps), and multimodal vision-language alignment (extending DPO-style losses to align image-conditioned response generation with human preferences about factual grounding and reduced hallucination).

Iterative and online extensions of DPO, in which a model is repeatedly fine-tuned on preference data generated by comparing its own most recent outputs (self-rewarding or self-play style pipelines), have been used in several notable model releases to achieve multiple successive rounds of alignment improvement without returning to a full PPO-based RLHF pipeline, illustrating how the simplicity of the direct alignment loss has enabled more rapid, lower-overhead iteration cycles in production model development.

## Integration with Other Methods

Direct preference optimization is commonly combined with parameter-efficient fine-tuning techniques such as LoRA (Low-Rank Adaptation), in which only a small set of low-rank adapter weights are updated during DPO training while the base model's weights remain frozen, substantially reducing the memory footprint of preference alignment and allowing multiple distinct preference-aligned variants of the same base model to be maintained cheaply as separate small adapter files.

DPO-style losses are frequently combined with constitutional-AI-style pipelines, in which an AI system (rather than a human annotator) generates preference labels by evaluating candidate responses against a fixed set of written principles or a constitution, producing large volumes of preference data without direct human labeling effort while retaining human oversight over the higher-level principles that guide the automated judgments; this combination has become a dominant pattern for scaling preference-based alignment beyond what human annotation budgets alone could support.

Retrieval-augmented and tool-using language model pipelines increasingly incorporate a DPO-style alignment stage specifically targeting faithfulness to retrieved evidence or correct tool invocation, using preference pairs that contrast a grounded, accurate response against a fluent but unfaithful or hallucinated one, extending the core direct-alignment methodology beyond general conversational quality into more structured correctness objectives relevant to retrieval-augmented generation and agentic tool use.

## Future Research Directions

An active area of research seeks to close the gap between the purely offline nature of standard DPO and the online exploration capability of PPO-based RLHF, through iterative and online direct-alignment variants that periodically regenerate fresh preference data from the current policy's own outputs, aiming to combine DPO's implementation simplicity with RLHF's ability to reinforce genuinely novel high-quality behaviors discovered during training rather than being confined to a fixed offline dataset.

Extending direct preference optimization to align on process-level, multi-step, or long-horizon objectives, rather than single-turn response preferences, is an emerging research direction relevant to aligning agentic systems that take many sequential actions or reasoning steps before producing a final outcome, where credit assignment across the full trajectory (deciding which intermediate steps deserve credit or blame for the final preference judgment) is substantially more difficult than in the single-response setting for which DPO was originally derived.

Theoretical work continues to refine the understanding of exactly when and why direct alignment methods succeed or fail relative to RLHF, including formal characterizations of the likelihood-displacement phenomenon, the precise conditions under which DPO's implicit reward matches a well-calibrated explicit reward model, and principled approaches to combining multiple, potentially conflicting sources of preference signal (human annotations, AI judgments, and rule-based verifiers) within a single, unified direct-alignment training objective.

## Summary & Key Takeaways

Direct Preference Optimization replaced the multi-model, RL-based RLHF pipeline with a single closed-form loss derived by substituting the KL-regularized RLHF objective's optimal policy expression directly into the Bradley-Terry preference model, algebraically eliminating the need for an explicit reward model and reducing preference alignment to a supervised-learning-style classification loss over chosen and rejected response pairs. The DPO family has since expanded to include IPO (bounded loss for robustness to noisy or non-separable preferences), KTO (unpaired desirable/undesirable labels drawn from prospect theory), and ORPO (single-stage training without a separate reference model), each targeting specific weaknesses of the original formulation.

Despite its simplicity, DPO requires careful attention to the beta regularization strength, a substantially reduced learning rate relative to SFT, early stopping to avoid overfitting the static preference dataset, and length-normalization or length-controlled evaluation to avoid reward hacking through verbosity. Its fundamentally offline nature, limited to refining behavior already represented in the training distribution, is the central tradeoff against RLHF's more expensive but more exploratory online training loop, motivating ongoing research into iterative and online direct-alignment variants.

Keywords: direct preference optimization, DPO, RLHF alternative, Bradley-Terry preference model, KL-regularized reward maximization, Identity Preference Optimization, IPO, Kahneman-Tversky Optimization, KTO, Odds Ratio Preference Optimization, ORPO, implicit reward model, reference policy, likelihood displacement, preference data alignment, constitutional AI feedback, LoRA fine-tuning, AlpacaEval MT-Bench, length-controlled win rate, iterative online DPO

---

## Appendix: Practical Labs

### Lab 1: Deriving and computing the DPO loss from policy log-probabilities

import numpy as np


def sigmoid(x):
    """Numerically stable sigmoid."""
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))


def dpo_loss(logp_policy_chosen, logp_policy_rejected,
             logp_ref_chosen, logp_ref_rejected, beta=0.1):
    """Computes the DPO loss and the implicit reward margin for a batch of
    preference pairs, given log-probabilities of the chosen and rejected
    responses under both the policy being trained and the frozen reference
    model. All inputs are 1D arrays of shape (batch_size,)."""
    policy_logratio_chosen = logp_policy_chosen - logp_ref_chosen
    policy_logratio_rejected = logp_policy_rejected - logp_ref_rejected

    logits = beta * (policy_logratio_chosen - policy_logratio_rejected)
    losses = -np.log(sigmoid(logits) + 1e-12)

    implicit_reward_chosen = beta * policy_logratio_chosen
    implicit_reward_rejected = beta * policy_logratio_rejected
    reward_margin = implicit_reward_chosen - implicit_reward_rejected

    return losses.mean(), reward_margin


def test_dpo_loss_decreases_as_margin_increases():
    rng = np.random.RandomState(0)
    batch_size = 500

    # Reference model log-probs: arbitrary fixed baseline
    logp_ref_chosen = rng.uniform(-20, -5, batch_size)
    logp_ref_rejected = rng.uniform(-20, -5, batch_size)

    # Case A: policy barely differs from reference (weak preference signal)
    logp_policy_chosen_weak = logp_ref_chosen + 0.01
    logp_policy_rejected_weak = logp_ref_rejected - 0.01
    loss_weak, margin_weak = dpo_loss(
        logp_policy_chosen_weak, logp_policy_rejected_weak,
        logp_ref_chosen, logp_ref_rejected, beta=0.1
    )

    # Case B: policy strongly upweights chosen, downweights rejected
    logp_policy_chosen_strong = logp_ref_chosen + 3.0
    logp_policy_rejected_strong = logp_ref_rejected - 3.0
    loss_strong, margin_strong = dpo_loss(
        logp_policy_chosen_strong, logp_policy_rejected_strong,
        logp_ref_chosen, logp_ref_rejected, beta=0.1
    )

    print(f"Weak-margin DPO loss:   {loss_weak:.4f} (mean margin {margin_weak.mean():.4f})")
    print(f"Strong-margin DPO loss: {loss_strong:.4f} (mean margin {margin_strong.mean():.4f})")

    assert loss_strong < loss_weak, "Larger implicit reward margin should yield lower DPO loss"
    assert margin_strong.mean() > margin_weak.mean()

    # Loss must be strictly positive (it is a negative log-sigmoid) and finite
    assert loss_weak > 0 and np.isfinite(loss_weak)
    assert loss_strong > 0 and np.isfinite(loss_strong)

    print("DPO loss margin-sensitivity test passed.")


if __name__ == "__main__":
    test_dpo_loss_decreases_as_margin_increases()

### Lab 2: Beta sensitivity and the reward/KL tradeoff

import numpy as np


def sigmoid(x):
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))


def dpo_gradient_magnitude(policy_logratio_chosen, policy_logratio_rejected, beta):
    """Computes the magnitude of the DPO gradient's scaling coefficient, which
    is beta * (1 - sigmoid(beta * margin)): this term multiplies the gradient
    of the log-probabilities themselves, and dictates how strongly the loss
    pushes the policy in either direction for a given preference pair."""
    margin = policy_logratio_chosen - policy_logratio_rejected
    logits = beta * margin
    weight = beta * (1.0 - sigmoid(logits))
    return weight


def test_beta_controls_gradient_and_kl_tradeoff():
    # A moderately well-separated preference pair (policy already prefers chosen
    # somewhat over rejected relative to the reference model)
    policy_logratio_chosen = 1.5
    policy_logratio_rejected = -0.5

    betas = np.array([0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0])
    weights = np.array([
        dpo_gradient_magnitude(policy_logratio_chosen, policy_logratio_rejected, b)
        for b in betas
    ])

    print(f"{'beta':>6} | {'gradient weight':>16}")
    for b, w in zip(betas, weights):
        print(f"{b:6.2f} | {w:16.6f}")

    # For a fixed, already-somewhat-correct margin, larger beta produces a
    # LARGER effective per-example weight when the margin is small relative
    # to 1/beta, since beta scales the logits fed into the sigmoid; but as
    # beta grows very large, the sigmoid saturates and the weight approaches
    # zero because the pair is already "confidently correct" under a large
    # effective KL penalty. We check the intermediate non-monotonic behavior.
    assert weights[0] < weights[1], "Weight should initially increase with beta from a very small value"

    # At very large beta, the model is already confident this pair is correctly
    # ordered (since margin is positive), so the gradient weight should shrink
    # back down as beta continues to increase past a moderate value. The peak
    # occurs around beta ~ 1/margin; beta must grow well past that point
    # before the sigmoid saturation dominates and the weight collapses toward
    # zero, so we compare the largest beta against the peak region rather
    # than against the immediately preceding value.
    peak_weight = weights.max()
    assert weights[-1] < peak_weight, "Weight should shrink well below its peak once beta is large enough to saturate the sigmoid"
    assert weights[-1] < weights[3], "Weight at the largest beta should fall back below the weight at a moderate beta"

    # Larger beta always implies a tighter effective KL constraint: verify the
    # policy log-ratio magnitude implied by a fixed target reward is inversely
    # proportional to beta (reward = beta * logratio => logratio = reward / beta)
    target_reward = 2.0
    implied_logratios = target_reward / betas
    assert np.all(np.diff(implied_logratios) < 0), "Larger beta must imply smaller policy deviation for the same reward"

    print("Beta reward/KL tradeoff test passed.")


if __name__ == "__main__":
    test_beta_controls_gradient_and_kl_tradeoff()

### Lab 3: Comparing DPO (sigmoid loss) and IPO (bounded squared loss) under label noise

import numpy as np


def sigmoid(x):
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))


def dpo_loss_from_margin(logratio_diff, beta):
    logits = beta * logratio_diff
    return -np.log(sigmoid(logits) + 1e-12)


def ipo_loss_from_margin(logratio_diff, beta, tau=1.0):
    """IPO's loss is a bounded squared-error objective pulling beta times the
    log-ratio difference toward a fixed target (typically 1/(2*tau)), rather
    than pushing it toward positive infinity as DPO's sigmoid loss does."""
    target = 1.0 / (2.0 * tau)
    return (beta * logratio_diff - target) ** 2


def simulate_training_step(logratio_diff, loss_fn, lr=0.1, **kwargs):
    """A simplified proxy for one gradient step: since d(logratio_diff)/d(theta)
    is treated as a fixed unit direction for this illustrative comparison, we
    numerically differentiate the loss with respect to logratio_diff itself
    and step in the descent direction."""
    eps = 1e-4
    loss_plus = loss_fn(logratio_diff + eps, **kwargs)
    loss_minus = loss_fn(logratio_diff - eps, **kwargs)
    grad = (loss_plus - loss_minus) / (2 * eps)
    return logratio_diff - lr * grad


def test_ipo_stays_bounded_while_dpo_diverges_on_easy_pairs():
    beta = 0.1
    # An "easy" (already well-separated, low-noise) preference pair: DPO's
    # sigmoid loss has no minimum and will keep pushing the margin toward
    # +infinity indefinitely if training continues, whereas IPO's target-based
    # squared loss stabilizes once the margin reaches its fixed target.
    logratio_diff_dpo = 0.5
    logratio_diff_ipo = 0.5

    n_steps = 200
    dpo_trajectory = [logratio_diff_dpo]
    ipo_trajectory = [logratio_diff_ipo]

    for _ in range(n_steps):
        logratio_diff_dpo = simulate_training_step(
            logratio_diff_dpo, dpo_loss_from_margin, lr=2.0, beta=beta
        )
        logratio_diff_ipo = simulate_training_step(
            logratio_diff_ipo, ipo_loss_from_margin, lr=2.0, beta=beta, tau=1.0
        )
        dpo_trajectory.append(logratio_diff_dpo)
        ipo_trajectory.append(logratio_diff_ipo)

    print(f"DPO logratio-diff after {n_steps} steps: {dpo_trajectory[-1]:.3f}")
    print(f"IPO logratio-diff after {n_steps} steps: {ipo_trajectory[-1]:.3f}")

    # DPO's margin should keep growing without bound (it has no stationary point)
    assert dpo_trajectory[-1] > dpo_trajectory[len(dpo_trajectory) // 2], (
        "DPO's unbounded sigmoid loss should keep increasing the margin throughout training"
    )

    # IPO's margin should converge to a finite target and stop growing
    late_change = abs(ipo_trajectory[-1] - ipo_trajectory[-20])
    assert late_change < 0.05, "IPO's bounded loss should converge to a stable margin, not keep growing"

    ipo_target = 1.0 / (2.0 * 1.0) / beta
    assert abs(ipo_trajectory[-1] - ipo_target) < 0.1, "IPO should converge near its fixed target margin"

    print("DPO-vs-IPO boundedness test passed.")


if __name__ == "__main__":
    test_ipo_stays_bounded_while_dpo_diverges_on_easy_pairs()

### Lab 4: KTO-style unpaired preference loss vs. paired DPO under limited data

import numpy as np


def sigmoid(x):
    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))


def dpo_paired_loss(logp_policy_chosen, logp_policy_rejected,
                     logp_ref_chosen, logp_ref_rejected, beta):
    logratio_chosen = logp_policy_chosen - logp_ref_chosen
    logratio_rejected = logp_policy_rejected - logp_ref_rejected
    logits = beta * (logratio_chosen - logratio_rejected)
    return -np.log(sigmoid(logits) + 1e-12).mean()


def kto_unpaired_loss(logp_policy, logp_ref, is_desirable, beta, reference_kl=0.0):
    """A simplified KTO-style loss: desirable examples are pushed to have a
    HIGHER implicit reward than a reference KL baseline, undesirable examples
    are pushed to have a LOWER implicit reward, using only a per-example
    binary label rather than a paired comparison."""
    logratio = logp_policy - logp_ref
    implicit_reward = beta * logratio

    desirable_loss = 1.0 - sigmoid(implicit_reward - reference_kl)
    undesirable_loss = 1.0 - sigmoid(reference_kl - implicit_reward)

    losses = np.where(is_desirable, desirable_loss, undesirable_loss)
    return losses.mean()


def test_kto_usable_with_unpaired_data_dpo_requires_pairs():
    rng = np.random.RandomState(2)
    n = 300

    # Simulate a realistic scenario: only 40% of prompts have a genuine paired
    # comparison (both a chosen and rejected response collected), the rest
    # only have a single response with a thumbs-up/thumbs-down label.
    logp_ref = rng.uniform(-15, -5, n)
    logp_policy = logp_ref + rng.normal(0.5, 1.0, n)  # policy has drifted somewhat
    is_desirable = rng.uniform(0, 1, n) > 0.4  # 60% desirable, 40% undesirable

    # KTO can consume ALL n examples directly, no pairing required
    kto_loss = kto_unpaired_loss(logp_policy, logp_ref, is_desirable, beta=0.1)
    assert np.isfinite(kto_loss) and kto_loss > 0

    # DPO requires explicit pairs; simulate that only examples which happen to
    # share overlapping context can be paired (here, arbitrarily, the first
    # 40% of desirable/undesirable examples are paired up)
    desirable_idx = np.where(is_desirable)[0]
    undesirable_idx = np.where(~is_desirable)[0]
    n_pairs = min(len(desirable_idx), len(undesirable_idx))

    assert n_pairs < n, "Only a subset of unpaired data can be reconstructed into pairs"

    chosen_logp_policy = logp_policy[desirable_idx[:n_pairs]]
    chosen_logp_ref = logp_ref[desirable_idx[:n_pairs]]
    rejected_logp_policy = logp_policy[undesirable_idx[:n_pairs]]
    rejected_logp_ref = logp_ref[undesirable_idx[:n_pairs]]

    dpo_loss = dpo_paired_loss(
        chosen_logp_policy, rejected_logp_policy,
        chosen_logp_ref, rejected_logp_ref, beta=0.1
    )
    assert np.isfinite(dpo_loss) and dpo_loss > 0

    print(f"KTO loss (using all {n} unpaired examples): {kto_loss:.4f}")
    print(f"DPO loss (using only {n_pairs} reconstructable pairs out of {n}): {dpo_loss:.4f}")
    print(f"Fraction of data usable by DPO's pairing requirement: {n_pairs / n:.1%}")

    assert n_pairs / n < 1.0, "DPO's pairing requirement discards some fraction of unpaired feedback"

    print("KTO unpaired-data utilization test passed.")


if __name__ == "__main__":
    test_kto_usable_with_unpaired_data_dpo_requires_pairs()