Model Extrapolation Expedites Alignment

Line 225-229: The authors are trying to show that |αδθ| is small. They argue that α decreases through training, then conclude that |αδθ| is small. But they do not discuss the fact that δθ is (presumably) increasing during training.

We have already discussed this fact in lines 227-228: "This suggests that models trained with more steps, whose $|\Delta\theta|$ is usually larger, require smaller optimal $\alpha$ values."

Section 3.4: The authors successfully demonstrate in this section that length bias can be an issue for the method. But this does not prove that length bias is the only issue causing reduced accuracy when training for longer. Also, shouldn't the same logic (length bias being an issue) apply to the standard training procedure? Yet, training for the full recipe produces optimal results, as opposed to training for a fraction of training steps. Why is length bias only an issue for ExPO?

We have not claimed or proved that length bias is the "only" issue affecting the accuracy of $\Delta\theta$. We want to point out that "This issue" in line 242 refers to "increasing the training steps makes the model more prone to learning spurious features from the training data", not length bias. If this wording caused confusion, we sincerely apologize. We use length bias here as an example of data quality mainly because it is a typical example that is simple and amenable to controlled experiments.

Finally, we hope our responses can address your concerns, and we look forward to your favorable consideration of our work.

Thank you for your review and feedback. We address your concerns and questions as follows.

For your main issue, we first clarify:

The theory presented essentially states, "we can partially train an RLHF model, then predict the result of fully training. The resulting predicted model will achieve the accuracy that would have been obtained by full training".

(We first clarify that our term "full training" refers to the baseline that uses the complete training steps. This does not imply that the model is trained to optimality.)

Our formalization in Section 2 does not imply that ExPO is "predicting full training". ExPO hypothesizes that: for a partially-trained model $M_1$, there may exist a better-performing model $M_2$ on the extrapolation trajectory from $M_0$ to $M_1$. This $M_2$ may not (and very likely cannot) be obtained by continuing to train $M_1$, given the inherent uncertainty and non-uniqueness of neural network optimization and convergence (for large language models, this also involves learning rate dynamics, gradient clipping, etc.). ExPO does not aim to, nor can it "predict" the results of full training.

... the results demonstrate that the extrapolation method results in strong gains even when M1 is a fully trained model. (This seems at first like a sign that M1 is undertrained, ...)

As you noted, we conjecture this is because $M_1$ has not been trained to optimality. Given the inherent challenges in training LLMs, alignment training is very likely to be suboptimal (as also evidenced in Section 4.1).

... Based on the theory, the extrapolated model M2 should be most accurate if the search for M2 occurs along the line between M0 and the "fully trained" M1.

We respectfully disagree. A fully-trained model (i.e., the baseline using complete training steps) is not necessarily trained to optimality, thus its indicated $\Delta\theta$ may not be optimal (e.g., it may overfit or capture spurious features in the data, as stated in lines 240-242). A partially-trained model may indicate a more accurate $\Delta\theta$, in which case the optimal model on this extrapolation trajectory would outperform the optimal model on the trajectory indicated by the fully-trained model (according to our discussion in lines 266-268, since $\theta_2=\theta_0+(1+\alpha)\Delta\theta$, the achievable performance of $M_2$ is uniquely determined by $\Delta\theta$).

This line is approximated by training a proxy (e.g. M_1^{20\\%}) and searching along the line between M0 and the proxy.

Based on the previous responses, ExPO is not "predicting full training", thus the "proxy" is not approximating the trajectory from $M_0$ to the fully-trained $M_1$.

But for some reason, this proxy yields better results than searching the line between M0 and M1 (22.7% win rate compared to 18.0% win rate). Thus, a misalignment of the search space is beneficial to the model. This seems like a clear sign that something not predicted by the theory is causing the benefits in accuracy.

Based on the previous responses, ExPO is not "predicting full training", and the "proxy" does not necessarily (and likely will not) indicate the same extrapolation direction $\Delta\theta$ as the fully-trained model. Therefore, the experimental results do not contradict Section 2. On the other hand, the performance improvements come from the accuracy of $\Delta\theta$, which is exactly what we discuss and analyze in detail in Section 3.3-3.5.

Given the above observations, it seems like the true gains in accuracy could be coming from the fact that ExPO models are selected using a method that finds the best-performing model on the UltraFeedback dev set as calculated by another reward model.

We respectfully disagree. As stated in lines 266-268, since $\theta_2=\theta_0+(1+\alpha)\Delta\theta$, and $M_0$ is fixed, the achievable performance of $M_2$ is uniquely determined by $\Delta\theta$. The development set and reward model only assist in searching for $\alpha$, they do not change $\Delta\theta$, and thus naturally cannot change the achievable performance of $M_2$ (although due to limitations of the reward model, the searched optimal $\alpha$ may vary slightly; however, improving reward models to enhance search results is beyond the scope of our work).

