Advantage-Guided Distillation for Preference Alignment in Small Language Models

TL;DR

This work tackles the challenge of aligning small language models (SLMs) with human preferences by transferring alignment knowledge from a preference-aligned teacher. It introduces Dual-Constrained Knowledge Distillation (DCKD) as a baseline and Advantage-Guided Distillation for Preference Alignment (ADPA), which uses an advantage function derived from a DPO-trained teacher and a reference teacher to supply distribution-level reward signals. Empirical results show that both methods improve SLM alignment, with ADPA-based variants often outperforming baselines and narrowing the gap to larger models; combining ADPA with DCKD (ADPA+) yields the strongest gains across MT-Bench, AlpacaEval, and Open LLM Leaderboard. The work demonstrates the practical viability of leveraging larger, well-aligned models to guide alignment of resource-constrained models, offering scalable pathways for safer and more helpful AI in constrained environments.

Abstract

Alignment techniques enable Large Language Models (LLMs) to generate outputs that align with human preferences and play a crucial role in their effectiveness. However, their impact often diminishes when applied to Small Language Models (SLMs), likely due to the limited capacity of these models. Instead of directly applying existing alignment techniques to SLMs, we propose to utilize a well-aligned teacher LLM to guide the alignment process for these models, thereby facilitating the transfer of the teacher's knowledge of human preferences to the student model. To achieve this, we first explore a straightforward approach, Dual-Constrained Knowledge Distillation (DCKD), that employs knowledge distillation with two KL-divergence constraints from the aligned teacher to the unaligned student. To further enhance the student's ability to distinguish between preferred and dispreferred responses, we then propose Advantage-Guided Distillation for Preference Alignment (ADPA), which leverages an advantage function from the aligned teacher to deliver more nuanced, distribution-level reward signals for the student's alignment. Our experimental results show that these two approaches appreciably improve the alignment of SLMs and narrow the performance gap with larger counterparts. Among them, ADPA demonstrates superior performance and achieves even greater effectiveness when integrated with DCKD. Our code is available at https://github.com/SLIT-AI/ADPA.

Figures (5)

• Figure 1: Preliminary results showing the "alignment tax" on smaller models and the impact of our ADPA+ method on MT-Bench zheng2023judging rating. Under DPO training, the larger model (Mistral-7B) improves notably (+0.34), while smaller models show limited gains (+0.07 for Danube2-1.8B) or even a drop (-0.02 for Danube3-500M). In contrast, ADPA+ enables small models to achieve greater gains (+0.21 for Danube3-500M and +0.60 for Danube2-1.8B).
• Figure 2: Overview of ADPA. Training involves two teacher models: a DPO teacher $\pi_{\text{dpo}}$ fine-tuned on preference data and a reference teacher $\pi_{\text{ref}}$ fine-tuned on instruction-tuning data. The student model is trained via instruction-tuning and then advantage-guided distillation with on-policy data.
• Figure 3: An illustration comparing the sample complexity of different reward granularities. Using a distribution-level advantage function (left), the model directly identifies the optimal action $a_t^* = \mathop{\mathrm{arg\,max}}\limits_{a \in \mathcal{A}} A(s_t, a)$ at state $s_t$, as shown by the red solid line, without needing to explore subsequent states or sample additional trajectories (dotted line). For token-level reward (middle), the model evaluates immediate rewards $r(s_t, a)$ for each action $a \in \mathcal{A}$, transitions to the corresponding next states $s_{t+1}$ (solid line), and may consider future rewards to determine $a_t^*$, leading to a sample complexity of $O(|\mathcal{A}|)$. With sequence-level reward (right), the model generates full trajectories starting from each possible action $a \in \mathcal{A}$, reaching EOS to receive the reward $R(\tau)$. This requires exploring all possible action sequences of length $T - t$, resulting in sample complexity $O(|\mathcal{A}|^{T - t})$.
• Figure 4: Comparison between ADPA and PPO-based methods on the validation set. The x-axis denotes the training epochs, and the y-axis indicates either the average scores (left) or the win rates (right) of responses generated by checkpoints during training, as evaluated using FsfairX.
• Figure 5: Impact of hyperparameters $\alpha$ in DCKD (left) and $\gamma$ in ADPA (right) on the preference alignment of Danube2-1.8B with Mistral-7B as the teacher on the DPO-MIX-7K dataset. The average scores evaluated by FsfairX and Reward Accuracy meng2024simpo are reported.