Advancing Model Refinement: Muon-Optimized Distillation and Quantization for LLM Deployment

원문 초록

Correction notice (v2). Two implementation issues affect these experiments. All LoRA training ran under DeepSpeed, which replaced the optimizer passed to the Trainer with AdamW at a constant learning rate, so the Muon and Adam arms used the same optimizer. The cross-entropy term was computed without the next-token shift, so the preference for alpha = 1 reflects that objective rather than the value of the teacher. Each comparison also uses a single GPTQ quantization, and repeated quantizations differ by up to several points. The findings on Muon's quantization robustness, the distillation weight and accuracy relative to GPTQ alone should not be relied on. See the notice on page 1. Large Language Models (LLMs) enable advanced natural language processing but face deployment challenges on resource-constrained edge devices due to high computational, memory, and energy demands. Optimizing these models requires addressing three key challenges: acquiring task-specific data, fine-tuning for performance, and compressing models to accelerate inference while reducing resource demands. We propose an integrated framework combining GPTQ-based quantization, low-rank adaptation (LoRA), and a specialized data distillation process to significantly reduce model size and complexity while preserving or enhancing task-specific performance. By leveraging data distillation, knowledge distillation via Kullback-Leibler divergence, Bayesian hyperparameter optimization, and the Muon optimizer, our pipeline achieves up to 2x memory compression (e.g., reducing a 6GB model to 3GB) and enables efficient inference for specialized tasks. Empirical results demonstrate superior performance on standard LLM benchmarks compared to GPTQ quantization alone, with the Muon optimizer notably enhancing fine-tuned models' resistance to accuracy decay during quantization.