Pith. sign in

REVIEW 3 major objections 7 minor 13 references

Optimal Transport-Based Token Weighting scheme for Enhanced Preference Optimization

T0 review · 3 major / 7 minor · reviewed 2026-08-07 · deepseek-v4-flash

Pith's one-line read OTPO replaces DPO's uniform token weights with an unbalanced optimal transport plan between chosen and rejected responses, upweighting shared, prompt-relevant tokens and downweighting noise, and reports up to 10.9% higher…

desk verdict Solid, well-engineered token weighting for DPO with a plausible mechanism, but the semantic-relevance premise is asserted rather than shown, so treat the headline interpretation with caution. read the letter →

arxiv 2505.18720 v1 pith:I2X47AKF submitted 2025-05-24 cs.CL cs.AIcs.LG

classification cs.CLcs.AIcs.LG
keywords optimaltransporttokenweightingdirectpreferenceoptimizationDPOalignmentlengthbiasunbalancedinstructionfollowing
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper claims that direct preference optimization (DPO) misaligns language models because it weights every token equally, so noisy or filler tokens can dominate the loss. The proposed fix, OTPO, computes an unbalanced optimal transport plan between the token sequences of the chosen and rejected responses, using Euclidean distances between the model's last-layer hidden states as the matching cost. The row and column sums of that plan become per-token weights that upweight shared, semantically relevant tokens and downweight divergent ones, and these weights replace the uniform weights inside the DPO reward difference. If the claim holds, preference optimization concentrates on the content humans actually compare, yielding more stable rewards and better instruction-following without explicit length penalties; the paper reports up to 10.9% higher length-controlled win rate on AlpacaEval2 over DPO and an 8.6% win-rate gain over the best baseline on TL;DR summarization.

What carries the argument

The central object is the unbalanced optimal transport plan $\Gamma^*$, defined in Eq. 9 as the minimizer of a matching cost $\sum_{i,j} \Gamma_{ij} M_{ij}$ plus an entropy regularizer and two KL penalties that keep the marginals near uniform. The cost is metric: $M_{ij} = \|h_i^c - h_j^r\|_2$, the Euclidean distance between last-layer hidden states of token $i$ in the chosen response and token $j$ in the rejected response. Its row and column sums (Eq. 7) produce the token-level weights $\omega_c^*$ and $\omega_r^*$, normalized by Eq. 10 to the common budget $\tau = \min(|y_c|, |y_r|)$. Inserting these weights into the DPO reward difference (Eq. 11) is what redirects the gradient toward semantically shared tokens; the mechanism therefore does not introduce any new supervision signal, only a geometric alignment of the two responses.

What would settle it

Train the same OTPO procedure but scramble the cost matrix, for instance by permuting the rows and columns of $M$ with a fixed random permutation per response pair, or by using distances from an untrained random-initialization layer. If the length-controlled win-rate gains over DPO persist under scrambled geometry, then the specific representation-driven weighting is not what carries the improvement. Alternatively, collect token-level human importance annotations on a sample of preference pairs and compare against the OT weights; a systematic disagreement would refute the semantic-relevance premise.

Watch

Extended reading notes

Core claim

On the paper's own terms, OTPO's discovery is that the reward difference in DPO should be weighted according to the semantic alignment between chosen and rejected tokens, and that this alignment can be read off an unbalanced optimal transport plan $\Gamma^*$ (Eq. 9). The cost matrix entry $M_{ij} = \|h_i^c - h_j^r\|_2$ measures how far token $i$ of the chosen response sits from token $j$ of the rejected response in the model's representation space, so similar tokens are cheap to couple. Summing the optimized plan along its two dimensions gives token weights $\omega_c^*$ and $\omega_r^*$ (Eq. 10), normalized to the budget $\tau = \min(|y_c|, |y_r|)$, which are then injected into the reward difference (Eq. 11) and the log-sigmoid loss (Eq. 12). The paper argues that this makes the reward difference more contrastive, more stable, and more interpretable, and reports consistent length-controlled win-rate improvements over DPO, SimPO, SamPO, and LDDPO across Llama-3-8B, Llama-3.2-3B, Qwen-2.5-3B, and Mistral-7B models.

