Transfer Q Star: Principled Decoding for LLM Alignment

Summary

This paper introduces 'Transfer Q*', which is a decoding strategy to align language models with target reward models without any fine-tuning. The main contribution of this paper is the estimation of the optimal Q* function, which is often unavailable in practice and is necessary for approximating the RL optimal policy, through DPO (Direct Preference Optimization). The paper first introduces an algorithm for direct transfer decoding using DPO-trained models as optimal baselines, and then presents indirect transfer decoding that uses these as conservative baselines. The authors describe the theoretical characteristics of this methodology and experimentally demonstrate that it achieves state-of-the-art performance on target reward models.

Strengths

Aligning to specific values through decoding algorithms without fine-tuning is highly useful. In particular, estimating the token-level Q* function in this process is very challenging, and existing works typically train external networks to estimate these values. This paper uses DPO-trained models as baselines to estimate the token-level Q* function, as DPO learns the Q* function in an offline manner. Given recent reports that DPO-trained models predict preferences better than classification-based reward models, this approach is very interesting and can be considered a valuable contribution. Additionally, since DPO-trained models are offline trainers and thus may be distant from the optimal Q* function, the introduction of an indirect transfer decoding algorithm that treats them as suboptimal, along with the presentation of its theoretical properties, is also a significant contribution.

Weaknesses

1. Equation 5 on line 121 is introduced roughly. Its derivation should be described in the Appendix, or references discussing it should be introduced. 2. The introduction of Controlled Decoding (CD) is incorrect. It states that CD uses $Q^{\pi_{sft}}$ as a proxy for Q*, but in reality, it uses an energy-based model of $Q^{\pi_{sft}}$ and a value function through an external network as Q*. The difference between this work and CD is that CD uses FUDGE and Q function as a baseline, while the proposed TQ* uses DPO as a baseline reward model. Therefore, Figure 1 and its explanation are incorrect. Considering that the indirect method, which conservatively estimates the target model, is the contribution, Figure 1 should be drawn in reverse. 3. Equation 12 on line 175 is believed to be incorrect. Theoretically, the $\pi(z|s)$ used in this equation should be the reference model $\pi_{ref}$ before exploration, not the aligned model $\pi_{BL}$. The equation development started with the premise that the optimal policy can be expressed as an energy-based model of the reference model and reward model, but in equation 11, it roughly transitions from the reference model to the aligned DPO, losing theoretical justification. In this case, it's expected that the original distribution will collapse significantly and over-optimize to the reward model. 4. A crucial experiment on the trade-off between KL divergence and the reward model is missing. In RL, there's a strong trade-off in reward model optimization depending on how much KL penalty is given, so it's not sufficient to only check the performance against the reward model as in this paper's experiments. While Figure 2(b) touches on this, it's not enough. An experiment showing whether it demonstrates pareto-frontier performance in the trade-off relationship according to KL penalty is also needed. This experiment is particularly important given the use of the DPO-trained model as the reference model, as mentioned in weakness 3.

Questions

1. In Figure 2 (b), when measuring the KL divergence for TQ*, which model was used to compare the decoding results? For a fair comparison, it should be measured against $pi_sft$, like the other baselines. 2. I'm curious about the inference cost. How do the memory requirements and time complexity compare to naive greedy decoding?

Rating

5

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

This paper has done truly promising work, but it has the following limitations: 1. Although DPO is used as a baseline, it is an offline trainer and there is a theoretical gap with the optimal Q*. This issue is addressed through the indirect transfer decoding part, but there is still room for improvement, and various discussions are likely to follow. 2. Due to the use of the DPO model as the initial policy, the theoretical justification is weakened, and a clear interpretation of the experiments is not possible. 3. In the process of training the DPO model, over-optimization towards the reward model is also controlled by the KL penalty. It would be interesting to add analysis and experiments on this aspect.

Summary

This paper addresses decoding for aligning large language models, which is a process of inference-time, token-level optimisation without updating the parameters of the LLM. Two scenarios are considered: 1. where a baseline policy is given and is aligned with the target trajectory-level reward. 2. where the baseline policy is given and aligned with a different trajectory-reward than the target one. Two approaches are provided in accordance: one named direct transfer, which as the name suggests directly derive a token-level policy from the baseline LLM, and the other named indirect transfer, whose key lies in an importance sampling trick to reweigh using the ratio between the target and baseline trajectory policies. The authors also provide theoretical analysis characterising the sub optimality gap and the divergence from the supervise-fine-tuned model. Empirical results demonstrates superior performance in several key metrics.

Strengths

