AKD : Adversarial Knowledge Distillation For Large Language Models Alignment on Coding tasks

TL;DR

AKD tackles the challenge of aligning Code-LLMs under limited high-quality training data by introducing adversarial knowledge distillation, which uses a teacher–student setup and Direct Preference Optimization to distill a larger model's capabilities into a smaller one. The framework constructs a three-way dataset $(p,c,r)$ and uses margin rewards $M=R(c)-R(r)$ to steer adversarial data selection with $P(p) = \exp(-M(p)) / \sum_{p'} \exp(-M(p'))$, while employing curriculum-based prompts to progressively challenge the student. Empirically, AKD matches or surpasses self-supervised fine-tuning on code benchmarks with far smaller synthetic datasets, demonstrating data efficiency and robustness; results suggest scalability across model families and potential for broader domain adaptation. Looking forward, the work envisions enhancements via larger datasets, compiler feedback signals, and cross-domain or cross-family speculative decoding to further improve practical and reliable automated code generation.

Abstract

The widespread adoption of Large Language Models (LLMs) for code generation, exemplified by GitHub Copilot\footnote{A coding extension powered by a Code-LLM to assist in code completion tasks} surpassing a million users, highlights the transformative potential of these tools in improving developer productivity. However, this rapid growth also underscores critical concerns regarding the quality, safety, and reliability of the code they generate. As Code-LLMs evolve, they face significant challenges, including the diminishing returns of model scaling and the scarcity of new, high-quality training data. To address these issues, this paper introduces Adversarial Knowledge Distillation (AKD), a novel approach that leverages adversarially generated synthetic datasets to distill the capabilities of larger models into smaller, more efficient ones. By systematically stress-testing and refining the reasoning capabilities of Code-LLMs, AKD provides a framework for enhancing model robustness, reliability, and security while improving their parameter-efficiency. We believe this work represents a critical step toward ensuring dependable automated code generation within the constraints of existing data and the cost-efficiency of model execution.

Figures (2)

• Figure 1: Adversarial Knowledge Distillation framework comprises three components, differentiated by color: models (green), training processes (red), and synthetic data generation elements (blue). In the initial step, topics and profession seeds are used to generate an initial dataset. Profession is acting as a persona to help diversify the dataset. For subsequent iterations, exercises are sampled based on margin rewards calculated from the previous training cycle. Once the exercises are selected or generated, the teacher model produces the "chosen" solutions, while the student model generates the "rejected" solutions. The resulting dataset, containing both chosen and rejected solutions, is then used to train the student model via Direct Preference Optimization (DPO). Margin rewards, which reflect the difference in performance between the teacher and student on specific tasks, are calculated to rank exercises. This ranking informs the sampling process for generating the next dataset, ensuring a targeted and iterative improvement cycle for the student model.
• Figure 2: Margin Rewards during training: At the 200-step mark, an adversarial step generates a new dataset, continuing training with the same optimizer states. A small performance dip is observed in some experiments due to increased exercise difficulty.