Load-bearing premise

The load-bearing premise is that the Euclidean distance between the last-layer hidden states of two tokens measures how semantically related they are, so that tokens shared or similar across chosen and rejected responses are exactly the ones preference optimization should emphasize; if that geometry is not meaningful, the weights are arbitrary and the reported gains would not transfer.

Editorial extensions

If this is right

  • Length-controlled win rate on AlpacaEval2 rises by 2.6–10.9% over DPO and by 1.0–3.8% over the best prior baseline across four on-policy settings, with the largest gain on Llama-3-8B + UltraFeedback.
  • The weighting is a generalization of earlier length-bias fixes: SimPO's length normalization, SamPO's random downsampling, and LDDPO's tail down-weighting all become special cases of the token-weighting form in Eq. 5.
  • On TL;DR summarization, OTPO outperforms the best baseline (SamPO) by 8.6% win rate under a GPT-4o judge, consistent with the idea that emphasizing shared content yields more concise summaries.
  • OTPO's per-token weights correlate at 0.76 with leave-one-out explanations of the ArmoRM reward model, meaning the weighting is not arbitrary but tracks what an external reward model attends to.
  • The extra training cost is small: an $O(n^2)$ OT solve plus $O(n^2)$ memory per pair, which the paper argues is negligible next to the transformer forward/backward pass.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If hidden-state proximity is a reliable indicator of semantic relevance, the same OT weighting scheme should transfer to other preference objectives (IPO, KTO, distributional OT approaches), since the weighting is a drop-in replacement for the reward-difference decomposition rather than a DPO-specific trick.
  • The reported weight–explanation correlation suggests a post-hoc diagnostic: OT weights could be used as a cheap, model-intrinsic explanation of preference judgments, avoiding expensive leave-one-out scoring in future work.
  • A sharp testable boundary of the premise is cross-lingual transfer: Euclidean geometry in the last layer may not align shared meaning across languages as well as within English, so OTPO's gains could shrink on non-English preference data; the paper explicitly lists English-only training as a limitation.
  • Because the cost matrix is recomputed from the policy being trained, the weighting is coupled to the model's current representation; this suggests a potential bootstrapping failure mode if the model's representations are poor early in training, and a possible mitigation would be a fixed reference-model representation for the cost while keeping the policy for the log-ratios.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 7 minor

Summary. The paper proposes OTPO, a token-level weighting scheme for Direct Preference Optimization. It decomposes the DPO reward difference into per-token log-ratio contributions (Eqs. 4-7), computes an unbalanced optimal transport plan over chosen-rejected token pairs using Euclidean distances between last-layer hidden states (Eq. 9), derives token weights from the transport plan marginals (Eq. 10), and substitutes these weights into the DPO loss (Eqs. 11-12). The method is evaluated on instruction-following (AlpacaEval2) and summarization (TL;DR) across several backbone models and datasets, with claims of up to 10.9% relative length-controlled win-rate improvement over DPO. The paper also provides ablations for the OT component and weight-normalization choices, hyperparameter sensitivity experiments, and a human evaluation.

Significance. If the central claim holds, OTPO offers a principled, context-aware alternative to uniform token weighting in DPO, with potential benefits for alignment quality and interpretability. The paper has notable strengths: the token-level decomposition in Sec. 2.2 is mathematically correct; the unbalanced OT problem in Eq. 9 is well-posed; the ablations in Sec. 4.2 support the design choices over simple smoothing; the code is available; and the experimental evaluation spans multiple models, datasets, and tasks, including off-policy summarization and on-policy instruction-following. However, the significance is tempered by two serious gaps: the gradient derivation in Appx. K.3 treats the OT weights as constants even though they depend on the policy parameters, and the core semantic-relevance premise is not supported by systematic evidence. These issues have load-bearing consequences for the method's theoretical soundness and its stated interpretation.

