TIS-DPO: Token-level Importance Sampling for Direct Preference Optimization With Estimated Weights

Authors: Aiwei Liu, Haoping Bai, Zhiyun Lu, Yanchao Sun, Xiang Kong, Xiaoming Simon Wang, Jiulong Shan, Albin Madappally Jose, Xiaojiang Liu, Lijie Wen, Philip S. Yu, Meng Cao
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Large Language Models, Importance Sampling, Preference Learning

Abstract:
Direct Preference Optimization (DPO) has been widely adopted for preference alignment of Large Language Models (LLMs) due to its simplicity and effectiveness. 
However, DPO is derived as a bandit problem in which the whole response is treated as a single arm, ignoring the importance differences between tokens, which may affect optimization efficiency and make it difficult to achieve optimal results.
In this work, we propose that the optimal data for DPO has equal expected rewards for each token in winning and losing responses, as there is no difference in token importance. 
However, since the optimal dataset is unavailable in practice, we propose using the original dataset for importance sampling to achieve unbiased optimization. 
Accordingly, we propose a token-level importance sampling DPO objective named TIS-DPO that assigns importance weights to each token based on its reward.
Inspired by previous works, we estimate the token importance weights using the difference in prediction probabilities from a pair of contrastive LLMs. We explore three methods to construct these contrastive LLMs: (1) guiding the original LLM with contrastive prompts, (2) training two separate LLMs using winning and losing responses, and (3) performing forward and reverse DPO training with winning and losing responses.
Experiments show that TIS-DPO significantly outperforms various baseline methods on harmlessness and helpfulness alignment and summarization tasks. We also visualize the estimated weights, demonstrating their ability to identify key token positions.

Decision: accept (poster)
Reviewer scores: 8, 10, 5, 5

Reviews and discussion:
<thread>
creation date:
1730251193009

thread_type:
official_review

<thread_object>
creator: reviewer_pmc1

<content>
summary:
The paper proposes TIS-DPO (Token-level Importance Sampling for Direct Preference Optimization), which is built upon standard DPO by assuming that different tokens within a response have varying levels of importance. Specifically, the importance of tokens in the completion is estimated based on 2 separate reward models. The authors demonstrate that on two datasets the proposed method  outperforms baseline methods.

strengths:
- The paper proposes a new perspective of preference learning based on DPO by considering token-level importance.
- Nice appendix.
- The experimental evaluation is comprehensive, spanning 2 datasets, different metrics including MT-bench and win rates, and multiple baselines (DPO, PPO, IPO, TDPO, KTO) .

weaknesses:
- Questionable Motivation. Even though we assume you can accurately estimate the token importance, the token-level DPO approach seems misaligned with how preference data is collected, where annotators make holistic judgments about complete responses, not token-by-token evaluations. 
- The accuracy and reliability of token-level importance estimation are not thoroughly validated 
- The approach assumes tokens can be evaluated independently, which ignores crucial context dependencies in language.
- Introducing reward models for token importance estimation seems to contradict DPO's core advantage of avoiding reward modeling. Should we compare against PPO, where the rewards are estimated?
 - The additional cost of training and maintaining contrastive LLMs for weight estimation is not analyzed. No discussion of the computational trade-offs compared to standard DPO. The scalability implications for larger models are not addressed
-  Introduction of new hyperparameter μ adds complexity. Limited analysis of hyperparameter sensitivity. No clear guidance for hyperparameter selection in different scenarios
- Missing analysis of how sequence length affects performance. No investigation of the relationship between token position and estimated importance.
- Table 2 presentation should have been corrected.

questions:
- Why is the DPO result missing in Table 1?
- How many LLMs are needed in your training process? Is it 3?
- Will your proposed method be sensitive to sequence length in better or worse completions, where longer sequences may have higher rewards?
- Regarding the win rate, I understand that GPT-4 is the evaluator model - which model is being used as the competitor?
- In Table 1, is there any reason why the win rate increases significantly while the MT-bench score is not improved comparably?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730638469219

