Transfer Q-star : Principled Decoding for LLM Alignment

Thank you for your positive feedback on our paper.

Weakness 1: Equation 5.... introduced.

Response to Weakness 1: We will add the derivation for the closed-form solution in Equation 4 (also shown in CD [B]) as a Lemma in the appendix for better clarity.

Weakness 2: The introduction of CD..... reverse.

Response to Weakness 2: We thank the reviewer for this point but there seems to be a confusion. We apologize for any oversight during our explanation of CD in our paper. We take this opportunity to expand on this point more below.

We note that the objective of Controlled Decoding which is a KL regularized RL problem (also highlighted in CD) is given by

$

\pi^*(\cdot|s\_t) := \arg \max\_{\pi} \mathbb{E}\_{z \sim \pi(\cdot|s\_t)}[Q^*(s\_t, z)] - \alpha KL (\pi(\cdot|s\_t), \pi\_{\text{SFT}}(\cdot|s_t)) \tag{4}.

$

However, it is extremely crucial to note that in order to obtain the optimal policy in (4), the Q-function in the equation should be $Q^*(s\_t, z)$ (the optimal action-value function i.e the reward return calculated with optimal policy) and not any arbitrary $Q$ function from any policy. However, the approach in CD indicates that the $Q$ estimation (either directly or through an external network) is done through the data generated using $\pi\_{sft}$ policy (refer to Equation 1, Equation 4 in [B]), leading to sub-optimality as demonstrated in our Figure 1. Our main contribution is to show that with an access to an aligned model, we can estimate $Q^{*}(s_t, z)$ in a much more efficient way than CD as demonstrated in all our experimental settings.

Possible source of confusion: However, we believe the source of confusion is that in the CD paper, the definition of action value function used in (4) above (which is equation (1) in CD paper [B]) utilized $\pi_{sft}$ to sample but denotes the value function with $V^*$ which is usually reserved for the optimal value function.

Weakness 3: Equation 12 on line 175 .... the reward model.

Response to Weakness 3: We thank the reviewer for providing us with this opportunity to clarify further.

Solution in (12) is theoretically correct and justified. We have stated our proposed decoding optimization problem in Equation (11), and solving that leads to the choice of having the reference policy as $\pi_{\text{BL}}$ in Equation 12 on line 175. This is a specific choice of our algorithmic solution design which we have also implemented in our work. However, we want to emphasize that our theoretical results in Sec. 3.3 (specifically the KL divergence upper bound in Theorem 1, statement (2)) derives the divergence of proposed algorithm's policy to the original reference policy $\rho_{\text{SFT}}$. The upper bound of Theorem 1 [statement 2] is given by

$

\mathbb{D}_{\text{KL}}(\rho^*_{Alg}(\cdot|\mathbf{x}),\rho_{SFT}(\cdot|\mathbf{x}) )\leq(\frac{1}{\beta}+\frac{1}{\alpha}T)r_{\max}.

$

By controlling the value of $\alpha$ and $\beta$, we can control the distance to the original reference model. Moreover, in our experimental ablation Figure 2(b) we show the KL divergence to the original SFT reference model which proves that the proposed method does comparable to other baselines and is not diverging from the SFT policy.

Weakness 4: A crucial experiment on t...... in weakness 3.

Response to Weakness 4: As suggested by the reviewer, we performed additional experiments to generate the pareto-frontier results. To be specific, in Figure-3 in the rebuttal PDF, we compare the tradeoff between the win-rate and the KL divergence to the base reference SFT policy. TQ* outperforms existing baselines. We will add this to the final version.

Question 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_{\text{SFT}}$

Response to Question 1: Yes, you are right. We have indeed used the reference model $\pi_{\text{SFT}}$ for all the algorithms while measuring the KL divergence for TQ*.

Question 2: I'm curious about the inference cost. How do the memory requirements and time complexity compare to naive greedy decoding?

Response to Question 2: We report the inference time-complexity of all the existing deocding algorithms. Ours is comparable to state-of-the-art decoding methods.

Reference:

[B] Sidharth Mudgal et al., Controlled decoding from language models, 2024

Comment 1.1: The author still lacks a theoretical foundation. Whether it's a Q-function or a Value function, they end up having the same representation in this algorithm due to the (both tractable or intractable) normalizer.

