---
title: Optimal Transport-Based Token Weighting (OTPO)
url: https://www.emergentmind.com/topics/optimal-transport-based-token-weighting-otpo
type: topic
---

# Optimal Transport-Based Token Weighting (OTPO)

Optimal Transport-Based Token Weighting (OTPO) frameworks constitute a class of techniques in preference optimization that leverage optimal transport (OT) theory to assign adaptive, semantic-aware importance to tokens when aligning large language models (LLMs) to human preferences. These approaches replace uniform token weighting found in standard Direct Preference Optimization (DPO) with OT-derived weights, focusing the optimization objective on the most informative and consequential token alignments. By formulating reward differences or preference losses as optimal transport problems over token embeddings, OTPO schemes mitigate the impact of irrelevant or noisy tokens, reduce length bias, and yield more contrastive and robust alignment between model outputs and human intent [2505.18720].

## 1. Motivation and Conceptual Foundations

Standard DPO directly optimizes the log-likelihood difference between a preferred (chosen) and a non-preferred (rejected) sequence. In this regime, each token’s log-ratio contribution is weighted uniformly:
\[
\Delta_r = \sum_i q_c^i - \sum_j q_r^j
\]
where $q_c^i$ and $q_r^j$ denote log-likelihood ratios of tokens in the chosen and rejected responses, respectively. Uniform weighting allows irrelevant or high-frequency tokens to dominate the difference, confounding alignment and inducing length bias.

OTPO paradigms introduce a semantically informed reweighting mechanism. Instead of treating all token positions equally, OTPO assigns higher weights to token pairs that are semantically aligned—typically assessed via hidden state similarities—and down-weights less relevant, less matched tokens. This weighting is determined by solving an OT problem between the token-level embeddings of candidate and reference responses, yielding an alignment plan that focuses model optimization on meaningful distinctions [2505.18720].

The approach connects to broader applications of OT-based token weighting, including interpretable semantic similarity [2202.13196] and global distributional alignment in preference learning [2604.01837], but distinguishes itself via integration into DPO-style direct preference objectives.

## 2. Mathematical Formulation and Algorithm

The OTPO workflow is formalized as follows [2505.18720]:

**Token Embedding Extraction**  
Obtain last-layer hidden states for the chosen $(y_c = [y_c^1, ..., y_c^m])$ and rejected $(y_r = [y_r^1, ..., y_r^n])$ responses:
\[
h_c^i,\, h_r^j \in \mathbb{R}^d
\]

**Cost Matrix Construction**  
Define $C\in\mathbb{R}^{m\times n}$ by Euclidean distances:
\[
C_{ij} = \| h_c^i - h_r^j \|_2
\]

**Regularized Unbalanced OT Objective**  
Solve for optimal transport $\Gamma^{*}\in\mathbb{R}_{\ge 0}^{m\times n}$:
\[
\Gamma^* = \arg\min_{\Gamma\geq 0} \sum_{ij} \Gamma_{ij}C_{ij}
+ \epsilon_1\sum_{ij} \Gamma_{ij}\log\Gamma_{ij}
+ \epsilon_2\Bigl( \mathrm{KL}(\Gamma\mathbf{1}_n \|\mathbf{1}_m)
+ \mathrm{KL}(\Gamma^\top\mathbf{1}_m \|\mathbf{1}_n)\Bigr)
\]
The entropy regularizer ($\epsilon_1$) and KL terms ($\epsilon_2$) control the smoothness and marginal fidelity, supporting “unbalanced” mass.

**Token Weight Computation**
Row and column sums yield raw weights:
\[
\omega_c^i = \sum_j \Gamma_{ij},\quad \omega_r^j = \sum_i \Gamma_{ij}
\]
Normalize both to a budget $\tau = \min(m, n)$ for stability.

**OT-Weighted Preference Difference**
Form the OT-weighted reward difference:
\[
\widehat{\Delta}_r = \sum_i \omega_{c}^{*i}\,q_c^i - \sum_j \omega_{r}^{*j}\,q_r^j
\]

**Loss Objective**  
Use the same sigmoid cross-entropy as DPO:
\[
\mathcal{L}_{\mathrm{OTPO}} = -\mathbb{E}_{(x,y_c,y_r)\sim D}\left[ \log\,\sigma\left( \beta \widehat{\Delta}_r \right) \right]
\]