This paper is very well written, poses a reasonable problem, namely how to more efficiently decode an existing trajectory level policy into a token level policy, and provides an elegant solution. The proposed approach draws on latest developments in the field (DPO) and enjoys a rigorous theoretical characterisation. The gap between the trajectory level learnt policy and the sft-policy is interpretably characterised in terms of the regularisation coefficients in front of the KL terms. The authors finished with a nice empirical analysis showing the efficacy of their method.

Weaknesses

Morally speaking, it is not clear to me why should decoding this way works much better than directly taking $\rho_{BL}$- at least in the case of direct decoding. After all, TQ* is just the 1-step-action-value function of $\rho_{BL}$, so I think the readers would benefit from an explicit explanation of why this process of *obtaining $\rho_{BL}$* --> *estimating its (1-step-action) value function* --> *compute the closed form of the optimal policy of this value function with a KL term* should do better than just directly using the trajectory level policy which is already the (regularised) optimal of the reward function $r$. More concretely, it would be nice to have this shown as a theoretical result - a theorem that shows that the suboptimality gap of the policy proposed by the authors is smaller than $\rho_{BL}$. I appreciate that the authors have shown this empirically, though. For the indirect case, the method proposed by the authors require access to the ground truth trajectory-level reward function. But if we had this in practice, would we not have just directly used it to compute $\rho_{BL}$? Thus, I think here it makes sense to add some explanation regarding why this is useful. Finally, I find it a bit of a misleading juxtaposition with DPO - the point of DPO is that we no longer need to model the reward function, but in this work we rely on both the trajectory level optimal policy learned by DPO as well as the reward function which also needs to be learned. Both of these suffer from statistical/optimisation instability in their own right. So it is strong to assume you have both of these.

Questions

equation (2) - This seems like a typo since for i > 0, the first argument of R should also take in the so-far generated sequence z_0 ... z_{i-1}. equation (15) - how is $Z_{BL}$ and $Z_r$ computed? This is needed later in the importance ratio, if I understood correctly, but I don't see how to compute these quantities which are usually intractable.

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Aside from the limitations I mentioned in Weaknesses, the authors also acknowledges that TQ* suffers from increased time complexity. They suggest that this can be mitigated by training a small value function adapter as discussed in prior work.

Summary

The paper Transfer Q⋆: Principled Decoding for LLM Alignment proposes a novel approach to aligning large language models by leveraging a principled decoding strategy. The authors propose a method to estimate the optimal Q-function for decoding using an existing aligned policy, addressing limitations in previous approaches like Controlled Decoding (CD). The paper presents an indirect transfer method, allowing for alignment even when the baseline model is trained on a different reward function. The authors provide theoretical analysis characterizing the sub-optimality gap and KL divergence to the reference policy. Extensive experiments across multiple datasets and model architectures demonstrate the effectiveness of TQ⋆ compared to existing methods.

Strengths

**Novelty**: The paper introduces an original approach to LLM alignment via decoding, leveraging existing aligned policies to estimate the optimal Q-function. **Theoretical foundation**: The authors provide a rigorous theoretical analysis of their method, including bounds on the sub-optimality gap and KL divergence. This adds credibility to the approach and helps explain its effectiveness. **Comprehensive evaluation**: The experimental section is thorough, covering multiple datasets, model architectures, and evaluation metrics. The inclusion of both synthetic and real transfer tasks demonstrates the method's robustness. **Practical relevance**: TQ⋆ addresses a significant challenge in LLM alignment, offering a computationally efficient alternative to fine-tuning approaches. This has potential implications for improving the deployment of aligned LLMs.

Weaknesses

**Comparison with Baseline**s: The comparisons with existing baselines like DPO are insightful, but additional baselines, especially those focusing on inference-time control, could strengthen the evaluation. **Hyperparameter sensitivity**: The paper does not thoroughly explore the sensitivity of TQ⋆ to its hyperparameters, particularly the decoding alignment parameter α. A more detailed analysis of this aspect would strengthen the work. **Scalability**: While the method is tested on 7B parameter models, it's unclear how well it scales to larger models that are increasingly common in practical applications. **Evaluation metrics**: While the paper uses several evaluation metrics, including GPT-4 based assessment, it lacks human evaluation studies. Given the subjective nature of language quality and alignment, human evaluation would provide valuable validation of the method's effectiveness.

Questions

1. Have you investigated the stability of the alignment achieved by TQ⋆ over extended generation sequences? Does the alignment quality degrade for longer outputs, and if so, how does this compare to other methods? 2. How sensitive is TQ⋆ to the choice of baseline model? If multiple baseline models are available, each aligned with different rewards, how might one optimally select or combine them to estimate Q⋆ for a given target reward?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors acknowledge some limitations of their work, such as the potential for hallucination in responses to very obscure queries. However, they could improve their discussion of limitations by addressing: 1. The reliance on existing aligned baseline models, which may not always be available or suitable for all target rewards. 2. The potential for errors or biases in the GPT-4 based evaluation, which is used as a proxy for human assessment. 3. The computational overhead of TQ⋆ compared to standard decoding methods, which may impact real-time applications.