major comments (3)
  1. [Appx. K.3, Eqs. 30-32] The gradient derivation in Appx. K.3 treats the token weights ω* as constants when differentiating the loss. However, the weights are computed from the policy's last-layer hidden states (Eq. 9), so they depend on the policy parameters θ. The full gradient of L in Eq. 12 therefore includes additional terms ∂ω*/∂θ multiplied by the log-ratio differences; these are omitted from Eq. 32. If the implementation detaches the weights (stop-gradient), the paper must state this explicitly and redefine the objective accordingly; otherwise, the training procedure does not optimize the loss as written. This is a load-bearing correctness issue for the method's description and its gradient analysis.
  2. [Sec. 2.3 (Eqs. 9-10) and Appx. I] The paper's central interpretation is that OTPO emphasizes 'semantically meaningful token pairs and de-emphasizes less relevant ones.' The mechanism relies on the cost M_ij = ||h_i^c - h_j^r||_2, which drives the OT plan to concentrate mass on similar token pairs; hence the marginal weights in Eq. 10 upweight shared tokens and downweight unique tokens. In near-miss preference pairs—where the chosen and rejected responses are identical except for a few decisive tokens (e.g., 'Paris' versus 'Berlin')—those decisive tokens are unique and would receive low weight, inverting the stated mechanism. The only external evidence for the semantic-relevance claim, the 0.76 correlation with ArmoRM in Appx. I, is computed on a single illustrative pair and does not probe near-miss cases. The paper needs systematic evidence (e.g., constructing or annotating preference pairs with known key tokens and checking whether OTPO upweights them) to support the claim that OT weights track preference-relevant semantics.
  3. [Sec. 2.3 (Eq. 9) and Appx. F] The optimal transport formulation introduces several free hyperparameters (ε1, ε2, τ) whose choices are justified only heuristically. The sensitivity analysis in Appx. F shows that ε1 and ε2 substantially change the token-weight distribution and the normalization value τ, although the reward margin remains relatively stable. The paper does not explain why τ = min(|yc|, |yr|) is the right choice, nor does it report how the final AlpacaEval win rates vary with ε1 and ε2. Tab. 3 ablates the normalization choice, but not the entropy or KL penalty strengths. This leaves open the possibility that the reported gains depend on a narrow hyperparameter window, which weakens the generality claim.
minor comments (7)
  1. [Eq. 9] The notation KL(Γ1, 1_{|yc|}) is ambiguous: KL divergence is defined for probability distributions, but '1_{|yc|}' appears to denote a vector of ones rather than the uniform distribution. Please clarify that the second argument is the uniform distribution (e.g., u_{|yc|} = 1/|yc|).
  2. [Eq. 10] The expression 'ω*_c = Γ1 |Γ| τ' appears to be missing a division; based on the text, it should read (Γ1 / |Γ|) · τ. Please correct the typesetting.
  3. [Sec. 2.3, paragraph after Eq. 9] The statement that the first term of Eq. 9 'corresponds to the Wasserstein distance' is imprecise. That term is a transport cost; the unbalanced formulation with marginal penalties does not yield the Wasserstein distance in the standard sense.
  4. [Fig. 3] The prompt in Fig. 3 is written as 'What is the capital of Paris?' which is presumably a typo; it should likely be 'What is the capital of France?'.
  5. [Throughout] The method name is inconsistent: Sec. 2.2 refers to 'our OTWPO algorithm' while the rest of the paper uses OTPO. Please unify the terminology.
  6. [Appx. I] The reported correlation of 0.76 between OT weights and ArmoRM leave-one-out explanations is based on a single example. Please report the number of examples and the variance of the correlation across a sample of pairs.
  7. [Tab. 2] The significance asterisks (***) are said to indicate 99% confidence, but no significance test is described. Please specify the test (e.g., bootstrap over AlpacaEval prompts) and report the associated p-values or confidence intervals.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: OTPO's weighting is computed from model representations but is not fitted to the evaluation targets, and the central preference-optimization objective is an independent modification of DPO.

