PaperlayerView exact revision on arXiv

Source-linked AI summary

Token-level Direct Preference Optimization

Yongcheng Zeng, Guoqing Liu, Weiyu Ma, Ning Yang, Haifeng Zhang, Jun Wang

arXiv:2404.11999v5cs.CLcs.AI

TL;DR

Existing preference-optimization methods evaluate complete responses even though language generation proceeds token by token, motivating more granular divergence control. TDPO introduces token-level optimization with forward KL constraints and a tokenized Bradley-Terry preference model; experiments report improved alignment, quality, and diversity relative to DPO and PPO-based RLHF. The paper also identifies limitations in its reformulation and points toward iterative and multi-turn alignment.

  • Problem

    Preference optimization commonly evaluates full responses with pairwise comparisons and KL divergence, although responses are generated sequentially at the token level.

  • Method

    TDPO optimizes policy at the token level, converts the Bradley-Terry model into a token-level preference model, and applies forward KL divergence constraints for each token.

  • Results

    TDPO improves the balance between alignment and generation diversity across text tasks and improves generated-response quality relative to DPO and PPO-based RLHF.

  • Takeaways & Limitations

    Token-level optimization offers a simple DPO-based route to better regulate divergence while improving alignment, quality, and diversity in the evaluated tasks.

  • Takeaways & Limitations

    The TDPO reformulation cannot use the same reward-difference cancellation strategy as DPO because the Bradley-Terry model does not reduce to reward differences in this formulation.

Abstract

from arXiv · show

Fine-tuning pre-trained Large Language Models (LLMs) is essential to align them with human values and intentions. This process often utilizes methods like pairwise comparisons and KL divergence against a reference LLM, focusing on the evaluation of full answers generated by the models. However, the generation of these responses occurs in a token level, following a sequential, auto-regressive fashion. In this paper, we introduce Token-level Direct Preference Optimization (TDPO), a novel approach to align LLMs with human preferences by optimizing policy at the token level. Unlike previous methods, which face challenges in divergence efficiency, TDPO incorporates forward KL divergence constraints for each token, improving alignment and diversity. Utilizing the Bradley-Terry model for a token-based reward system, TDPO enhances the regulation of KL divergence, while preserving simplicity without the need for explicit reward modeling. Experimental results across various text tasks demonstrate TDPO's superior performance in balancing alignment with generation diversity. Notably, fine-tuning with TDPO strikes a better balance than DPO in the controlled sentiment generation and single-turn dialogue datasets, and significantly improves the quality of generated responses compared to both DPO and PPO-based RLHF methods. Our code is open-sourced at https://github.com/Vance0124/Token-level-Direct-Preference-Optimization.

1. Introduction

DPO aligns language models using pairwise preferences and sentence-level KL divergence, but this does not match autoregressive token-level generation. TDPO addresses this mismatch by controlling token-level KL divergence to balance alignment and diversity while retaining DPO’s simplicity.

  • DPO trains from pairwise comparisons without explicit reward modeling and uses KL divergence to keep the policy close to a reference language model.
  • Sentence-level KL divergence may be less effective because responses are generated sequentially, token by token.
  • Reverse KL divergence can reduce generation diversity because its mode-seeking behavior favors narrower output distributions.
  • TDPO optimizes human-preference alignment sequentially by controlling KL divergence for each token and connecting sentence-level rewards to token generation through the Bellman equation.
  • TDPO converts the Bradley-Terry preference model to token level and incorporates forward KL constraints into its objective for improved divergence regulation.
  • Across text tasks, TDPO improves generated-response quality relative to DPO and PPO-based RLHF while addressing excessive KL divergence.

2. Related Works

LLM training commonly includes pre-training and supervised fine-tuning, but fine-tuned models can retain biases, ethical concerns, toxicity, and hallucinations. Alignment methods such as RLHF and DPO are therefore used to make outputs more responsive to human expectations.

  • LLM training generally progresses from unsupervised pre-training to supervised fine-tuning with task-specific data.
  • Fine-tuned models may retain societal biases, ethical concerns, toxicity, and hallucinations, motivating a subsequent alignment phase.
  • RLHF refines model outputs using qualitative human feedback to improve responsiveness to human expectations and ethical considerations.