Summary

The paper proposes a new estimation of Q function for controlled decoding. Instead of using a Q function derived from SFT model, the paper propose to use Q function derived from separate aligned models. Evaluation shows the proposed method can achieve higher reward in benchmarks.

Strengths

1. The idea of using aligned models to estimate Q is quite intuitive as these models closer to optimal policy than SFT model, also the paper proposed mathematical derivations and explanations on how and why aligned models are better for estimating Q. 2. The benchmarks showed the proposed method tend to achieve higher reward, especially it outperforms the controlled decoding baseline along with other alignment methods.

Weaknesses

1. The proposed approach seems computationally expensive to do decoding. 2. The approach require accessing an already aligned model, this adds more requirements for using the method and limits it use cases.

Questions

1. Did the authors compare decoding speed for different approaches? It would be useful to report even thought the proposed approach is slower. 2. For gpt4 based evaluation, why the authors only report Win-Tie instead of splitting win and ties? In particular I am curious what is the percentage of ties for each comparison.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are discussed in the appendix.

Thanks for the response

Thank you for addressing my concerns with additional experiments and baselines for larger models. I appreciate your commitment to adding comprehensive human evaluation. Based on these improvements, I am satisfied with maintaining my original score for your paper.

Additional New Experimental Evaluations

Dear Reviewers and ACs, Thank you for your time and efforts in reviewing our paper and rebuttal discussions. ***Additional New Experimental Results:*** To further strengthen our empirical evaluations, we ran experiments and obtained Pareto front results for two more evaluation setups: Evaluation 2 and 3 (details in the paper Table 1). We present the results in the form of tables here. We are committed to adding them for all the evaluations in the final version of our paper. - ***For Evaluation 2 Setup*** (detailed in Paper Table 1): This table shows the value of KL and the corresponding win rate and shows that our proposed method outperforms the existing methods. | Method | | | | | | | | | | |------------------------------|----------|-------|-------|-------|-------|-------|-------|-------|-------| | ARGS-DPO | KL | 0.40 | 1.15 | 2.20 | 3.80 | 5.75 | 7.05 | 8.20 | 9.15 | | ARGS-DPO | Win-Rate | 50.50 | 57.80 | 61.75 | 65.90 | 67.10 | 67.70 | 68.20 | 68.15 | | $\text{CD}^{--}$ | KL | 0.50 | 1.25 | 2.35 | 4.20 | 6.50 | 8.75 | 9.35 | 10.75 | | $\text{CD}^{--}$ | Win-Rate | 50.75 | 62.85 | 68.90 | 72.70 | 75.40 | 76.00 | 76.15 | 76.30 | | $\texttt{TQ}^{\star}$ (Ours) | KL | 0.42 | 1.20 | 2.18 | 3.85 | 5.95 | 7.90 | 8.85 | 10.40 | | $\texttt{TQ}^{\star}$ (Ours) | Win-Rate | **54.30** | **70.90** | **75.70** |**79.60** | **80.55** | **81.95** | **82.95** | **83.25** | - ***For Evaluation 4 Setup*** (detailed in Paper Table 1): This table shows the value of KL and the corresponding win rate and shows that our proposed method outperforms the existing methods. | Method | | | | | | | | | | |------------------------------|----------|-------|-------|-------|-------|-------|-------|-------|-------| | ARGS-DPO | KL | 0.37 | 1.26 | 2.05 | 3.71 | 5.86 | 7.14 | 8.36 | 9.23 | | ARGS-DPO | Win-Rate | 50.10 | 58.32 | 62.10 | 66.13 | 67.32 | 67.89 | 67.41 | 66.02 | | $\text{CD}^{--}$ | KL | 0.45 | 1.32 | 2.39 | 4.36 | 6.58 | 8.85 | 9.50 | 10.89 | | $\text{CD}^{--}$ | Win-Rate | **51.05** | 63.12 | 69.44 | 73.16 | 75.80 | 76.25 | 77.00 | 77.17 | | $\texttt{TQ}^{\star}$ (Ours) | KL | 0.38 | 1.27 | 2.11 | 3.95 | 6.05 | 8.03 | 8.98 | 10.58 | | $\texttt{TQ}^{\star}$ (Ours) | Win-Rate | 50.86 | **69.45**| **73.28** | **76.19** | **79.20** | **80.19** | **81.00** | **82.16**| We have thoroughly addressed all concerns and are more than happy to engage in further discussions if any additional issues remain. Thank you so much again for your consideration.