Paper: Fine-Tuning Language Models with Advantage-Induced Policy Alignment (ICLR-2024)

Reviewer comment (Theory):
The authors have pointed out discrepancies between the original target and the final loss of AWR as well as the changing of π_old in the online case as motivating factors for APA. However, it appears that APA still faces these issues.

Author response:
Thank you for your comments. We would like to clarify first that APA does not face the same issue introduced by changing $\pi_{old}$ as AWR. Intuitively, this is because minimizing $\pi_{old} \log(\pi / \pi_{init})$ (actual implementation of AWR) is different from minimizing $\pi_{init} \log(\pi / \pi_{init})$ (the ideal target of AWR), while minimizing  $\pi_{old} (\pi - \pi_{init})^2$ (actual implementation of APA) is similar to minimizing $\pi_{init} (\pi - \pi_{init})^2$ (the ideal target of APA) as long as the policy $\pi_{old}$ and $\pi_{init}$ covers enough state-action pairs. This benefit is naturally introduced by replacing KL with squared loss, and is our main motivation for designing APA. We have also provided theoretical guarantees to formalize this intuition.

Reviewer comment (Theory):
The authors claim to provide a comparison among APA, PPO, and AWR theoretically in the introduction, but only the upper bound of APA is presented in Theorem 1. This leaves readers unable to derive the theoretical advantages of the proposed method.

Author response:
Thank you for the comment. We discuss in Appendix A why Z(s) can be approximated as 1, which is validated in experiments. This assumption is also used implicitly in the derivation of AWR. 

In terms of theoretical guarantees for PPO and AWR, it can be hard to provide bounds due to their complicated nature, especially for PPO which relies on estimating both value, advantage and updating the model. For AWR, we show in Equation (7) that it converges to the wrong target that is inconsistent with Equation (3). Thus theoretically, we know that AWR is not convergent while APA is convergent.

Reviewer comment (Theory):
The assumption for explaining why Z(s) approximates to 1 is deemed overly strong.

Author response:
Thank you for the comment. We discuss in Appendix A why Z(s) can be approximated as 1, which is validated in experiments. This assumption is also used implicitly in the derivation of AWR. 

In terms of theoretical guarantees for PPO and AWR, it can be hard to provide bounds due to their complicated nature, especially for PPO which relies on estimating both value, advantage and updating the model. For AWR, we show in Equation (7) that it converges to the wrong target that is inconsistent with Equation (3). Thus theoretically, we know that AWR is not convergent while APA is convergent.

Reviewer comment (Experiments):
Recent research [2] has highlighted the significance of considering various divergences, each offering its unique balance between alignment and diversity trade-offs. Given that the core contribution of this paper revolves around the introduction of a novel f-divergence for distribution matching, it is imperative that the authors conduct a comparative analysis by pitting the new f-divergence against commonly used f-divergences, such as reverse KL, total variation (TV), and Jensen-Shannon (JS) divergences.

Author response:
Thank you for your comments. We choose to mainly focus on KL in Equation (3) since that is a traditional divergence where both APA and PPO rely on. We are happy to include more comparisons on different f-divergences in the final version of the paper.

In Figure 2, we control the KL divergence of APA and PPO to be on the same level to ensure a fair comparison. In this case, we can observe that APA is more KL-efficient and sample-efficient than PPO.

Reviewer comment (Experiments):
Figure 2 suggests that PPO may not have converged due to a large KL penalty.

Author response:
Thank you for your comments. We choose to mainly focus on KL in Equation (3) since that is a traditional divergence where both APA and PPO rely on. We are happy to include more comparisons on different f-divergences in the final version of the paper.

In Figure 2, we control the KL divergence of APA and PPO to be on the same level to ensure a fair comparison. In this case, we can observe that APA is more KL-efficient and sample-efficient than PPO.

Reviewer comment (Experiments):
Recent research [3] argues that the KL penalty can be removed, and the authors themselves find that PPO without KL penalty converges to a higher reward compared to APA, as shown in Figure 7. Does this imply that PPO outperforms APA when using smaller or no KL penalty terms? It is important to note that the reviewer does not believe that the higher cost of KL divergence is the issue.

Author response:
Thank you for the comment. In fact, the results in [3] suggest that removing KL penalty will not affect the “gold reward-KL frontier”. This suggests that no matter whether we have KL regularization, if KL is much larger, the gold reward will still be worse than the case of small KL. In Appendix, we find that PPO achieves a higher reward than APA, but at the cost of a much higher KL divergence. According to [3], this means that the PPO may have lower gold reward even though the proxy reward is higher. 

[3] Gao L, Schulman J, Hilton J. "Scaling laws for reward model overoptimization." International Conference on Machine Learning. PMLR, 2023: 10835-10866.

Reviewer comment (Experiments):
It seems that the modification also applies to the DRL scenarios beyond the LLMs. But as far as I know, the prominent approach in research and industry is still the KL-based PPO. I am wondering whether this is due to the special case of LLMs or this modification could also outperform PPO in classical applications.

Author response:
Thank you for your comments. We have included in Appendix D.5 on the experiments in 4 traditional RL tasks in Mujoco, where we observe that APA outperforms PPO for 3 of them. This shows the potential of APA being applied to traditional RL tasks as well.