3. Preliminaries

DPO formulates preference optimization from an RLHF objective under a reverse KL penalty, maps rewards to an optimal policy, and trains from Bradley-Terry pairwise comparisons. Its formulation uses a reference model, a trainable policy, and a preference dataset.

  • DPO starts from an RLHF objective in which prompts produce preferred and dispreferred completions.
  • The objective uses a human-preference dataset, reward function, reference policy, trainable policy, and reverse KL divergence penalty coefficient.
  • DPO derives a reward representation from the optimal policy under reverse KL divergence, including a prompt-dependent partition function.
  • The Bradley-Terry model represents pairwise preference probability using the rewards of two completions.
  • Substituting the reward representation into Bradley-Terry and applying negative log likelihood yields the DPO objective and its derivative.

4. Methodology

TDPO reformulates preference optimization at the token level, connecting token-wise rewards, policy mappings, and Bradley–Terry preferences while regulating sequential KL divergence. Its loss functions explicitly combine forward and reverse KL considerations to balance alignment and generation diversity.

  • Token-Level Formulation: Text generation is modeled sequentially, with each state comprising the prompt and prior response tokens, and each action representing the next token.The token-wise reward is defined for each state-action pair, and the language model predicts the next-token distribution conditioned on the prompt and generated prefix.
  • Token-Level Optimization: The proposed token-level objective replaces DPO’s sentence-level objective while retaining a policy-optimization formulation.The method derives a token-level objective from sequential decision-making and relates it to state-action, state-value, and advantage functions.
  • Preference Modeling: The Bradley–Terry preference model is converted into a token-level representation equivalent to the Regret Preference Model.The derivation assumes response rewards decompose into discounted token-wise rewards and uses the logistic sigmoid to express preference probabilities.
  • Preference Modeling: The reformulated Bradley–Terry model ties preference probabilities to the policy and reference policy through reward differences and sequential forward-KL differences.The advantage-function derivation naturally offsets response-dependent partition functions before producing the policy-based preference relation.
  • Loss Function: TDPO explicitly introduces sequential forward KL divergence while retaining reverse-KL regulation to balance alignment with generation diversity.The method’s objective uses KL constraints at the token level, and its parameters regulate deviation from the reference model and the balance of sequential KL divergence.
  • Loss Function: TDPO2 modifies TDPO1’s loss, using a separate scaling parameter to control the balance of sequential KL divergence.The derivation also analyzes gradient weighting, where the KL-related term increases emphasis when divergence disparities between preferred and dispreferred responses grow.

5. Experiments

Experiments on IMDb, Anthropic HH, and MT-Bench evaluate TDPO against DPO, PPO, f-DPO, and SFT for alignment, diversity, KL regulation, and response quality. Across these settings, TDPO2 generally provides a stronger balance between alignment and generation diversity while regulating KL divergence more effectively.

  • Datasets and evaluation: Experiments cover controlled sentiment generation on IMDb, single-turn dialogue on Anthropic HH, and multi-turn evaluation on MT-Bench.IMDb uses positive-sentiment movie-review prefixes; Anthropic HH uses human queries; MT-Bench evaluates open-ended questions across eight knowledge domains.
  • IMDb sentiment generation: TDPO1 and TDPO2 outperform DPO on the IMDb reward–KL frontier, with TDPO2 improving further over TDPO1.Higher reward is achieved while maintaining low KL divergence; larger α values in TDPO2 made reward optimization difficult.
  • KL-divergence regulation: TDPO2 exhibits better sequential KL-divergence control than DPO and TDPO1 on IMDb during training.DPO’s divergence grows faster on dispreferred than preferred responses, while TDPO1 increases divergence on preferred responses more than TDPO2.
  • Anthropic HH dialogue: TDPO2 surpasses DPO, f-DPO, and TDPO1 in both accuracy and entropy on Anthropic HH, achieving a superior alignment–diversity balance.Accuracy measures agreement with chosen completions, while entropy measures diversity from 25 nucleus-sampled responses at p = 0.95.
  • Anthropic HH dialogue: DPO, TDPO1, and TDPO2 achieve win rates of at least 50% against chosen responses at temperature 0.75, outperforming SFT in preference alignment.Win rates are computed for generated responses against chosen responses on the Anthropic HH test set.
  • MT-Bench evaluation: MT-Bench compares SFT, PPO, DPO, TDPO1, and TDPO2 using GPT-4-evaluated win, tie, and lose rates.Responses are sampled at temperature 0.7 with a maximum of 512 newly generated tokens; the reported result states that TDPO2 achieves a higher win rate.

