Optimal-Transport Direct Preference Optimization

• OTPO is a design pattern that reformulates preference learning as transport between distributions, emphasizing semantic token differences.
• It replaces uniform token weighting in DPO with context-aware weights from an unbalanced transport plan, reducing length bias in reward estimates.
• Empirical studies show OTPO variants improve alignment performance, reward stability, and interpretability in language model training.

Searching arXiv for the specified papers and closely related optimal-transport preference-learning work. {"queries":[{"query":"id:(Zhu et al., 2 Apr 2026)"},{"query":"id:(Li et al., 24 May 2025)"},{"query":"id:(Melnyk et al., 2024)"}]} I retrieved the relevant arXiv records for the three cited works and will ground the article in those sources: "PLOT: Enhancing Preference Learning via Optimal Transport" (Zhu et al., 2 Apr 2026), "Optimal Transport-Based Token Weighting scheme for Enhanced Preference Optimization" (Li et al., 24 May 2025), and "Distributional Preference Alignment of LLMs via Optimal Transport" (Melnyk et al., 2024). Optimal-Transport Direct Preference Optimization (OTPO) denotes a family of alignment methods that incorporate optimal transport into direct preference optimization and adjacent preference-learning objectives for LLMs. In the most specific current usage, OTPO is the method introduced in "Optimal Transport-Based Token Weighting scheme for Enhanced Preference Optimization," which replaces uniform token weighting in DPO by an unbalanced optimal-transport coupling over chosen and rejected token representations (Li et al., 24 May 2025). Closely related work uses the same optimal-transport perspective to formulate token-distribution alignment in PLOT, explicitly describing the reformulation of direct preference optimization as an optimal-transport problem ("OTPO") (Zhu et al., 2 Apr 2026), and to formulate reward-distribution alignment in Alignment via Optimal Transport (AOT) (Melnyk et al., 2024). Across these variants, the common objective is to make preference optimization globally informed rather than purely local, while preserving useful distributional structure.

1. Scope, terminology, and variant structure

The term Optimal-Transport Direct Preference Optimization is not fully standardized across the recent literature. One paper uses OTPO as the method name for a context-aware token weighting scheme on top of DPO, another describes PLOT as a reformulation of direct preference optimization as an optimal-transport problem, and a related distributional method is presented as AOT but also exposited as an OTPO-style approach. This suggests that OTPO is best understood as a design pattern for OT-based preference optimization rather than a single canonical loss (Li et al., 24 May 2025, Zhu et al., 2 Apr 2026, Melnyk et al., 2024).

Variant	Core OT object	Preference setting
OTPO	Transport plan between chosen and rejected token hidden states	Paired DPO
PLOT / OTPO reformulation	Transport between model output distribution $Q_\theta$ and target preference distribution $P_t$	Fine-tuning-based alignment
AOT	One-dimensional OT between reward distributions	Paired and unpaired

Despite their differences, these methods share several structural commitments. First, they recast preference learning as transport between distributions or empirical measures rather than as a purely sample-level scalar comparison. Second, they use the transport plan or transport cost to emphasize semantically meaningful differences. Third, they frame OT as a way to control alignment updates without arbitrarily distorting the original model distribution.

2. Motivation from the limitations of conventional DPO

The OTPO literature starts from a critique of standard DPO. In the token-decomposed DPO objective, each token contributes equally to the reward difference, so longer but semantically irrelevant spans can inflate the margin and dominate the gradient. The 2025 OTPO paper identifies this as a source of “length bias,” and argues that uniform weighting does not distinguish tokens that carry the core answer or instruction-following content from noisy or low-value tokens (Li et al., 24 May 2025).

PLOT motivates OT from a broader alignment perspective. It states that existing methods remain limited by modest performance gains, high computational costs, hyperparameter sensitivity, and insufficient modeling of global token-level relationships. The proposed remedy is a token-level loss derived from optimal transport that adds a global, distribution-level penalty rather than penalizing only local token-by-token KL deviations (Zhu et al., 2 Apr 2026).

A related limitation appears at the distributional level. AOT observes that pairwise human preferences at the sample level do not imply alignment on the distributional level. It therefore formulates alignment as first-order stochastic dominance of the positive reward distribution over the negative reward distribution, then relaxes that condition into an OT objective (Melnyk et al., 2024). Taken together, these criticisms define the main OTPO agenda: mitigate uniform-token artifacts, encode semantic structure, and move from local pairwise comparisons toward global distributional alignment.

3. OTPO as context-aware token weighting in DPO

In the 2025 OTPO formulation, the chosen and rejected responses are treated as two discrete distributions in the model’s hidden space, using last-layer token representations $h_c^i$ and $h_r^j$. The cost matrix is defined by Euclidean distance,

$M_{i,j} = \|h_c^i - h_r^j\|_2,$

and the transport plan $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$ is obtained from an unbalanced, entropic OT problem with both entropy and marginal-KL regularization (Li et al., 24 May 2025): $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$

This transport plan induces token weights through marginalization,

$$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$

which are then normalized to sum to a fixed budget $\tau$, with $\tau=\min(|y_c|,|y_r|)$ by default. The weighted reward difference is