Reviewer comment (Experiments):
Why do you choose to finetune only the last two layers (on the HH dataset in appendix C)? It seems that with the used GPT-J-6B reward model, we might get much higher reward with PPO [2] (although it does depends on the implementation and starting checkpoint). I am not sure whether this is because you only choose to finetune the last two layers but it seems that in practice, we tends to use low-rank adapter instead of freezing most of the layers.

Author response:
Thank you for your comments. In fact, we use a low-rank adapter for our stack-exchange dataset fine-tuning to ensure a fair comparison with both LoRA and frozen-layer-based methods. The reward with PPO highly depends on the chosen evaluation dataset.

Reviewer comment (Experiments):
In figure 6, the PPO does not work normally as the reward seems to be decrease from the beginning. I understand that PPO often drops suddenly but this could be mitigated by using larger KL penalty and also smaller learning rate and less epochs for each iteration. I think the PPO is not tuned to the best performance so far and using the default parameters in trlx is not enough.

Author response:
Thank you for your comments. The beginning drop of PPO can be normal due to the inaccurate estimation of value functions. After the value function converges or stabilizes, the performance of PPO will begin to increase. We are happy to further fine-tune the performance of PPO in the final version. Althoughwe adopt the default hyperparameters for PPO, we indeed have conducted extensive hyperparameter search for all the three algorithms. In Appendix D, we list the most significant hyperparameter that changes the result, which is the KL control coefficient. We provide results for APA when the KL coefficient is set to be $0.1, 1, 10$, and the results for PPO when the initial KL coefficient is $0$ or $0.05$.  It turns out that if we run PPO without KL control, the reward will always increase, but at a cost of much higher KL divergence. But once we increase the KL coefficient, we observe some decrease in the reward, potentially due to both stringent KL control and the instability of the PPO algorithms. Besides this, we do find that the default hyperparameters for PPO give the most robust experimental results for our settings. However, we hope to mention that one motivation of proposing APA is that there are so many hyperparameters to tune in PPO, while APA only requires the KL coefficient, which is another huge benefit.

Reviewer comment (Writing):
A minor typo: I think the state space of the LLM should be the product of the token spaces instead of the union.

Author response:
Thank you for your suggestion! In fact, the state space is defined as the union of the product of the token spaces with variable number of tokens. So our first union is union over all possible lengths in the token.

Reviewer comment (Theory):
In the main text, there is a reference to a connection to soft-q learning but it seems like only the f-divergence interpretation is discussed.

Author response:
We have included a new Appendix C discussing an alternative loss inspired by soft-Q learning. In fact, the soft-Q learning loss is our first attempt before APA. Unfortunately, it does not give as good performance as APA or PPO. So we only provide intuitions and discussions in the Appendix. We have also updated the main text to connect to the Appendix better.

Reviewer comment (Experiments):
For the 125M parameter experiments, it seems like PPO is not properly tuned to optimize for the reward.

Author response:
We have optimized the PPO performance as hard as we can. In fact, it is interesting to see that with a larger language model (3B or 6B), the performance of PPO seems to be more stabilized than that of 125M parameters. We have also conducted ablation study in Appendix D for the effect of different hyperparameters for PPO and APA algorithm.

Reviewer comment (Experiments):
How important was it to estimate a good critic? Was there anything else that was needed to learn a better A rather than having a separate model? This seems to be an interesting angle to experiment some more since the provable guarantees also relies on having a good advantage function for $\pi_{old}$.

Author response:
It is very important to learn a good critic / advantage function, which is required for both PPO and APA. In fact, during our experiment, the value function converges very fast, which is a good sign that the advantage might be estimated very well due to the same advantage estimation technique used in PPO.

Reviewer comment (Experiments):
The major reason I tend to reject is the scope of evaluate tasks. As RLHF (with PPO) has been well verified on ChatGPT, which has great generalization ability, PPO can be utilized to train language models in scale. In addition to training language models, PPO has stand the test of time in many other area such as robotics control. I believe that is why researchers use PPO to fine-tune many large language models. However, the evaluation tasks in this paper locate in a specific domain. I think it is not enough to be an alternative to PPO.

Author response:
Thank you for your comments. We have included in Appendix D.5 of the updated draft on the experiments in 4 traditional RL tasks in Mujoco, where we observe that APA outperforms PPO for 3 of them. This shows the potential of APA being applied to traditional RL tasks as well.

Reviewer comment (Experiments):
The paper does not provide much discussion on the choice of the hyperparameter $\lambda$, which controls the trade-off between the expected advantage and the KL divergence. Besides, as this paper proposes an optimization algorithm, the ablation study is required.

Author response:
Thank you for your comments. We have ablation study included in Appendix D. In Appendix D, we list the most significant hyperparameter that changes the result, which is the KL control coefficient. We provide results for APA when the KL coefficient is set to be $0.1, 1, 10$, and the results for PPO when the initial KL coefficient is $0$ or $0.05$.  It turns out that if we run PPO without KL control, the reward will always increase, but at a cost of much higher KL divergence. But once we increase the KL coefficient, we observe some decrease in the reward, potentially due to both stringent KL control and the instability of the PPO algorithms. Besides this, we do find that the default hyperparameters for PPO give the most robust experimental results for our settings. However, we hope to mention that one motivation of proposing APA is that there are so many hyperparameters to tune in PPO, while APA only requires the KL coefficient, which is another huge benefit.