Advantage-Guided Distillation for Preference Alignment in Small Language Models

Thank you for your thoughtful review and for recognizing the importance of our contribution to leveraging preference signals in knowledge distillation.

W1: Unclear how the improvement is related to the motivation of the method; ADPA seems not to need preference labels during training; how is ADPA related to preference alignment?

Response:
Thank you for the insightful question. While ADPA does not directly use preference labels during the student model's training phase, it crucially relies on the advantage function: $A\_{\text{dpo}}(s\_t, a\_t) = \beta \left( \log \pi\_{\text{dpo}}(a\_t|s\_t) - \log \pi\_{\text{ref}}(a\_t|s\_t) \right),$ which requires log probabilities from both the DPO teacher and the reference teacher.

The DPO teacher is trained using preference data, and thus, the advantage function inherently captures the preference information learned by the teacher. By distilling this advantage function, the student model learns to prioritize actions that are preferred according to the teacher's alignment with human preferences.

In essence, ADPA transfers the preference alignment from the teacher to the student, even though the student does not directly use the preference labels during its own training. We have clarified this relationship in the revised paper, emphasizing how the advantage function serves as a channel for preference information.

---

W2: Why consider the difference between $\pi\_{\text{dpo}}$ and $\pi\_{\text{ref}}$ instead of $\pi\_{\text{dpo}}$ alone? What happens if $\pi\_{\text{dpo}}$ is removed from Equation (11)?

Response:
Thank you for raising this important point and for providing us with the opportunity to explain our design choice in detail. The reference model $\pi\_{\text{ref}}$ serves as a baseline to quantify the improvement in alignment achieved by the DPO teacher $\pi\_{\text{dpo}}$. The advantage function: $A\_{\text{dpo}}(s\_t, a\_t)$ measures how much more preferable an action is compared to the reference policy. This relative measure is crucial because it highlights the specific adjustments made during preference alignment.

If we remove $\pi\_{\text{ref}}$ and use $\pi\_{\text{dpo}}$ alone, the training objective becomes minimizing the reverse cross-entropy between the student model and the DPO teacher. Therefore, student models would attempt to mimic the teacher's policy without understanding which aspects are aligned with human preferences. In our experiments, we observed that removing $\pi\_{\text{ref}}$ from the advantage function leads to diminished performance, as shown in an ablation study (Table 2: ADPA w/o Ref teacher).

If we remove $\pi\_{\text{dpo}}$ and use $-\pi\_{\text{ref}}$ alone, the student model's objective becomes maximizing the reverse cross-entropy with the reference teacher. This encourages the student model to produce outputs that diverge as much as possible from the reference teacher's outputs, actively pushing it away from the reference teacher without positive guidance on where to move towards. Without the positive guidance provided by $\pi\_{\text{dpo}}$, the supervised fine-tuning loss term $\mathcal{L}\_{\text{SFT}}$ alone is insufficient to regularize the training process. Consequently, this approach is prone to training instability or even collapse, as the student model lacks direction on which actions are preferable and may deviate in unintended ways, leading to degraded performance. In practice, we found that student models trained in this setting often produce repetitive or meaningless content. Therefore, we did not show their performance on the test set.

We appreciate your encouraging remarks and are glad that you find our work interesting and effective.

Question: Experiments on more popular models.

Response:
Thank you for this valuable suggestion. In the revised paper, we have added experiments using LLaMA-3.2-1B as the student model and LLaMA-3.1-8B as the teacher model. We chose LLaMA-3.2-1B over other models like Qwen-0.5B due to its wider usage in the research community.

Below, we present the evaluation results for LLaMA-3.2-1B trained on DPO-MIX-7K and HelpSteer2. We can observe that our proposed ADPA+ surpasses all baseline methods, indicating that ADPA+ is a robust and generalizable approach for the preference alignment of small language models.

Thank you for your positive assessment and constructive suggestions regarding our work.

W1: Limited baselines of alignment methods; only DPO is included as an alignment baseline.