**Algorithmic Steps**  
1. Compute $h_c^i, h_r^j$ for each paired response.
2. Form $C_{ij}$.
3. Solve OT via unbalanced Sinkhorn iterations.
4. Sum and normalize weights.
5. Compute $q_c^i, q_r^j$ from model/reference.
6. Compute $\widehat{\Delta}_r$.
7. Backpropagate $\mathcal{L}_{\mathrm{OTPO}}$.

**Hyperparameters**  
- $\epsilon_1$: entropy regularization, typically $[0.1, 1.0]$.
- $\epsilon_2$: KL marginal fidelity, typically fixed at $0.2$.
- $\tau$: weight budget, set to $\min(m, n)$.
- Sinkhorn iterations: $20$--$50$.
- Complexity: $O(L^2)$ per instance ($L$ is sequence length); little overhead compared to Transformer forward pass.

## 3. Comparison to Alternative OT-Based Token Weighting

OTPO is part of a broader movement toward leveraging optimal transport for token- or distribution-level alignment in NLP:
- **RCMD/CLRCMD** [2202.13196]: “Relaxed Contextualized Mover’s Distance” computes OT-based sentence distances for semantic similarity by cost matrices over token embeddings with (typically cosine) distance, but enforces looser (relaxed) marginal constraints. While RCMD emphasizes interpretability and efficient, sparse matching, OTPO emphasizes precise semantically weighted contributions to the reward difference in RLHF.
- **PLOT** [2604.01837]: “Preference Learning via Optimal Transport” constrains OT at the vocabulary distribution level, comparing model output pseudo-distribution $Q_\theta$ to a data-defined target $P_t$ using a cost matrix derived from absolute pairwise differences of embedding norms. The final loss is a (Wasserstein-1) scalar, not a position-weighted objective. While PLOT exploits global OT-loss, OTPO explicitly focuses on token-level assignment and gradient routing.

A summary table contrasting representative OT-based schemes:

| Method      | OT Domain            | Weighting Granularity     |
|-------------|---------------------|--------------------------|
| OTPO [2505.18720]   | Token embeddings      | Token/position-level      |
| RCMD [2202.13196]   | Token embeddings      | Token/nearest-neighbor   |
| PLOT [2604.01837]   | Vocab distributions   | Token-frequency/global   |

## 4. Implementation Details and Practical Considerations

OTPO leverages established OT numerical libraries such as PythonOT’s unbalanced-OT and POT Sinkhorn. The chosen distance is typically Euclidean, although others (e.g., cosine) may be substituted. Budget normalization ($\tau$) ensures reward difference magnitudes remain consistent across length-variable responses, addressing length bias.

Algorithmic stability is maintained through regularization hyperparameters ($\epsilon_1$, $\epsilon_2$) and through robust normalization of weights. Empirically, OTPO loss rewards remain stable across wide hyperparameter intervals.

Runtime overhead is minimal: each step incurs $\sim1.05\times$ the compute of standard DPO, dominated by the $O(L^2)$ cost of the cost matrix and Sinkhorn, which is marginal compared to Transformer forward/backward passes.

Ablation studies reveal that naive weighting (e.g., uniform or heuristic embedding similarity) fails to match performance or reliability of full OT-based weighting.

## 5. Empirical Benchmarks and Comparative Performance

Extensive experiments in [2505.18720] demonstrate that OTPO delivers robust gains in instruction-following, summarization, and alignment. Key findings include:
- On AlpacaEval2 (Llama-3-8B + UltraFeedback), OTPO improves LC-win-rate from $48.14\%$ (DPO) to $53.37\%$.
- For Llama-3.2-3B, the improvement is from $26.02\%$ (DPO) to $26.97\%$ (OTPO).
- In TL;DR summarization with Qwen-2.5-3B, OTPO achieves an $+8.6\%$ win-rate improvement over the best baseline.
- Additional wins of $1$--$3\%$ absolute over other DPO variants are observed on HelpSteer2 across Qwen-2.5-3B and Mistral-7B.

Empirical reward margins and model alignment are robust to regularizer settings. Replacing OT with non-adaptive or purely similarity-based

Source: https://www.emergentmind.com/topics/optimal-transport-based-token-weighting-otpo