Distilling large language models for code generation via ranking supervision

Area of Science:

• Artificial Intelligence
• Machine Learning
• Natural Language Processing

Background:

• Large language models (LLMs) excel at code generation but are computationally expensive.
• Knowledge distillation is a technique to compress LLMs, but traditional methods like KL divergence have limitations.
• Existing methods struggle with long-tail noise and can underperform supervised fine-tuning.

Purpose of the Study:

• To develop a more efficient knowledge distillation method for LLMs in code generation.
• To reduce computational cost and memory consumption during LLM compression.
• To improve the performance of distilled LLMs compared to existing techniques and supervised fine-tuning.

Main Methods:

• Proposed a novel distillation approach using ranking supervision with a ListNet-based loss.
• Focused on aligning the teacher's token ranking preferences rather than the full probability distribution.
• Evaluated the method on benchmarks like HumanEval, MBPP, DS-1000, and MultiPL-E.

Main Results:

• The proposed ranking distillation method consistently outperformed supervised fine-tuning and KL divergence baselines (FKL, RKL).
• Achieved comparable training times with significantly reduced GPU memory consumption.
• Demonstrated stable performance across various model scales and families (Qwen2.5-Coder, DeepSeek-Coder).

Conclusions:

• Ranking distillation is an effective technique for compressing LLMs for code generation tasks.
• This method offers a practical solution for deploying powerful LLMs in resource-constrained environments.
• The findings provide valuable insights for future research in LLM compression and model optimization.