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

TL;DR

The paper tackles the challenge of deploying large language models on resource-constrained edge devices by proposing an integrated pipeline that blends synthetic data distillation from a strong teacher (T1), logit-based knowledge distillation from a tokenizer-aligned teacher (T2) to a LoRA-tuned student (S1), and GPTQ 4-bit quantization, guided by Bayesian hyperparameter optimization. The Muon optimizer is incorporated to improve robustness to quantization, resulting in up to $2\times$ memory compression and substantial latency reductions while preserving or enhancing task-specific performance across multiple benchmarks. Empirical results demonstrate superior accuracy compared to GPTQ quantization alone, with Muon-enhanced fine-tuning showing reduced accuracy decay under quantization across most tasks. Overall, the framework provides a practical path to deploy compact, task-specialized LLMs on edge devices without sacrificing critical performance metrics.

Abstract

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.

Figures (5)

• Figure 1: Pipeline Overview
• Figure 2: Self-Instruct Pipeline Detail
• Figure 3: Win Count across 3 methods: GPTQ Alone vs Adam-Optimzed vs Muon-Optimized
• Figure 4: Comparison of Accuracy Decrease due to Quantization across 8 Benchmarks; Adam-Optimized vs Muon-optimized. Note: in SIQA benchmark, accuracy increased when quantized
• Figure 5: Throughput Comparison: Pre- and Post- Quantization