DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss
Ruixuan Tu, Yuyang Dong
The University of Tokyo · RIKEN AIP · SB Intuitions Corp.
Published October 13, 2026 · on the feed October 19, 2026 · ▲ 3 · arXiv 2610.17644
The abstract
Knowledge distillation (KD) frameworks place teacher and student on the same GPUs, so per-GPU memory limits the teachers that fit. DisaggKD separates them across nodes: teacher replicas stream hidden states to the student, a KD loss over row windows and vocabulary shards equals the full-vocabulary loss in exact arithmetic without forming the full-vocabulary logits, and a prefetch ring paces producer and trainer. Closed-form memory and traffic models locate the memory limits of each layout. On the Qwen3.5-35B-A3B to 9B pair, the disaggregated layout matches the LM loss of the colocated run with the same loss and is faster than that run; our fused loss also lets a colocated run train on as few nodes at faster speed, while the split is memory efficient when the student spans several nodes. On the larger pairs, colocated runs with a fused loss still run OOM where the split trains, because teacher weights and student states fill the budget.
The Brief
1:51A short podcast summary of the paper. Written and voiced by AI.
Transcript
DisaggKD. Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss. By Ruixuan Tu and one other, from the University of Tokyo, the RIKEN Center for Advanced Intelligence Project, and one other institution. Published on the thirteenth of October. No code repository is linked to it.
Knowledge distillation frameworks put the teacher and the student on the same GPUs. That arrangement quietly sets a ceiling. Whatever memory one GPU has is shared between the two models, so the size of teacher you can distill from is limited by what fits alongside the student.
DisaggKD pulls the two apart and places them on separate nodes. Teacher replicas stream hidden states across to the student, rather than living next to it.
Three pieces make the split workable. A distillation loss is computed over row windows and vocabulary shards, and in exact arithmetic it equals the full-vocabulary loss, without ever forming the full-vocabulary logits. A prefetch ring paces the producer and the trainer against each other. And closed-form models of memory and traffic locate where each layout runs into its memory limits.
The main pairing distills from Qwen3.5-35B-A3B into a nine billion parameter student. There, the disaggregated layout matches the language modelling loss of a colocated run using the same loss, and runs faster than it. The fused loss helps the colocated layout as well, letting it train on as few nodes at a faster speed, while the split layout is the memory efficient choice once the student spans several nodes. On larger pairs, colocated runs still run out of memory even with the fused loss, because teacher weights and student states fill the budget between them. The split trains.
Two contributions sit side by side here. The fused loss is useful on its own, whichever layout you pick. Separating teacher from student is what removes the ceiling, and its case is strongest exactly where the models are too big to share a GPU.
More with an account
Listen to every day’s papers
Daily playlists, deep dives, quests for the papers, and quizzes on their key terms.
Join free