On-policy distillation reduces train-inference mismatch by supervising student-generated trajectories, but full-vocabulary token-level distribution matching remains expensive for large language models. A common approximation truncates reverse-KL distillation to the student’s top-k tokens, which preserves efficiency but drops tail-token supervision and ignores support-level mass mismatch between the student and teacher. We analyze these approximation gaps and propose Stochastic Support Sampling with Teacher Mass Reweighting, which samples fixed-size supports from the student distribution and uses teacher support mass as a reliability weight for the conditional distillation loss. Under the same support budget, our method outperforms the strongest baseline by +1.4 hard-math average points on Qwen3-8B→Qwen3-1.7B-Base, and by +1.1 math average points and +2.8 code average points on DeepSeek-R1-Distill-Llama-8B→Meta-Llama-3.1-8B-Base.

Mixture of experts (MoE) is a popular technique to improve capacity of Large Language Models (LLMs) with conditionally-activated parallel experts. However, serving MoE models on memory-constrained devices is challenging due to the large parameter size. Typical solutions such as memory swapping or expert pruning may lead to significantly higher latency or severe accuracy loss.In this paper, we introduce SwapMoE, a framework for efficient serving of MoE-based large language models with tunable memory budgets. The main idea of SwapMoE is to keep a small dynamic set of important experts, namely Virtual Experts, in the main memory for inference, while seamlessly maintaining how the Virtual Experts map to the actual experts. Experiments have shown that SwapMoE can reduce the memory footprint while maintaining reasonable accuracy. For example, on text summarization tasks with Switch Transformer, SwapMoE can reduce the memory consumption from 14.2 GiB to 4.7 GiB, together with 50% latency reduction and a slight Rouge-2 score drop of 0.041.