full rationale

The paper's derivation chain (Eq. 4 token-level decomposition, Eq. 5 weighted reward, Eq. 6-7 pair factorization, Eq. 9 unbalanced OT objective, Eq. 10 marginal weights, Eq. 11 weighted reward difference, Eq. 12 final loss) is algebraically self-contained: each step is an identity or a clearly defined objective, and no step assumes the conclusion it is used to support. The OT weights are computed from hidden-state Euclidean distances, which is a modeling assumption rather than a fitted parameter, and the reported AlpacaEval2 / TL;DR gains are measured against external benchmarks under fixed hyperparameter choices; there is no fitted-input-called-prediction relation between the OT plan and the evaluation metric. The motivating claim that shared or similar tokens are semantically relevant is an empirical premise, and the paper supplies independent, albeit limited, grounding via the 0.76 correlation with ArmoRM leave-one-out explanations (Appx. I); even if that premise were false, it would be a correctness risk, not circularity. The claim that OTPO 'emphasizes semantically meaningful token pairs' is in part true by construction because the cost matrix is a Euclidean distance and Eq. 9 minimizes transport cost, but the substantive claim, that this weighting improves preference alignment, is not derived from that definition. The only self-citation found (Han et al. 2024 in Related Work, sharing author Meng Li) is an example of OT applications and is not load-bearing. No circular step can be exhibited from the text, so the appropriate finding is no significant circularity.

Assumptions & free parameters 4 free parameters · 5 assumptions · 0 invented entities

The method inherits standard DPO assumptions and adds three hand-chosen hyperparameters plus a heuristic cost metric. The strongest additional assumptions are that Euclidean distance in the policy's hidden state space tracks semantic relevance and that shared tokens are the important ones; neither is proven. No new physical or conceptual entities are introduced, as the transport plan and token weights are mathematical constructions over existing hidden representations.

free parameters (4)
  • epsilon_1 (entropy regularization coefficient) = 1 for UltraFeedback and TL;DR; 0.1 for HelpSteer2
    Chosen per dataset in Appx A; controls sparsity of the transport plan and directly changes token weight differences (Appx F, Fig 10).
  • epsilon_2 (marginal KL penalty coefficient) = 0.2 for all configurations
    Fixed in Eq 9; penalizes deviation of transport marginals from uniform and affects the normalization value tau (Appx F).
  • tau (total weight budget) = min(|yc|, |yr|)
    Hand-chosen normalization level in Eq 10; ablation in Sec 4.2.2 shows mean, max, and per-response length normalizations are worse.
  • beta (DPO temperature) = 0.01 for Llama-3-8B and Qwen-2.5-3B; 0.1 for Llama-3.2-3B
    Inherited from DPO and tuned per model, not introduced by OTPO, but the final weighted loss depends on it (Appx A).
assumptions (5)
  • standard math DPO reward reparameterization and Bradley-Terry preference model (Eqs 1-3)
    Adopted from Rafailov et al. 2024b; the shared baseline for all methods compared.
  • ad hoc to paper Shared or semantically similar tokens between chosen and rejected responses are more likely to be relevant to the prompt and should be upweighted
    Stated in Sec 2.3 as the motivating observation; no formal proof or systematic validation beyond examples and one correlation.
  • ad hoc to paper Euclidean distance between last-layer hidden states is a meaningful semantic metric for token matching
    Used to build cost matrix M in Eq 9; the paper cites metric-space usage in Wasserstein GANs but does not validate that this distance tracks preference-relevant semantics.
  • ad hoc to paper Unbalanced OT with entropy and KL penalties yields token weights that improve preference optimization
    Design choice in Eq 9; supported only by ablation experiments and benchmark results.
  • ad hoc to paper OT weights can be treated as constants when differentiating the OTPO loss
    Appx K.3 computes the gradient ignoring theta-dependence of the weights; the paper does not state or justify a stop-gradient.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Optimal Transport-Based Token Weighting scheme for Enhanced Preference Optimization." pith.science (2026). https://pith.science/paper/I2X47AKF