$P_t$0

where

$P_t$1

The final loss becomes

$P_t$2

The training algorithm follows the standard DPO workflow up to the reward decomposition, then inserts an OT solve. For each mini-batch, the model extracts last-layer hidden states, computes the cost matrix, solves the unbalanced entropic OT problem, converts the coupling into row and column weights, normalizes those weights, computes the weighted log-likelihood-ratio margin, and updates $P_t$3 with AdamW. The reported hyperparameters are batch size $P_t$4, epochs $P_t$5, learning rate $P_t$6, $P_t$7, $P_t$8, $P_t$9, and $h_c^i$0.

Conceptually, this OTPO variant does not replace DPO’s Bradley–Terry-style preference objective; it replaces DPO’s implicit uniform token weighting. The claimed advantages are a more contrastive reward-difference estimate, greater reward stability, improved interpretability, and better focus on meaningful differences between responses.

4. PLOT and the OTPO reformulation of token-distribution alignment

PLOT formulates preference learning as an optimal-transport problem directly over token distributions rather than over pairwise token hidden states. Let $h_c^i$1 denote the model’s current output distribution over $h_c^i$2 tokens, $h_c^i$3 a target preference distribution, and $h_c^i$4 a cost matrix with entries $h_c^i$5. The discrete OT problem is (Zhu et al., 2 Apr 2026)

$h_c^i$6

subject to

$h_c^i$7

In expanded form,

$h_c^i$8

The fine-tuning loss augments a baseline objective by the OT term: $h_c^i$9 The stated interpretation is that PLOT preserves the overall shape of the original distribution because mass is only “re-routed,” not arbitrarily zeroed or blown up.

A distinctive element of PLOT is how it constructs the target preference distribution. From a dataset $h_r^j$0, token frequencies $h_r^j$1 and $h_r^j$2 are counted in positive and negative answers. A raw difference is then formed: $h_r^j$3 followed by non-negativity and normalization: $h_r^j$4

PLOT also builds semantic information into the transport cost through the model’s token-embedding matrix $h_r^j$5, with $h_r^j$6. The paper uses

$h_r^j$7

while noting that one could more generally use $h_r^j$8 or cosine distance. This design is intended to make mass travel more cheaply between semantically close tokens.

The implementation path is correspondingly direct. After precomputing $h_r^j$9 and the cost matrix, each training batch produces $M_{i,j} = \|h_c^i - h_r^j\|_2,$0, solves the OT problem against $M_{i,j} = \|h_c^i - h_r^j\|_2,$1, computes $M_{i,j} = \|h_c^i - h_r^j\|_2,$2, and backpropagates the combined loss. For one-dimensional transport, the paper notes that one can compute the Wasserstein distance after sorting tokens by embedding norms $M_{i,j} = \|h_c^i - h_r^j\|_2,$3: $M_{i,j} = \|h_c^i - h_r^j\|_2,$4 If $M_{i,j} = \|h_c^i - h_r^j\|_2,$5 is large, entropy-regularized OT or Sinkhorn may be used with $M_{i,j} = \|h_c^i - h_r^j\|_2,$6–$M_{i,j} = \|h_c^i - h_r^j\|_2,$7. The reported hyperparameters are $M_{i,j} = \|h_c^i - h_r^j\|_2,$8, $M_{i,j} = \|h_c^i - h_r^j\|_2,$9, batch size $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$0, epochs $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$1, and either direct $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$2 for one-dimensional vectors or Sinkhorn for full $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$3, with $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$4 per batch.

5. Relation to distributional OT alignment and first-order stochastic dominance

AOT places OT preference optimization on a more explicitly statistical footing by defining alignment through first-order stochastic dominance of one reward distribution over another. In this formulation, if $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$5 and $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$6 are reward random variables induced by positive and negative data, then the desired condition is

$\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$7

where $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$8 is the $\Gamma \in \mathbb{R}_{\ge 0}^{|y_c|\times |y_r|}$9-th quantile of $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$0 (Melnyk et al., 2024). Rather than enforcing this exact condition, AOT introduces a convex surrogate

$$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$1

with examples such as the $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$2-squared hinge, $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$3-logistic, and least-squares-with-margin surrogates.

The key OT identity is

$$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$4

so the stochastic-dominance relaxation becomes a one-dimensional optimal-transport problem on reward distributions. For empirical measures with equal sample size, the solution has a closed form by sorting: $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$5 where $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$6 and $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$7 are order statistics. The paper further states that sorting is $$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$8-Lipschitz and piecewise-constant in its Jacobian, so one can backpropagate by sorting at each step and differentiating with the permutation held fixed; a soft-sort variant via entropic OT is also possible.

This formulation is not DPO with token reweighting. Instead, it generalizes the OT perspective from token-level alignment to reward-distribution alignment. Its relevance to OTPO lies in the shared claim that preference optimization can be expressed more naturally through transport between structured distributions than through isolated scalar margins. The statistical analysis is also distinctive: under stated boundedness and complexity assumptions, the empirical AOT objective achieves the parametric rate

