Student guided distillation improves efficient large language model task routing
Student-Guided Teacher Distillation for Efficient LLM Task Routing: Positioning Against Jev-Style System-1 Classifiers
Machine LearningArtificial Intelligence
Summary
Routing user requests to the right AI task can be slow and expensive when the system has many categories to check. The authors study a way to speed this up by training a smaller model (student) to quickly suggest a few likely categories, then having a bigger, more powerful but slower model (teacher) refine those choices. This process repeats, improving the smaller model's suggestions over time. They find that using the teacher's full knowledge is important when training the student and that their method keeps almost all important decision info while being cheaper.
What this means in practice
- •For llm service operators: Improve efficiency of routing user queries to specialized large language models by deploying a trained student classifier to reduce evaluation costs.
- •For enterprise chatbot developers: Enhance response accuracy and speed by integrating a task-specific student-teacher classification pipeline for directing user intents among many specialized tasks.
Authors
Haifeng Wu, Srinivasan Manoharan, Jian Wan, Fangbo Tu, Junhua Zhao, Xin Chen
Abstract
Zero-shot classifiers are useful for routing user requests to specialized LLM tasks, but scoring every request against a large candidate set is expensive: a zero-shot NLI classifier must evaluate one premise-hypothesis pair per label, so cost scales linearly with taxonomy size. We study a student-guided teacher distillation pipeline for a fixed taxonomy of 60 LLM task categories: a compact ModernBERT classifier predicts the full category distribution in one forward pass and retrieves a small top-k candidate set, and a larger DeBERTa-v3 zero-shot NLI classifier reranks only those candidates rather than all 60 labels; the resulting teacher labels iteratively improve the student, which produces sharper candidates for the next round. Unlike generic embedding retrieval or clustering-derived shortlists used in extreme multi-label classification, our candidate generator is trained end-to-end on the target taxonomy and is the same model serving production traffic, distinguishing it from LLM-routing work that routes between candidate models, and from concurrent System-1 encoder-classifier proposals (e.g. TypeSafe AI's Jev and the open-source Laya project) whose training methodology is undocumented or RL-based. Our best student checkpoint reaches 77.5% teacher agreement on a 200-example evaluation set, and preliminary coverage measurements show Coverage@16 of 91-100%, suggesting top-k sets retain most of the teacher's decision-relevant information. We further show truncated top-k teacher scores should not be treated as full 60-class soft targets for KL distillation: zeroing untruncated classes destroys the dark knowledge soft-label distillation depends on, introducing systematic bias rather than a harmless sparse approximation. A complete evaluation, including coverage at multiple k on a held-out set, an embedding-retrieval baseline, and a larger human-reviewed test set, remains in progress.