@misc{pith2026250518720,
  author       = {Pith},
  title        = {Pith review of: Optimal Transport-Based Token Weighting scheme for Enhanced Preference Optimization},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/I2X47AKF}},
  note         = {Machine review of arXiv:2505.18720}
}
read the original abstract

Direct Preference Optimization (DPO) has emerged as a promising framework for aligning Large Language Models (LLMs) with human preferences by directly optimizing the log-likelihood difference between chosen and rejected responses. However, existing methods assign equal importance to all tokens in the response, while humans focus on more meaningful parts. This leads to suboptimal preference optimization, as irrelevant or noisy tokens disproportionately influence DPO loss. To address this limitation, we propose \textbf{O}ptimal \textbf{T}ransport-based token weighting scheme for enhancing direct \textbf{P}reference \textbf{O}ptimization (OTPO). By emphasizing semantically meaningful token pairs and de-emphasizing less relevant ones, our method introduces a context-aware token weighting scheme that yields a more contrastive reward difference estimate. This adaptive weighting enhances reward stability, improves interpretability, and ensures that preference optimization focuses on meaningful differences between responses. Extensive experiments have validated OTPO's effectiveness in improving instruction-following ability across various settings\footnote{Code is available at https://github.com/Mimasss2/OTPO.}.

Figures

Figures reproduced from arXiv: 2505.18720 by the authors.