Response: Thank you for your comment. We apologize if we missed anything but this comment is not very clear for us to respond to. We have followed the standard definitions from the reinforcement learning literature [A]. We start by explicity defining the token level MDP in Section 2.1. We emphasize that the Value function is defined for each state (defined in line 210 in our paper), and the Q function is for each state and action (defined in Equation (2) in our paper).

Request to Reviewer: We kindly request the reviewer to please expand on why value and Q function will be the same and what is the meaning of the term "normalizer" in this context. We want to clearly understand the concern before responding. Thank you for your feedback and engagement.

Our Focus and Contributions: In our work, we want to emphasize that our focus is on the estimation of optimal Q* and show that it is better as compared to all existing decoding methods such as CD, ARGS, etc. We show this in theory as well as experiments.

[A] Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.

Comment 1.2: Furthermore, although a theoretical guarantee has been presented regarding the excessively diverging KL-div, sufficient experiments and defense have not been conducted. Therefore, I maintain my current assessment.

Response: Thank you for your comment. We want to highlight that our work contributes on both theory and experiments side.

On the theory side, our work is the first to derive such a theoretical upper bound for both (1) suboptimality and (2) KL divergence for a decoding algorithm. There are no theoretical results in any of the existing works (such as CD, ARGS etc), which constitutes a unique and novel contribution of our work on its own.

On the experimental side, we have tested our proposed approach on six evaluations (in the main submission) and added two more large-scale evaluations in the rebuttal pdf [link to pdf in openreview] .

• For the KL divergence plot, we have Figure 2(b) in the main body and as the reviewer suggested, we added a Pareto front plot for evaluation 1 (Figure 3 in the rebuttal pdf) as well in the rebuttal pdf [link to pdf in openreview] which clearly shows the superior performance of our method.
• Additionally, our current comparison includes win rate, coherence, reward, and diversity, which are designed to approximate human preferences.

We are running more experiments and committed to adding them in the final version of our work. We believe we have addressed all the concerns, and are happy to engage in further discussions if any remain.

Thank you once again for your time and consideration. Looking forward to your feedback.

Thank you for your comments.

Weakness 1: Morally speaking, it is not ...... empirically, though.

Response to Weakness 1: We believe the reviewer wants to understand the advantage of directly generating the response using $\rho_{\text{BL}}$ over TQ*. We explain in detail below.

Our proposed method enables optimal action at the Token Level. In our algorithm (TQ*), for every state ($s_t$) we compute the action value function $Q^*(s_t,z)$ using $\rho_{\text{BL}}$ and select the optimal action by maximizing the action value function $Q^*(s_t,z)$, ensuring the optimal performance for any token as highlighted in Theorem 1. Intuitively, it ensures the possibility of performing better credit assignment with token level MDP while decoding, when compared to $\rho_{\text{BL}}$ which is trained as a contextual bandit for trajectory responses. This aspect has also been highlighted in the very recent work [C]. As a result, the majority of the prior decoding methods demonstrate improvement over directly using the DPO policy

Theoretical Justification: From Theorem 1 in the paper, we can extract the theoretical justification of the improvement of our proposed algorithm over $\rho\_{\text{BL}}$. The sub-optimality gap of our algorithm, defined in Equation 18, can be decomposed into two terms ($\Delta \leq \Delta_1 - \Delta_2$) as detailed in Appendix G.1 (Equations 28 and 29).

The first term, $\Delta_1$, is a function of parameter $\beta$ and represents the sub-optimality from the optimal value function due to $\rho\_{\text{BL}}$, with its upper bound $\Delta\_1\leq \beta\mathbb{D}_{\text{KL}}[\rho^*(\cdot|\mathbf{x}) \mid \mid \rho\_{\text{sft}}(\cdot|\mathbf{x})]$. This indicates the sub-optimality incurred at the token level when using $\rho\_{\text{BL}}$ alone.

However, the second term, $\Delta_2 = \alpha h_{\alpha}$ , is always non-negative, illustrating that our algorithm provides an additional benefit of improving the suboptimality gap. By tuning the value of $\alpha$, we can improve the performance beyond what is possible with $\rho_{\text{BL}}$ alone. This improvement is consistently demonstrated across all our experimental results, as also highlighted by the reviewer.

Weakness 2: For the indirect case, .... is useful.

Response to Weakness 2: We believe what the reviewer is asking is if we have access to the ground truth reward, why can't we directly have an aligned $\rho_{\text{BL}}$ for that ground truth reward? We expand on this in detail as follows.