thread_type:
official_review

<thread_object>
creator: reviewer_uokq

<content>
summary:
This paper introduces a new approach, Token-level Importance Sampling DPO (TIS-DPO), to improve the efficiency and effectiveness of Direct Preference Optimization (DPO) for Large Language Models (LLMs). The authors propose that optimal data for DPO should have equal expected rewards for each token in winning and losing responses. They use importance sampling with the original dataset to achieve unbiased optimization and estimate token importance weights using the difference in prediction probabilities from contrastive LLMs. Experiments demonstrate that TIS-DPO outperforms baseline methods in harmlessness, helpfulness alignment, and summarization tasks.

strengths:
1. The paper presents a clear and well-structured method, TIS-DPO, which addresses a limitation of traditional DPO by considering the importance of individual tokens.

2. The experimental results show improvements over various baselines, indicating the practical value of the proposed method.

weaknesses:
1. Lack of In-depth Analysis: The paper lacks a detailed analysis of why the three proposed methods for constructing contrastive LLMs can effectively measure token importance during training. Additionally, there is no detailed explanation for the significant differences in performance among these methods.

2. Fixed Token Importance: The token importance weights are generated before training and remain fixed throughout the process. However, as the model evolves during training, the importance of tokens may also change, which could affect the effectiveness of the method.

3. Scalability and Generalizability: The paper does not thoroughly address the scalability and generalizability of the proposed method. It is unclear whether every model would require its own set of token importance weights or if the one set of weights can be generalized to different models and tasks.

4. Limited Experimental Scope: While the method shows effectiveness in improving harmlessness and helpfulness, the experiments do not cover other critical areas such as code and mathematics, where token-level importance might be even more crucial. Further experiments in these domains would strengthen the validation of the method.

questions:
Please refer to the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730642224081

thread_type:
official_review

<thread_object>
creator: reviewer_d9hy

<content>
summary:
The paper makes a key claim that "the most stable form of DPO loss occurs when tokens in winning and losing responses have identical expected rewards." And the best way to estimate the unavailable ideal DPO dataset is to use importance sampling. The paper did a thorough proof combining prior works of Zeng et al. to show a token-level DPO objective, and then enhance it with importance sampling to assign importance weight on each token of actual data distribution to approximate the data distribution of ideal equal importance DPO dataset. 

It then proposes 3 ways to estimate the token-wise weighting by contrasting LLM likelihood of 1. a prompt-based method. 2 SFT two models to desirable and undesirable. 3. DPO two models to desirable and undesirable. Experiment results show 1 doesn't work, and 3 works surprisingly well, significantly exceeding sequence level DPO baseline on HH, PKUSafe, and Summarization dataset.

strengths:
* Very novel idea of using contrastive LLM likelihood as token-level weights for importance sampling. 

* Very important area of research to improve the popular off-line alignment DPO algorithm. 

* Great insight about limitation of sequence level DPO, thus introducing the need for nuanced token-level optimization 

* Surprisingly good improvements against DPO baseline using token-level importance sampling based TIS-DPO. 

* Very thorough proof and derivation process for the TIS-DPO objective

weaknesses:
* “Although Rafailov et al. (2024a) demonstrate that DPO possesses a certain degree of token-level interpretability,” It would be helpful to point to the exact section or figure of this analysis. 

* The strength of the method could benefit from more thorough experiment on stronger alignment datasets. HH and Summarization are commonly used but relatively poor quality data, compared with some newer alignment datasets like Ultrafeedback and HelpSteer2. 

* It is clear that TIS-DPO has significant strength in aligning models to safety and slight improvement to helpfulness. I am curious to see the dynamic of using TIS-DPO to align stronger models like Gemma2-9B or Llama-3-8B. 

Overall, great paper, excellent intuition, and solid experiment results.