Figure 1
Figure 1. The uniform weighting in DPO leads to sub [PITH_FULL_IMAGE:figures/full_fig_p001_1.png] view at source ↗
Figure 3
Figure 3. Weights assigned to the responses given different methods. Here, given the prompt [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗
Figure 4
Figure 4. TL;DR summarization win rates compared to [PITH_FULL_IMAGE:figures/full_fig_p006_4.png] view at source ↗
Figures from the paper (8 more)
Figure 5
Figure 5. Figure 5: Trend of gradient norm during training. fairness by considering token-level relationships between the chosen and rejected responses. How￾ever, it is slightly worse than OTPO in both length￾controlled win rate and win rate, as it only con￾siders the relationship of each…
Figure 6
Figure 6. Figure 6: Comparison of training time across differ [PITH_FULL_IMAGE:figures/full_fig_p009_6.png]
Figure 7
Figure 7. Figure 7: Dataset length difference (|yc| − |yr|) distri￾bution. Dash lines indicated the quartiles. with more negative length differences, having a median of -30 compared to -16 for UltraFeedback with Llama-3.2-8B. Notably, the length differences in the long-tail regions are si…
Figure 8
Figure 8. Figure 8: Estimated response log probability v.s. re [PITH_FULL_IMAGE:figures/full_fig_p018_8.png]
Figure 9
Figure 9. Figure 9: Reward difference v.s. the length differ [PITH_FULL_IMAGE:figures/full_fig_p018_9.png]
Figure 10
Figure 10. Figure 10: Hyperparameter sensitivity analysis of ϵ1 (top), ϵ2 (bottom). (left) Absolute token weight difference distribution between OTPO and DPO. (middle) The distribution normalization value τ in Eq.10. (right) Changes in reward margin. tor Hugo” are assigned with lower weigh…
Figure 11
Figure 11. Figure 11: Spearman correlation between token weight [PITH_FULL_IMAGE:figures/full_fig_p020_11.png]
Figure 12
Figure 12. Figure 12: Human Evaluation Guideline [PITH_FULL_IMAGE:figures/full_fig_p022_12.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

13 extracted references · 9 canonical work pages

  1. [1]

    Evaluators need to compare the two responses based on evaluation criteria listed below and select the better response

    Evaluation Format Each evaluation task presents a prompt and two responses (Response A and Response B). Evaluators need to compare the two responses based on evaluation criteria listed below and select the better response. If the responses are equally good or bad, evaluators can choose "Tie." Example format: Prompt: [Displayed] Response A: [Displayed] Res...

  2. [2]

    Evaluation Criteria Evaluators should compare the responses based on the following aspects: Relevance & Accuracy: Does the response correctly address the prompt? Is the information factually accurate and relevant? Coherence & Fluency: Is the response well-structured and grammatically correct? Does it read naturally and make logical sense? Completeness: Do...

  3. [3]

    Evaluation Options For each evaluation task, evaluators must select one of the following options:  Response A is better (A outperforms B in most criteria)  Response B is better (B outperforms A in most criteria)  Tie (Both responses are equally good or equally bad)

  4. [4]

    Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bi- lal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello

    A framework for few-shot language model evaluation. Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bi- lal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. 2024. A general theoretical paradigm to understand learning from human pref- erences. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics...

  5. [5]

    Payment is based on task completion and quality control measures to ensure reliable judgments

    Payment Evaluators will be compensated $0.25 per completed evaluation. Payment is based on task completion and quality control measures to ensure reliable judgments. Evaluators with consistently low-quality judgments may be disqualified from further participation. Thank you for contributing to this evaluation! Your judgments help improve AI model performa...

  6. [6]

    Advances in neural in- formation processing systems, 35:27730–27744

    Training language models to follow instruc- tions with human feedback. Advances in neural in- formation processing systems, 35:27730–27744. Richard Yuanzhe Pang, Weizhe Yuan, Kyunghyun Cho, He He, Sainbayar Sukhbaatar, and Jason Weston

  7. [7]

    In Advances in Neural Information Processing Sys- tems, volume 37, pages 116617–116637

    Iterative reasoning preference optimization. In Advances in Neural Information Processing Sys- tems, volume 37, pages 116617–116637. Curran As- sociates, Inc. Ryan Park, Rafael Rafailov, Stefano Ermon, and Chelsea Finn. 2024. Disentangling length from qual- ity in direct preference optimization. In Findings of the Association for Computational Linguistics...

  8. [12]

    The waves dance under the moon, a melody soft and bright

    Examples Example 1 (Clear Difference in Relevance & Accuracy) Prompt: What is the capital of France? Response A: Paris. Response B: Berlin. Correct Choice: Response A is better (B is factually incorrect). Example 2 (Tie due to Equal Performance) Prompt: Write a short poem about the ocean. Response A: "The waves dance under the moon, a melody soft and brig...

Show all 13 references
  1. [2017]

    In International conference on machine learning , pages 214–223

    Wasserstein generative adversarial networks. In International conference on machine learning , pages 214–223. PMLR. Yuki M. Asano, Christian Rupprecht, and Andrea Vedaldi. 2020. Self-labelling via simultaneous clus- tering and representation learning. In International Conferen...

  2. [2021]

    In International Conference on Learning Representations

    Measuring massive multitask language under- standing. In International Conference on Learning Representations. Haozhe Ji, Cheng Lu, Yilin Niu, Pei Ke, Hongning Wang, Jun Zhu, Jie Tang, and Minlie Huang. 2024. Towards efficient exact optimization of language model alignment. In...

  3. [2022]

    arXiv preprint arXiv:2204.05862

    Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862. Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. 2020. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of t...

  4. [2023]

    alignment-tax

    Unpacking the ethical value alignment in big models. arXiv preprint arXiv:2310.17551. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. HellaSwag: Can a ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Associat...

  5. [2024]

    In First Conference on Language Modeling

    Length-controlled alpaca‘eval: A simple de- biasing of automatic evaluators. In First Conference on Language Modeling. Rémi Flamary, Nicholas Courty, Davis Tuia, and Alain Rakotomamonjy. 2016. Optimal transport for do- main adaptation. IEEE Trans. Pattern Anal. Mach. Intell, 1...

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.