Obtaining aligned $\rho_{\text{BL}}$ to ground truth would require fine tuning. We note that even if we have access to ground truth reward, obtaining an aligned model $\rho_{\text{BL}}$ using standard DPO would require fine-tuning of parameters which is not the focus of this work. Our motivation is in tuning-free alignment settings, where we don't update the parameters of the language model.

Regarding our approach and decoding methods. We emphasize that any decoding method would require access to the target reward function (trajectory level). However, for indirect transfer, we leverage any open-sourced baseline model $\rho_{\text{BL}}$ aligned to an arbitrary reward function $r_{\text{BL}}$ which is different from the target reward. Thus, for principled decoding even though we have $r_{\text{target}}$, still we can't directly use $\rho_{\text{BL}}$ to estimate action value function and decode, which will result in sub-optimal decoding due to the distribution shift. Thus, we need to estimate $\rho_{\text{target}}$ using Equation 16.

Weakness 3: Finally, I find it a bit of a ....... So it is strong to assume you have both of these.

Response to Weakness 3: We remark that any decoding method requires the access to target reward function for alignment. As reviewer mentioned, if the target reward function is learned through preferences, there will be statistical error due to coverage and suboptimality of optimization methods. Furthermore, the DPO policy could also have it's own instability issues which might affect the final performance of our decoding method. We agree with the reviewer and will highlight this limitation in the revised final version of our work.

Examples of Access to Ground truth reward: We note that there can be true rewards/scores that don't need to be always trained from data, for example - coding, and mathematical tasks where one might have fixed ground truth avoiding above mentioned statistical errors.

Question 1: equation (2) - This seems like a typo ....sequence z_0 ... z_{i-1}.

Response to Question 1: Thanks for pointing out the typo, we will fix it in the final version of the paper. The first argument should be $s_{t+i}$ to make sure it take in so far generated sequence.

Question 2: equation (15) - how is $Z_{\text{BL}}$ and $Z_r$ computed.... intractable.

Response to Question 2: This is a good question. We note that for the theoretical analysis in this work, we assume access to the ratio mentioned in Equation (15). For the experimental purposes, we utilized unbiased estimates. For instance, we can estimate the partition from collecting samples from $\rho_{\text{SFT}}$ and evaluating the empirical estimate of $\mathbb{E}\_{y \sim \rho\_{\text{SFT}}}[\exp(\frac{1}{\beta}r(x,y))]$. Regarding this ratio explicitly, we had an interesting observation that for several realistic scenarios of transfer, as observed in our empirical experiments, when either the reward difference $r_1(x,y) - r_{\text{BL}}(x,y)$ is small or the reference policy is not aligned to any specific reward $r(x,y)$, the ratio of the partition functions will be 1. This would effectively relax the computational bottleneck for the implementations.

[B] Sidharth Mudgal et al., Controlled decoding from language models, 2024 [C]. Rafailov, R et al., From r to q*: Your language model is secretly a q-function. arXiv preprint arXiv:2404.12358.

Response to Reviewer Summary: We sincerely thank the reviewer for the thoughtful review of our work, highlighting the strength and novelty in both our formulation and experimental design in leveraging available aligned models for principled decoding. We address all other comments one by one as follows.

Weakness 1: The proposed approach seems computationally expensive to do decoding.

Response to Weakness 1: We thank the reviewer for this important point. We agree with the reviewer that performing the proposed principled decoding would indeed add an additional computational overhead at the inference time. But we would like to emphasize that our main focus and contribution in this work is to provide principled method to perform LLM alignment via decoding and provide theoretical gaurantees.

Moreover, our decoding method has a comparable inference time to the existing state-of-the-art Controlled Decoding (CD) [A]. For the practical implementation, similar to CD method in [A], one can train a small Q-function adaptor offline, which would allows for faster inference time. This significantly reduces time complexity, similar to ARGS [B], which only introduces a constant factor of k (top-k tokens) over classical decoding methods.

Another important point to highlight is that there are several critical scenarios—such as reasoning, mathematical puzzles, chess, and coding where accuracy is more crucial than inference time. In these situations, principled decoding can provide substantial benefits.

Weakness 1: The approach require accessing an already aligned model, this adds more requirements for using the method and limits it use cases.