My point here is, it seems like the theory presented doesn't match the given results, and there's some other reason for the gains in accuracy.

Based on the above discussion, we respectfully suggest that you may have misunderstood the idea of ExPO (which is not "predicting full training"), and we have provided clarification and explanation. We hope this addresses your concerns about ExPO's soundness.

Thank you for your review and feedback. We address your concerns and questions as follows.

More in-depth discussion of the method is necessary (Why does it work? When does it fail? etc.), For example: it would be helpful to discuss why extrapolation can enhance alignment while being ineffective for general training.

We discussed this in Section 4.3. Based on our analysis in Section 3.3, alignment training typically involves smaller $\Delta\theta$, while general training (like SFT) has larger $\Delta\theta$ that invalidates ExPO's first-order approximation. Additionally, alignment training only involves adjusting the response distribution, while SFT's $\Delta\theta$ usually includes adaptation to chat templates - amplifying this part of information in $\Delta\theta$ can easily lead to model collapse and affect normal response generation.

Analysis on the selection choice and importance of \alpha across various model sizes

Due to computational resource constraints, we were unfortunately unable to train larger models, but instead applied ExPO directly to open-source models (Section 4.1). We found that the value of $\alpha$ is not solely determined by model size, but is closely related to how $M_1$ was trained (based on our review of these open-source models' documentation or technical reports, such as tulu2, snorkel, starling, internlm2, etc.). The performance trends of different $M_2$ models with varying $\alpha$ are similar, as can be seen in Figure 6 of the appendix (the green curve for 40% training steps).

Finally, we hope our responses can address your concerns, and we appreciate your favorable consideration of our work.

Thank you for your review and feedback. We address your concerns and questions as follows.

Given that α in EXPO operates over an open range [0,+∞), what strategies do the authors suggest for effectively managing and narrowing down the search space for α?

The paper mentions that EXPO’s α search process takes less than 0.5 GPU hours, which appears efficient. Could the authors provide a more detailed breakdown of how the search process unfolds? Specifically, how many steps or iterations are typically required, and what is the distribution of time spent in binary search versus grid search?

We answer these two questions together. We use the hyperparameter search in Table 1 as an example (the search process is actually reflected in Figure 6 in the appendix).

• Starting with M_2^{10\\%}:
  • First, with an interval of 5, we tried $\alpha=5$ and $\alpha=10$. We found both significantly outperformed $M_1$, but $(\alpha=5)>(\alpha=10)$.
  • Then, setting the search range to $[5, 10]$ with an interval of 1, we applied binary search and tried $\alpha=7$ and $\alpha=8$. We found $(\alpha=8)>(\alpha=7)$. We then tried $\alpha=9$ and found $(\alpha=8)>(\alpha=9)$.
  • We thus determined $\alpha=8$ as optimal (smaller search intervals might yield better results, but we deemed this unnecessary in practice).
• Moving to M_2^{20\\%}:
  • With previous experience, we first tried $\alpha=2$ and $\alpha=4$ with an interval of 2. We found $\alpha=2$ significantly outperformed $M_1$, but $\alpha=4$ performed worse than $M_1$.
  • Then, setting search ranges to $[0, 2]$ and $[2, 4]$ with an interval of 1, we applied binary search and tried $\alpha=1$ and $\alpha=3$. We found $(\alpha=2)>(\alpha=3)>(\alpha=1)$.
  • Next, with an interval of 0.5 in $[2, 3]$, we tried $\alpha=2.5$ and found $(\alpha=2.5)>(\alpha=2)$.
  • We thus determined $\alpha=2.5$ as optimal. This took 5 searches total, each taking about 5min (using one A100 80GB, including inference on dev set and reward model scoring), totaling about 0.5 GPU hours.
• For M_2^{40\\%}:
  • Based on previous experience, we first tried $\alpha=0.5$ and found it outperformed $M_0$
  • Then with an interval of 0.1, we applied grid search and tried $\alpha=0.6$ and $\alpha=0.4$. We found $\alpha=0.6$ performed worse than $M_1$ (Note that this is a key motivation for using $[0, 0.5]$ as search range with 0.1 interval for M_2^{100\\%} and models in Section 4.1), while $(\alpha=0.5)>(\alpha=0.4)$.
  • We thus determined $\alpha=0.5$ as optimal.

Overall, we (and in practice) don't search blindly, but flexibly combine binary search, grid search, and dynamically adjusted search intervals. These strategies are simple, practical, and represent consensus in practice. Also note that the above search only requires inference-level GPU hardware (e.g., A10 24GB). Therefore, we believe that compared to the reduced training overhead (from 12 GPU hours for M_1^{100\\%} to 2.5 GPU hours for M_1^{20\\%}) and training-level GPU hardware (from eight A100 80GB to one A10 24GB), the $\alpha$ search process in ExPO is more economical and efficient.

Finally, we hope our responses can address your concerns, and we look forward to your favorable consideration of our work.