$$\Gamma^* = \arg\min_{\Gamma \ge 0} \left\{ \sum_{i,j}\Gamma_{i,j} M_{i,j} +\epsilon_1\sum_{i,j}\Gamma_{i,j}\log\Gamma_{i,j} +\epsilon_2\Bigl[ \mathrm{KL}(\Gamma\mathbf{1},\mathbf{1}_{|y_c|}) +\mathrm{KL}(\Gamma^\top\mathbf{1},\mathbf{1}_{|y_r|}) \Bigr] \right\}.$$9

6. Empirical findings, stability, and unresolved issues

The empirical record for OTPO-style methods is distributed across several task families. In instruction-following experiments, the 2025 OTPO paper evaluates Llama-3-8B and Llama-3.2-3B on UltraFeedback (61 k pairs) and HelpSteer2 (9 k pairs), and Qwen-2.5-3B on TL;DR (92 k pairs). It reports LC-WR improvements of $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$0–$$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$1 over DPO across four instruction settings, including $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$2 versus $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$3 on Llama-3-8B/UltraFeedback; a summarization win-rate improvement of $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$4 over the best baseline; consistent or minor changes below $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$5 on general benchmarks; human pairwise evaluation in which OTPO is rated best by both experts at approximately $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$6 versus $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$7 for DPO; and interpretability in the form of a $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$8 correlation between OT weights and leave-one-out contributions from an explicit reward model (Li et al., 24 May 2025).

PLOT reports gains across safety, helpfulness, and reasoning. On HarmBench ASR for Llama3.2-3B + DPO, the baseline is approximately $$\omega_c^i = \sum_j \Gamma_{i,j}, \qquad \omega_r^j = \sum_i \Gamma_{i,j},$$9 and $\tau$0PLOT reaches $\tau$1, a $\tau$2 relative change. Similar gains are reported for SFT ($\tau$3), PRO ($\tau$4), and AOT ($\tau$5). In Human Values, helpfulness measured by a reward model improves from $\tau$6 to $\tau$7, helpfulness under GPT-4 from $\tau$8 to $\tau$9, and humanity under GPT-4 from $\tau=\min(|y_c|,|y_r|)$0 to $\tau=\min(|y_c|,|y_r|)$1. In Logic & Problem Solving, GSM8K improves from $\tau=\min(|y_c|,|y_r|)$2 to $\tau=\min(|y_c|,|y_r|)$3, MATH from $\tau=\min(|y_c|,|y_r|)$4 to $\tau=\min(|y_c|,|y_r|)$5, reasoning from $\tau=\min(|y_c|,|y_r|)$6 to $\tau=\min(|y_c|,|y_r|)$7, coding from $\tau=\min(|y_c|,|y_r|)$8 to $\tau=\min(|y_c|,|y_r|)$9, and STEM from $P_t$00 to $P_t$01. The ablations report that removing embeddings from the cost matrix degrades ASR, with an example of $P_t$02 on SFS, and that DEFT’s reward-based distribution term gives ASR approximately $P_t$03 whereas the OT reformulation gives $P_t$04. Varying $P_t$05 from $P_t$06 to $P_t$07 of its default keeps PLOT’s ASR within approximately $P_t$08–$P_t$09, versus DEFT at $P_t$10–$P_t$11. On AlpacaEval LC Win Rate, DPO drops from $P_t$12 to $P_t$13, whereas $P_t$14PLOT recovers to $P_t$15, which the paper presents as less erosion of fluency and coherence. The reported wall-clock overhead is approximately $P_t$16–$P_t$17 (Zhu et al., 2 Apr 2026).

AOT reports strong results in the 7B regime. On UltraFeedback with Merlinite-7B, AlpacaEval (GPT-4) scores are $P_t$18 for the base model, $P_t$19 for DPO, $P_t$20 for KTO, $P_t$21 for IPO, $P_t$22 for paired AOT, and $P_t$23 for unpaired AOT. The paper states that AOT is state-of-the-art in the 7B family as of May 2024, and its ablations indicate that larger batch sizes improve performance, the logistic surrogate outperforms hinge or squared surrogates, soft-sort and hard-sort give virtually identical scores, and the best margin is approximately $P_t$24 (Melnyk et al., 2024).

Two points of clarification recur in discussions of OTPO. First, optimal transport is not presented as computationally prohibitive in these papers, but its cost profile depends on the formulation: the 2025 OTPO method states that the OT solve per pair is $P_t$25 in tokens, with end-to-end training time approximately DPO plus $P_t$26 overhead, while PLOT reports approximately $P_t$27–$P_t$28 wall-clock overhead for its setup. Second, OT-based preference optimization is not limited to paired preference data: AOT explicitly treats both unpaired and paired settings. The main unresolved issues are also formulation-specific. The 2025 OTPO paper notes additional OT hyperparameters requiring mild tuning, experiments limited to mid-sized models, unexplored fully iterative on-policy loops, and evaluation dependence on GPT-4 judges with potential biases. PLOT frames OT as more stable and robust, but its reported demonstration uses batch size $P_t$29 and one epoch. A plausible implication is that OTPO remains an active methodological category whose internal variants differ materially in what is transported, how the cost is defined, and where computational and statistical bottlenecks appear.