Response to Weakness 2: Thank you for raising this critical point. We remark that our proposed decoding method does not require access to a model specifically aligned to a particular reward function. Instead, interstingly, we can leverage any existing aligned model (BL), even if it is aligned to a different reward (Indirect transfer) and thus is not inherently restrictive. Therefore, utilizing already existing aligned model (such as on Hugging face) for improved decoding is the strength of our approach. We show this via extensive experiments in our paper.

Question 1 : Did the authors compare decoding speed for different approaches? It would be useful to report even thought the proposed approach is slower.

Response to Question 1: Thank you for pointing this out. We have now included a comparison of the inference times for all baseline decoding algorithms in the table below. Specifically, we report the inference time for decoding a single prompt, based on the system configuration detailed in Appendix C.

It is evident that $`TQ`^{\star}$ perform comparably with CD in terms of the time complexity, whereas it is slower than ARGS which suffers from low reward. However, with an offline trained adpator we can scale the inference time to be same as ARGS.

Question 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.

Response to Question 2:
We follow the same evaluation protocal as followed in prior decoding approaches like ARGS [B], thereby reporting the Win-tie in Table 2 for uniformity in comparison. However, as mentioned by the reviewer, we also provide the win-rate and tie-rate separately for Evaluation-1 and Evaluation-3 Setup below.

Table: Win, Tie and Lose-rate for Evaluation Setup-1

Table: Win, Tie, and Lose-rate for Evaluation Setup-2

[A]. Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, Jilin Chen, Alex Beutel, and Ahmad Beirami. Controlled decoding from language models, 2024

[B]. Maxim Khanov, Jirayu Burapacheep, and Yixuan Li. Args: Alignment as reward-guided search,2024.

Response to Reviewer Summary: We thank the reviewer for the encouraging remarks and recommending acceptance of our work. We provided detailed responses to other comments one by one as follows.

Weakness 1: Comparison with Baselines: The comparisons with existing baselines like DPO are insightful, but additional baselines, especially those focusing on inference-time control, could strengthen the evaluation.

Response to Weakness 1: Thank you for your comment. We have now included a comparison with additional inference time control baselines (aka decoding methods) as well in the table below. Specifically, we report the time taken to decode a single prompt, using the system configuration described in Appendix C in the main paper.

Weakness 2: 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.

Response to Weakness 2: We agree with the reviewer that doing more ablations with respect to hyperparameters would strengthen the work. To this end, as suggested by the reviewer, we did additional ablations experiments to understand the effect of varying the decoding alignment parameter $\alpha$ on the quality of the generated text. Specifically, we performed decoding by varying $\alpha$ such that $\frac{1}{\alpha} \in [0.1, 0.25, 0.5, 0.75, 1, 2, 5, 7.5]$. We report the results in Figure-3 in the rebuttal PDF. To be specific, we compare the tradeoff between the win-rate and the KL divergence to the base reference SFT policy for different values of decoding alignment parameters.

Weakness 3: 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.

Response to Weakness 3: Thank you for raising this concern. As the reviewer suggested, we performed an evaluation on two additional setups using larger models, as detailed in table below.

We report the results in Figure-2 in the rebuttal PDF.

Weakness 4: 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.

Response to Weakness 4: Our current comparison includes win rate, reward model performance, coherence, and diversity, which are designed to approximate human preferences. However, we agree that incorporating direct human evaluation would provide more robust validation of our method's effectiveness. We plan to include human evaluation in the final version and are currently in the process of obtaining the necessary permissions.

Question 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?

Response to Question 1: Thank you for your suggestion. We performed additional evaluation by varying the length of generated text. We report the results in Figure-1 in the rebuttal PDF. We observed that irrespective of the length of generated text, $`TQ`^{\star}$ consistently outperforms all the compared baselines.

Question 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?

Response to Question 2: This is a great question and is a valid scope of future research work. We can try by selecting argmax Q not only over the action token by as well as models itself. We are working as a future research.

Limitations: The authors acknowledge ..... impact real-time applications.

Response to Limitations: Thank you for your insightful comments. We will expand our discussion of limitations in the final version to include: (1) The reliance on existing aligned baseline models, noting their availability and suitability challenges for different target rewards. (2) Potential errors or biases in the GPT -4-based evaluation, addressing its limitations as a proxy for human assessment. (3) The computational overhead of TQ⋆ compared to standard decoding methods, considering its impact on real-time applications.

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.

• 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.

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.