Response:
We sincerely appreciate your insightful feedback regarding the selection of alignment baselines in our study. To clarify, our original submission included additional baselines beyond DPO. Specifically, we evaluated DPKD (utilizing the teacher model as the reference in DPO), DDPO (aligning implicit rewards derived from the student model with those from the teacher model), and PLAD (leveraging teacher-student generated pseudo-preference pairs). These methods investigate diverse dimensions of alignment through knowledge distillation or model guidance, and are closely related to the preference alignment objective of our work.

We fully agree that incorporating more recent preference alignment techniques can further strengthen the robustness of our comparative analysis. In response to your suggestion, we have extended our experiments to include the latest preference alignment methods, such as Simple Preference Optimization (SimPO) [1] and Weighted Preference Optimization (WPO) [2]. Our updated empirical results indicate that while these advanced preference alignment approaches yield moderate improvements in alignment metrics over the DPO baseline, our proposed ADPA framework consistently demonstrates more substantial performance gains.

References:
[1] SimPO: Simple preference optimization with a reference-free reward. 2024. In NeurIPS.
[2] WPO: Enhancing RLHF with Weighted Preference Optimization. 2024. In EMNLP.

---

W2: Computational overhead of the KD methods compared to directly applying alignment techniques.

Response:
We appreciate the reviewer’s insightful comment regarding the computational overhead of KD-based alignment methods compared to directly applying alignment techniques.

To address this, we provide the following clarifications. First, while KD methods involve the use of additional teacher models during training, the associated computational cost can be mitigated by precomputing the required teacher outputs offline before training begins. Second, once the student model is trained, the inference stage operates without incurring any additional computational overhead. Third, importantly, the primary challenge in aligning small language models lies not in the computational cost during training, but in overcoming the inherent limitations in their capacity and capabilities. Therefore, KD has been widely used as a crucial technique to address these challenges. Notable examples include its application in models such as LLaMA3.2 [3] and Gemma2 [4].

References:
[3] LLaMA 3.2: Revolutionizing edge AI and vision with open, customizable models. Meta AI.
[4] Gemma 2: Improving Open Language Models at a Practical Size. Gemma Team, Google DeepMind.

Dear Reviewer 346S,

We sincerely appreciate the time and effort you have devoted to reviewing our work, especially considering your busy schedule. As the discussion phase between authors and reviewers draws to a close, we kindly request your attention to our responses. In light of your feedback, we have incorporated additional alignment baselines, including recent methods like SimPO and WPO, and found that our ADPA framework consistently outperforms them. Regarding your concern about computational overhead, we have clarified that our method involves minimal additional computations during training, which can be mitigated through offline precomputations, and there is no added cost during inference.

We hope that we have effectively addressed your concerns and would welcome any additional questions or points you might like to discuss. We highly value your feedback and look forward to further discussion.

Thank you for your thoughtful consideration.

Best regards,

The Authors

W2 (the latter part): Absence of experiments assessing the impact of teachers of varying proficiency levels.

Response: We appreciate the reviewer's insightful comment regarding the impact of teachers with varying proficiency levels on small language (student) models. To address this, we conducted additional experiments to directly assess this factor. Specifically, we compared the effects of two teacher models with distinct proficiency levels: H2o-danube2-1.8B, representing a lower-proficiency teacher, and Mistral-7B, a higher-proficiency teacher previously used in our experiments. In both cases, the student model was H2o-danube3-0.5B, and all experiments were performed on the DPO-MIX-7K dataset.

The empirical results are summarized in the following table. We observe that the teacher model's proficiency influences knowledge transfer results. Notably, higher-proficiency teacher models, such as Mistral-7B, consistently enable superior performance in student models through the distillation process. These results validate the critical role of teacher proficiency in knowledge transfer and will be included in the revised manuscript for clarity and completeness.

---

Q1: In Formula 11, $L\_s$ seems to be undefined. Does it refer to $L\_{\text{SFT}}$?

