Accelerated Preference Optimization for Large Language Model Alignment

Thanks for your valuable feedback!

Q1: The theoretical results in this paper are not particularly strong and do not conclusively show that APO is superior to previous methods. Specifically, in Theorem 4.4, compared to DPO (α=0), APO achieves a smaller optimization error by a factor of 1−α, but incurs a larger statistical error by a factor of 1/(1−α). Therefore, in the finite-data case, it remains unclear whether APO is theoretically better than DPO. Similar concerns also arise in Theorem 4.8.

A1: Thank you for highlighting this trade-off. We agree that adding momentum increases statistical error. However, as noted in Remark 4.5, when working with sufficiently large datasets, which is widely studied in the analysis of preference-based optimization method, the optimization error becomes the dominant term in determining the sub-optimality gap. In these practical cases, our algorithm's ability to reduce this dominant term by a factor of $(1−α)$ represents a meaningful improvement. This improvement is particularly significant because in many real applications (e.g., $N=2K$ samples and $T=O(1)$ iterations in our experiments), the reduction in optimization error outweighs the increase in statistical error. Our theoretical analysis thus provides valuable insights into how momentum can be effectively leveraged to enhance convergence in preference-based optimization, especially when working with substantial amounts of preference data. This is also backed up by our empirical experiment results.

Q2: There is a mismatch between the theoretical results and the experiments in this paper. Theoretically, the paper proves that APO achieves a convergence rate of the order O~((1−α)/t). This rate is meaningful only in the large-iteration regime, as it ignores certain constants. However, in the experiments, the LLM is trained with only a small t=3. This discrepancy creates a gap between the theoretical analysis and the empirical results. As a result, I believe that the improved convergence rate in theory does not fully explain APO's empirical performance, and the experimental results do not adequately reflect its theoretical convergence rate.

A2: We appreciate this insightful observation regarding the gap between our theoretical analysis and the current empirical setup. Increasing the number of iterations $T$ improves performance but also leads to higher computational time and cost, requiring a careful balance between the benefits and the additional overhead. To address this, we are conducting further experiments with more iterations to better validate the theoretical convergence rate and will update our response with the results as soon as they are available.

Q3: In Equation (3.1), the loss function for reward models is denoted as $ℓ(r,x,y^w,y^l,π_t)$. However, it appears that this function does not actually take πt as an input.

A3: In our work, we consider the general form of loss functions. While DPO loss does not depend on $π_t$, other loss functions like SPPO (Example 3.2) do utilize $π_t$ in their formulation. Therefore, we maintain $π_t$ in the general loss function notation to accommodate these cases. This generalized notation allows our framework to encompass a broader class of preference-based optimization methods, including those where the loss function explicitly depends on the current policy. By maintaining this general form, our theoretical analysis and algorithmic framework can be applied to various preference-based optimization scenarios, providing flexibility in the choice of loss functions while maintaining theoretical guarantees.

Q4 There is a trade-off between optimization error and statistical error in the developed theory. Therefore, it is unclear whether Nesterov's acceleration offers a true advantage over the naive method.

A4 Thank you for highlighting this trade-off. We agree that adding momentum increases statistical error. However, as noted in Remark 4.5, when working with sufficiently large datasets, which is widely studied in the analysis of preference-based optimization method, the optimization error becomes the dominant term in determining the sub-optimality gap. In these practical cases, our algorithm's ability to reduce this dominant term by a factor of $(1−α)$ represents a meaningful improvement. This improvement is particularly significant because in many real applications (e.g., $N=2K$ samples and $T=O(1)$ iterations in our experiments), the reduction in optimization error outweighs the increase in statistical error. Our theoretical analysis thus provides valuable insights into how momentum can be effectively leveraged to enhance convergence in preference-based optimization, especially when working with substantial amounts of preference data. This is also backed up by our empirical experiment results.

Q5 To the reviewer’s understanding, new preference data is required in each iteration t=0,1,…,T. Is this correct? Could the authors clarify the preference collection process? Is the preference collected using the BT model with random noise?

A5 Yes, for each iteration, we collect new preference data with the reference policy $π_t$. The details of the preference collection process are discussed in section 5.2. Specifically, in each iteration, we sample five pairs of responses under the current policy πt and rank them using their PairRM score. The top-ranked response is designated as the winner while the bottom-ranked response is designated as the loser, which is not related to the BT model.

Q6 Related to this, in the experiments, a reward model is used. Why not study direct reward maximization (e.g., PPO) both empirically and theoretically? To the reviewer's knowledge, direct reward maximization methods can handle such tasks effectively and demonstrate strong empirical evidence in practice.

A6 First, we want to emphasize that in our work, we only use the reward model to behave as an oracle for the preference data, and our APO algorithm operates solely on the collected preference data without accessing the underlying reward values. This preference-based setting is more aligned with human feedback scenarios where explicit reward values may not be available, and is particularly relevant in real-world applications where feedback often comes in the form of comparisons rather than absolute scores. Therefore, our work is not directly comparable with direct reward optimization methods, as we specifically address the challenges and improvement unique to preference-based learning in language model alignment.