6. Conclusion

TDPO is a token-level fine-tuning approach that uses token-wise optimization and forward KL constraints to align LLMs with human preferences. It surpasses DPO and PPO-based RLHF on controlled sentiment generation and single-turn dialogue tasks.

  • TDPO employs token-wise optimization with forward KL divergence constraints and a token-level Bradley-Terry preference model.The approach targets alignment, output quality, and diversity.
  • TDPO surpasses DPO and PPO-based RLHF in controlled sentiment generation and single-turn dialogue tasks.
  • Future alignment work is anticipated to benefit from iterative refinement and multi-turn conversational strategies.These approaches are presented as prospects for improving alignment with human values.

A.1. Proving the Relationship between Maximizing the Advantage Function and Enhancing the Expected Returns

The appendix establishes that nonnegative expected advantage under a policy implies a relationship between policy optimization and expected returns, then applies KL-constrained optimization to text generation.

  • The proof assumes nonnegative expected advantage for every state under the comparison policy.
  • The trajectory formulation evaluates actions sampled from the comparison policy while generating token sequences.
  • Maximizing the parameterized policy’s expected return is formulated through the state-action advantage under the reference policy.
  • A reverse KL divergence constraint is introduced to prevent excessive degradation of language models.
  • The constrained optimization problem admits a closed-form solution involving a partition function.
  • The Bradley-Terry preference model is connected to token-wise rewards through an overall reward decomposed across generated tokens.

A.4. Deriving the TDPO Objective Under the Bradley-Terry Model

This appendix derives the TDPO objective under the Bradley-Terry model by expressing preference probabilities through the optimal policy and sequential forward-KL differences.

  • The theorem states that the Bradley-Terry model expresses human preference probabilities in terms of the optimal policy for the KL-constrained advantage problem.
  • The derivation defines u as the implicit reward difference between the policy and reference models.
  • δ represents the β-weighted difference in sequential forward KL divergence between preferred and dispreferred responses.
  • Substituting the policy and KL expressions into the Bradley-Terry formulation yields the TDPO preference objective.

B. TDPO Implementation Details and Hyperparameters

The implementation provides PyTorch code for the TDPO loss, computing token-position KL terms between policy and reference logits and combining them with preference rewards. Default training uses α = 0.5, β = 0.1, batch size 64, RMSprop, and a warmed-up learning rate.

  • The implementation defines a PyTorch tdpo_loss function using policy logits, reference logits, preferred and dispreferred indices, labels, β, α, and a TDPO variant flag.
  • Token-position KL values are computed from reference and policy vocabulary log-probabilities for preferred and dispreferred completions.
  • The TDPO objective combines preferred-minus-dispreferred rewards with a KL difference between dispreferred and preferred completions.
  • Default settings use α = 0.5, β = 0.1, batch size 64, RMSprop, and learning rate 5e-6.The learning rate is warmed up linearly from 0 to 5e-6 over 150 steps.

C. Additional Experimental Results

The IMDb experiment examines expected reward and forward KL divergence relative to a reference model, alongside sequential forward KL divergence across preferred and dispreferred responses during training.

  • Figure 6(a) shows the frontier of expected reward and forward KL divergence with respect to the reference model.
  • Figures 6(b) and 6(c) show sequential forward KL divergence over training steps for preferred and dispreferred response subsets, respectively.
  • Figure 6(d) shows the difference between sequential forward KL divergence for dispreferred and preferred responses throughout training.
Loading 2404.11999v5…