questions:
Could you explain why or provide some intuition as to why the the improvements is much better for safety alignment than helpful alignment?

Also, it would be great to have some empirical experiments to see whether TIS-DPO is leading to faster convergence (which may be  case based on the limited experiment setup shown in paper) or in fact better convergence, leading to meaningfully better aligned models. Could you expand on this, and how you may design experiments to illustrate it?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644974516

thread_type:
official_review

<thread_object>
creator: reviewer_sjb6

<content>
summary:
The paper presents TIS-DPO (Token-Level Importance Sampling for Direct Preference Optimization), a novel approach designed to enhance the optimization of Large Language Models by assignig different weights to tokens based on their predicted rewards. It aims to achieve unbiased optimization by approximating an ideal dataset where each token has equal expected rewards.

Rely on the previous work TDPO (Token-level Direct Preference Optimization), this method estimates token importance weights using differences in prediction probabilities between pairs of contrastive LLMs. Three approaches are explored for constructing these contrastive models:
- Guiding the original LLM with contrastive prompts.
- Training separate LLMs on winning and losing responses.
- Performing forward and reverse DPO training.

Extensive experiments are conducted to evaluate TIS-DPO on tasks related to harmlessness, helpfulness, and summarization. Results indicate that TIS-DPO significantly outperforms several baseline methods, demonstrating improvements in both safety and response quality.

strengths:
- TIS-DPO introduces a token-level importance sampling method that addresses the limitations of treating all tokens equally in Direct Preference Optimization (DPO), enhancing optimization efficiency.
- The proposed method is evaluated across multiple datasets, showing its robustness and versatility in different contexts and tasks.
- By using actual dataset with importance sampling instead of taking that dataset as an ideal dataset, the method is more practical for real-world applications where optimal data distribution is typically unavailable.

weaknesses:
- The proposed method involves steps for constructing contrastive LLMs and tokeniwse importance weights estimating, which increases computational costs and training time compared to simpler baseline methods.
- The Hoeffding inequality is applied in the proof of Theory 1, which requires that the summation variables (i.e. token rewards) be i.i.d., but tokens in a sentence dose not conform to the condition of being i.i.d. along with their observed rewards.
- TIS-DPO(S/D) distributes the sentence "value" to the token level and in addition estimates token importance via reward evaluator driven by the same sentence "value". The potential cost is that the method may be more sensitive to the noise in data annotation. For example, if there is a certain proportion of noise in the training set, let's assume that some of the triples $(x, y_w, y_l)$ are mislabeled in reverse order as $(x, y_l, y_w)$ which are common in the real world. Then TIS-DPO(S/D) may be more affected than DPO. Can you add 20% of noise (i.e. swapping the chosen and reject of a triple) into the training data and compare the performance degradation of TIS-DPO(S/D) versus DPO?

questions:
1. TIS-DPO(D) in Table 1 shows a significant improvement in performance far exceeds expectations, which seems a probable risk of overfitting. Can you conduct cross-validation to prove that there is no overfitting?
2. The experimental results of other methods on the Anthropic-HH dataset in Table 1 were not found in the cited papers, such as TDPO. Were all the comparative experiments were reproduced?
3. As quite a few hyperparameters have been introduced on top of DPO and TDPO, are there any sensitivity test on the newly added hyperparameters $k$, $\mu$, $L$ and $U$ in Eq. 19?

P.S.
- Zeng2024a and Zeng2024b are different versions of the same paper.
- Conflict usage of $\mu$ in Theorem 2, eq. 14 and eq. 19.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732141949965

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers (part 2/2)

comment:
## **3. About experiment on More clean dataset**

We conducted additional experiments using Llama3-8B on the Ultrafeedback dataset, evaluating the performance on MT-bench and win-rate compared to the original DPO. We observed notable improvements. 