Q7 How should (3.6) be implemented in practice? Could the authors provide more details? Should the parameter or the distribution be extrapolated? Additionally, how can this result be extended to cases where y is sequential? Is it straightforward to apply this closed-form solution?

A7 As we discussed in Section 5.2, when the policy is a softmax linear function, update (3.6) reduces to a momentum extrapolation in the parameter space. Therefore, we implement this momentum step by performing the extrapolation directly in the parameter space to approximate the corresponding momentum update (3.6) that is defined in the probability space. In addition, in our implementation, our output response y is already sequential and our method naturally handles this case through the same parameter space extrapolation.

Thank you for your careful review and detailed feedback!

Q1: Assumption 4.3 is strong. Essentially, it requires $π_t$ to be a purely stochastic policy. On the other hand, the optimal policy $π^*$ in Theorem 4.4 is deterministic. Therefore, the theory with this assumption seems problematic in the asymptotic case.

A1: It is very common to use stochastic policy in RL algorithms even though there always exists a deterministic optimal policy. For example, in one of the most widely used RLHF algorithms, PPO, the output policy is also stochastic. From the theoretical perspective, we want to clarify that our realizability does not require $π^*$ to fall into the function class, and the stochastic policy cannot fully match the deterministic optimal policy $π^*$ . Nevertheless, according to our Theorem 4.8, we have shown that the TV-distance between the stochastic policy $\hat{\pi}_{t+1}$ and the optimal policy $\pi^*(x)=\arg \max _{y\in Y} r^*(x,y)$ will converge to 0 after enough iterations ($T$) and samples ($N$), and this holds exactly in the asymptotic sense. In addition, as we stated before Assumption 4.3, similar assumptions have also been used in the previous analysis of preference-based optimization method [1]. We believe this is a standard assumption and not strong at all. We hope the explanation above helps to clarify any misunderstandings from the reviewer.

[1] Direct nash optimization: Teaching language models to self-improve with general preferences.

Q2: The implications of the theory are based on upper bounds and are weak. Could you provide simulations for a simple case to validate the tightness of the developed method? Additionally, could you use simulations to study the trade-off between optimization sub-optimality and statistical error?

A2: Using upper bounds to measure the efficiency is common in theoretical analysis, as it provides the guarantee even in the worst case and helps us understand the fundamental limitations of the algorithm.

In addition, for our large language model alignment work with reparameterized reward functions, conducting appropriate simulations to validate these theoretical bounds and study the trade-offs presents significant challenges. To our knowledge, there exists no simulation setting for aligning LLMs. We therefore rely on our empirical results to demonstrate the method's effectiveness and shows an improvement upon the baselines including DPO and KTO, which supports our theoretical analysis about the benefits of incorporating momentum in preference optimization.

Q3: It is recommended to explain in the main text why momentum increases the statistical error.

A3: Momentum improves optimization by accumulating historical gradient information in a velocity term, which helps navigate challenging loss landscapes. Specifically, in flat regions like valleys or plateaus where gradients are small, standard gradient descent makes slow progress since updates are directly tied to gradient magnitude. Momentum helps by "remembering" consistent directional trends across iterations, allowing continued progress even with small individual gradients. This memory effect helps reach optima faster.

However, this accumulation mechanism that provides benefits also amplifies any errors in gradient estimation. When we incorporate previous gradients into our update step through momentum, any statistical noise or errors in these gradient estimates get propagated and potentially amplified over time. This amplification of estimation errors is what leads to the increased statistical error term we observe. However, when working with large datasets, the optimization error becomes the dominant term in determining performance and the reduction in optimization error outweighs the increase in statistical error.

Thanks for your helpful suggestions and comments!

Q1 To the best of my understanding, Equation 3.6 in the APO algorithm does not update the neural network weights. Instead it is derived by keeping reference of the policy from the previous and the current iterations. It is not clear how the final policy in step 7 is used in the next iterations in steps 5 and 6.

A1 First, we want to clarify that in implementation, we perform updates on parameters rather than keeping reference of the policy from the previous and current iterations. As we discussed in Section 5.2, when the policy is a softmax linear function, this update reduces to a momentum extrapolation in the parameter space. Therefore, we implement this momentum step by performing the extrapolation directly in the neural network parameter space to approximate the corresponding momentum update (3.6) that is defined in the probability space. Second, for the next iteration, the final policy in step 7 plays two roles: it serves as the reference policy for collecting new preference data and also sets the reparameterized function class for optimization.

Q2 This work might benefit from additional baselines in the empirical studies - IPO, SPPO, SimPo, KTO and especially ExPO

A2 We have conducted comprehensive experiments with several baselines. Here are the detailed results:

Compared to our results in Table 1, these results highlight the superior performance of our method over these baselines.