Response: Thank you for pointing this out. You are correct that in Equation (11), $L\_s$ refers to the supervised fine-tuning loss $\mathcal{L}\_{\text{SFT}}$. We apologize for the oversight in the notation. In the revised manuscript, we have thoroughly reviewed and corrected the notation to ensure all terms are clearly defined and consistent throughout. Specifically, we have replaced $L\_s$ with $\mathcal{L}\_{\text{SFT}}$ in Equation (11) and throughout the manuscript where applicable.

We appreciate your positive feedback and thoughtful review of our work.

W1: Computational overhead introduced by fine-grained signals.

Response: We appreciate the reviewer's concern regarding the computational overhead introduced by the fine-grained reward signals in our ADPA framework. We would like to clarify that the training process of ADPA does not significantly increase computational overhead, especially when compared to traditional knowledge distillation (KD).

Specifically, the computation involving the fine-grained reward signals in ADPA requires calculating $E\_{a\_t \sim \pi\_{\text{stu}}(\cdot|s\_t)} \left( \log \pi\_{\text{dpo}}(a\_t|s\_t) - \log \pi\_{\text{ref}}(a\_t|s\_t) \right),$ which is somewhat similar to the KLD objective commonly employed in traditional KD. Both objectives require computing an expectation over the difference between two log probability distributions. Although ADPA introduces an additional teacher, the computational cost associated with this is minimized through precomputation of these values offline prior to training.

During our experiments training a h2o-danube2-1.8B student model on the DPO-MIX-7K dataset using 8 RTX 3090 GPUs with a per-device batch size of 8 (global batch size of 128) and DeepSpeed ZeRO-2, APKD resulted in only a slight increase in both memory usage and training time, compared to Vanilla KD. Specifically, memory usage rose by less than 3% (from 21,957 MB to 22,541 MB), and training time increased by less than 4% (from 78 minutes to 81 minutes). These minor overheads are well within acceptable limits for practical applications.

Thus, the overall impact on computational resources remains minimal and well within acceptable bounds.

---

W2: Utilizing a rather singular evaluation dataset; lack of results on a broader range of models, domains, and data.

Response: Thank you for this valuable feedback. To address this concern, we have expanded our experimental scope by incorporating additional models and evaluation datasets, as outlined below:

1. Additional Models: We have included experiments with an additional student model, LLaMA-3.2-1B, trained under the guidance of the teacher model LLaMA-3.1-8B on the DPO-MIX-7K dataset. This addition aims to demonstrate the adaptability of our proposed method to popular models widely utilized within the research community.

2. Broader Evaluation Datasets: To ensure a more comprehensive assessment, we evaluated our trained models on the Open LLM Leaderboard, which benchmarks open-source language models across six diverse domains, including scientific reasoning, commonsense reasoning, multi-task understanding, and factual accuracy. This broader evaluation provides a robust analysis of our approach's effectiveness across various tasks.

We present the evaluation results for LLaMA-3.2-1B on MT-Bench, AlpacaEval (ADPA as reference) and the Open LLM Leaderboard below. Our proposed methods, ADPA and ADPA+, consistently outperform all baseline approaches across most benchmarks, demonstrating exceptional performance on ARC-Challenge, HellaSwag, and TruthfulQA. This underscores their effectiveness in tackling diverse tasks. These results will be included in the revised manuscript for clarity and completeness.

---

Dear Reviewer PkN8,

We sincerely appreciate the time and effort you have devoted to reviewing our work, especially given your undoubtedly busy schedule. As the discussion phase between authors and reviewers draws to a close, we kindly invite you to review our responses to your comments.

Our primary aim has been to address your feedback comprehensively and to engage further on any remaining questions or suggestions you may have. Your insights have been invaluable, and we are eager to continue this constructive dialogue.

To address your comments and strengthen our findings, we have undertaken the following actions:

1. Demonstrated that the additional cost is minimal.
2. Expanded our experiments to include additional models, such as LLaMA-3.2-1B, and evaluated them on the OpenLLM Leaderboard.
3. Conducted experiments with teacher models of varying proficiency levels, which reaffirmed the critical role of teacher quality in effective knowledge transfer.

We sincerely thank you once again for your thoughtful feedback and careful consideration. We look forward to any further insights you may have.

Best regards,
The Authors