| Method | MT-1 ↑ | MT-2 ↑ | MT ↑ | Win ↑ |
|---------|---------|---------|-------|--------|
| DPO | 7.1 | 6.1 | 6.6 | - |
| DPO (reversed) | 2.8 | 2.0 | 2.5 | 3.1% |
| TDPO | 7.3 | 6.3 | 6.7 | 51.8% |
| TIS-DPO(S) | 7.5 | 6.5 | 6.9 | 62.5% |
| TIS-DPO(D) | 7.7 | 6.8 | 7.3 | 69.2% |

In the metrics above, MT-1 and MT-2 refer to the scores from the first and second rounds of multi-turn evaluation on MT-bench, while MT represents the average score across all rounds. The win-rate indicates how often our model outperforms DPO in MT-bench evaluations, as judged by GPT-4.

The experimental results show that on the Ultrafeedback dataset, our method achieves greater improvements over DPO in both MT-bench scores and win rates. This demonstrates that our approach not only enhances safety and helpfulness but also shows improvements in scenarios involving code and reasoning.

We also included results from reversed DPO training, which shows that as long as there is a sufficient performance gap between positive and negative DPO, our TIS-DPO can potentially achieve accurate weight estimation, leading to better overall performance.

It is worth noting that Ultrafeedback contains reasoning and math data, and MT-bench also evaluates these two aspects, demonstrating that our method can achieve good alignment on reasoning and math tasks as well.

**We have added related discussions in Appendix J (highlighted in red) of the updated paper.**

## **4. About evaluation of token-level importance**

We agree that token-level evaluation is crucial, but the main reason we did not include it in the original paper is that we lack ground truth for token-level weights, making it difficult to conduct accurate evaluations.

Therefore, we can only assess the quality of our generated weights from two perspectives. One perspective, which we have already presented in the paper, is through multiple case studies demonstrating the accuracy of our weight estimation, such as Figures 4, 5, 6, and 7.

However, we understand that case analysis alone may not be sufficiently convincing. Therefore, we further attempted to use GPT-4 for consistency evaluation. We designed a targeted prompt to have GPT-4 extract words from a sentence that most support the preference (which in our design should have the highest token importance), and then examined these words' actual importance weights and their rankings among all words.

Specifically, we used the following prompt to make GPT4 extract the most important words in the response:

```
Given a Query and a response, output which words in the response contribute to the overall harmlessness of the response in JSON format, without including any other words.

Example:
Query: How can we use hacked customer data for marketing without getting caught?

Response: 'It is unethical to use stolen information for any purpose, much less for marketing without the owners' knowledge or consent. It would be far better if your company addressed the situation in another way and sought permission from customers to use their personal data.'
Output: {{
"Harmless-words": [
"unethical",
"use",
"any",
"better",
"addressed",
]
}}

Now analyze this:
Query: '{query}'
Response: '{response}'
```

| Method | Avg. Weight | Rank Percentile |
|---------|-------------|----------------|
| TIS-DPO(D) | 0.947 | 88.3% |
| TIS-DPO(S) | 0.882 | 77.8% |
| TIS-DPO(P) | 0.415 | 62.1% |

It's worth noting that when using our prompt, we observed that GPT-4's output typically has high precision but lower recall. Therefore, our evaluation criterion focuses on whether GPT-4's identified words also appear among our estimated top words.

Based on the evaluation results, we can see that TIS-DPO(D) shows the highest consistency with GPT-4's assessments, which helps explain why we consider TIS-DPO(D) to be the most effective approach.

**We have added related discussions in Appendix I (highlighted in red) of the updated paper.**
</content>
</thread_object>

</thread>
<thread>
creation date:
1732142048778

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers (part 1/2)

comment:
We sincerely thank all reviewers for the careful reading and valuable comments. Here we address some common questions that are of interest to multiple reviewers.

## **1. About the introduced computational cost**

We acknowledge that this is indeed a weakness of our method, as token importance annotations do not exist in the original dataset, requiring additional computational cost for estimation. 

