Align → KD Pipeline in LLM Training

• Align → KD Pipeline is a training strategy where alignment precedes knowledge distillation to preserve target behaviors in large language models.
• The methodology leverages a high-recall reference model to effectively recover rare and valuable outputs, ensuring improved target precision and reward metrics.
• Practically, adopting the Align → KD sequence mitigates sampling and learning traps, thus optimizing model performance even when transitioning to compact, distilled models.

The Align → KD Pipeline refers to the principled sequence in which alignment (notably preference alignment or other task-specific objectives) is performed before knowledge distillation (KD) in training machine learning models, particularly LLMs. This order is motivated by the insight that a reference model's “distributional recall”—that is, its coverage of rare or low-probability but desirable behaviors—critically determines the effectiveness of subsequent alignment. Applying KD before alignment typically results in a compact model with diminished recall, constraining the aligned policy's ability to recover or prefer rare target outputs. Recent research establishes that to robustly align desired behaviors and maximize target-oriented metrics, alignment must always precede distillation.

1. Theoretical Foundations: Reference Model Recall and Alignment Objectives

The core theoretical observation underpinning the Align → KD pipeline is that preference alignment objectives—such as those in reinforcement learning from human feedback (RLHF) or Direct Preference Optimization (DPO)—anchor the learned policy to a fixed reference model via a divergence penalty, typically the Kullback–Leibler (KL) divergence. This reference acts as a support: it delimits which outputs the learning algorithm can access or reinforce.

Formally, in RLHF with proximal policy optimization and KL regularization, the agent minimizes:

$$\mathcal{J}(\theta) = \mathbb{E}[R(y|x)] - \beta \, D_{\mathrm{KL}}(\pi_\theta(y|x) \;\|\; \pi_{\text{ref}}(y|x))$$

where $R(y|x)$ encodes reward, $\pi_\theta$ is the policy, and $\pi_{\text{ref}}$ the reference. If $\pi_{\text{ref}}$ assigns near-zero probability to a rare but desirable behavior $y^*$, the KL penalty $\log(\pi_\theta(y^*|x)/\pi_{\text{ref}}(y^*|x))$ diverges, producing a negative-infinite penalty and precluding $y^*$ from being learned, regardless of the reward.

In DPO, for a preference pair $(y_{\text{winner}}, y_{\text{loser}})$, the loss includes the term

$$- \log \sigma \bigg(\beta \Big[ \log \frac{\pi_\theta(y_{\text{winner}}|x)}{\pi_\theta(y_{\text{loser}}|x)} + \log \frac{\pi_{\text{ref}}(y_{\text{loser}}|x)}{\pi_{\text{ref}}(y_{\text{winner}}|x)} \Big] \bigg)$$

If the reference model's probability on $y_{\text{winner}}$ is vanishingly small, the “reference log–ratio” saturates the loss, yielding vanishing gradients (gradient starvation) and preventing learning.

This establishes an inherent coupling between the alignment algorithm and the distributional support—i.e., recall—of the reference model used. High recall is necessary for the acquisition (or recovery) of rare behaviors via alignment.

2. Experimental Evidence: Synthetic and LLM Case Studies

The minimal working explanation is empirically validated in two settings:

Mixture-of-Gaussians (MoG) Synthetic Benchmark:

A ground-truth multimodal “reward” distribution $p^*$ (with eight modes) is approximated in two steps:

• A high-recall pretrained model ($p'$) covers six modes.
• A low-recall, knowledge-distilled model ($p''$) covers only a subset owing to reduced temperature during KD.

When aligning toward a reward concentrated in a rare mode, alignment anchored to $p''$ frequently fails: samples from the mode are absent or suppressed, and the KL penalty dominates. In contrast, using $p'$ as the reference in alignment enables effective preference recovery. Metrics such as Target Precision and Final Average Reward are substantially higher for the Align → KD pipeline (align first using $p'$, then distill into a compact model) compared to the KD → Align baseline.

LLM Alignment with the SmolLM2 Family:

With SmolLM2 models, the larger 360M model serves as the high-recall reference. Two pipelines are compared:

• KD → Align: Small model is first distilled, then aligned.
• Align → KD: Alignment is performed with the high-recall reference, followed by distillation.

Empirically, Align → KD robustly achieves higher rewards, target precision, and training stability (lower variance across seeds). The KD → Align pipeline underperforms, especially on rare behaviors, confirming the theoretical mechanism described above.

3. Distributional Support, Sampling Traps, and Optimization Pathologies

The key limitation of KD → Align is twofold:

• Sampling Trap: If rare behaviors are pruned away by distillation, subsequent alignment never samples or reinforces those behaviors due to their absence in the low-recall reference.
• Learning Trap: Even if sampled, strong regularization penalizes deviations from the (pruned) reference, overwhelming the preference reward or loss signal. In extreme cases, this results in “gradient starvation,” with the learning signal for rare outputs forced to zero.

The Align → KD sequence sidesteps these traps: alignment is anchored on a broadly supported, high-recall model first, so desired behaviors remain within reach both for sampling and optimization.

4. Practical Recommendations for Pipeline Design

Given this analysis, the following design principle is mandated: always perform preference alignment or similar objectives on a high-recall reference model before any model compression via KD. Only after the target behaviors are robustly aligned should the solution be knowledge-distilled into a compact student, with distillation temperature, capacity, and other hyperparameters tunable for desired trade-offs.

Best practices include:

• Selecting a reference model for alignment that maximally covers the full behavioral support required by downstream tasks.
• Employing early stopping criteria on reward differences and diversity metrics to avoid mode collapse during alignment.
• Monitoring target precision and recall in both the alignment and distillation phases.
• Avoiding knowledge-distilled reference models for alignment anchoring, as these structurally preclude recovery of rare but desirable modes.

5. Impact and Broader Implications

This revised pipeline fundamentally alters accepted practice in both the scaling and deployment of LLMs and related models where alignment (e.g., preference tuning, RL with KL regularization, DPO-based methods) is necessary. The principle that “alignment must precede distillation” directly mitigates phenomena such as behavior loss, mode collapse, and poor reward orientation in compact models.

Beyond LLMs, this guiding framework generalizes to any setting where alignment objectives operate via reference regularization anchored to a model's support: multi-modal generative modeling, structured prediction, and possibly reinforcement learning in the presence of policy constraints.

This approach also highlights the importance of evaluating reference model recall not just for downstream accuracy, but as a first-order constraint determining alignment success. Any model distillation, pruning, or quantization schedule that truncates reference support prior to alignment introduces irrecoverable learning barriers for the target behaviors.

6. Open Challenges and Further Directions

Key open questions include:

• Automated measures of distributional recall to guide reference model selection.
• Dynamic or adaptive reference strategies during the alignment process.
• Methods for maintaining or selectively expanding recall during distillation for complex or evolving behavior sets.
• Integration with advanced sequence modeling objectives and potential extensions to non-textual domains.

The conclusions drawn provide a robust prescription: design alignment pipelines to operate on the highest recall model available, and only afterwards transfer the aligned behaviors to a distillate model of the desired capacity.

Summary Table: Pipeline Comparison

Pipeline	Reference Recall	Alignment Behavior	Downstream Target Precision
Align → KD	High	Recovers all desired behaviors	High
KD → Align	Low	Prunes rare/target behaviors	Low

This summary draws directly from formal, experimental, and practical findings, establishing the Align → KD pipeline as essential for robust preference alignment in model training (Cha et al., 28 Sep 2025).