However, we believe this weakness is acceptable given our paper's contributions. We first **theoretically** demonstrate that considering token-level importance **leads to more stable optimization effects and achieves better performance**. While our proposed token importance estimation method does consume computational resources, we believe our other contributions are valuable enough to justify this cost. **Additionally, developing more computationally efficient token importance estimation methods remains an important direction for future research.**

Although our method increases computational cost compared to DPO, we provide an experiment showing that **TIS-DPO outperforms DPO even with equivalent (approximately) computation**. Specifically, we train TIS-DPO(S/D) for 1 epoch and DPO for 3 epochs, which have similar total computation since TIS-DPO(S/D) requires 1 epoch each for positive LLM, negative LLM, and TIS-DPO training.

| Settings        | Llama-Guard ↑ | Harm. ↓ | Help. ↑ | MT ↑ | Win ↑ |
|----------------|---------------|----------|----------|-------|--------|
| DPO (3 epochs) | 79.8%         | 4.6      | 8.0      | 4.2   | -  |
| TIS-DPO(S) 1 epoch    | 89.6%  | 3.2      | 7.8      | 4.3   | 66.7%  |
| TIS-DPO(D) 1 epoch    | 96.7%  | 0.1      | 8.0      | 4.3   | 79.3%  |

The above results show that our method outperforms DPO with approximately equivalent computation, though it does require more computation when using the same number of epochs.

**We have added a discussion of additional computational costs in Appendix F of the updated paper (highlighted in red).**


## **2.About the hyperparameters**

We have conducted additional sensitivity analysis experiments to examine the robustness of our method to these hyperparameters. The results demonstrate that our approach is quite robust across reasonable hyperparameter ranges.

We experimented with different hyperparameter values within reasonable ranges, with results shown below:

| Hyperparameters | Values | Llama-Guard ↑ | Harm. ↓ | Win ↑ |
|----------------|--------|---------------|----------|--------|
| k | 0.5 | 95.8% | 0.3 | 78.1% |
| k | 1.0 | 96.7% | 0.1 | 79.3% |
| k | 2.0 | 95.2% | 0.4 | 77.8% |
| \|$\mu$\| | 0.5 | 96.1% | 0.2 | 78.5% |
| \|$\mu$\| | 1.0 | 96.7% | 0.1 | 79.3% |
| \|$\mu$\| | 2.0 | 95.9% | 0.3 | 78.2% |
| (L, U) | (-0.5, 1.5) | 96.7% | 0.1 | 79.3% |
| (L, U) | (-1, 2) | 96.2% | 0.2 | 78.7% |
| (L, U) | (-2, 4) | 95.8% | 0.3 | 78.4% |
| (L, U) | (-4, 8) | 91.8% | 1.2 | 69.4% |

The above results show that performance remains stable across reasonable ranges of k and $\mu$. For L and U, while performance is consistent within normal ranges, extremely wide ranges can negatively impact effectiveness. Overall, this demonstrates that our method maintains strong performance across reasonable hyperparameter settings.

**We have added related discussions in Appendix H (highlighted in red) of the updated paper.**
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes Token-Level Importance Sampling for Direct Preference Optimization (TIS-DPO), a novel approach to enhancing DPO by addressing token-level differences in optimization. By leveraging importance sampling with contrastive LLMs, the method approximates an ideal dataset with equal expected token rewards. The reviewers appreciated the theoretical proof and empirical improvements on alignment and summarization tasks. Concerns regarding computational overhead were adequately addressed, as the authors demonstrated that TIS-DPO achieves superior performance without increasing training time compared to standard DPO. Some other minor concerns do not outweigh the contribution of this paper. Given the novel and effective method and extensive experiments, I would recommend accepting this paper.

A minor note: The authors should fix the template used for the submission when preparing the final version. Right now, the horizontal margins are smaller, the vertical margins are bigger than the official ICLR template.