MuSID-KD: Multigranularity Knowledge Distillation for Generative Recommendation with Semantic IDs

To diagnose candidate-group ordering, the paper first examines a student whose distillation objective matches only the teacher's probability distribution over candidate items. This paper finds that this is not enough. A student aligned only on the item distribution disagrees with the teacher on the order of candidate groups at intermediate SID levels in 53.8% of comparisons on average. And for a pair of items, matching the teacher at different SID positions helps at some positions and hurts at others in roughly seven out of ten cases. The authors propose MuSID-KD (a multi-granularity distillation method that works level by level over the SID). On three Amazon datasets and three lightweight students, its NDCG@10 is 8.31% higher on average than basic distillation and 1.59% to 5.13% higher than the strongest other method in each setting, with the online inference structure completely unchanged.

Teams trying to cut inference cost for SID-based generative recommendation should read this. It points out an easily overlooked fact: SIDs are hierarchical, so the teacher's knowledge lives not only in the item probabilities at the last level but also in the ranking of prefixes at every level, and reusing an item-level distillation loss throws that part away. The core idea in one sentence: one loss makes the candidate-group probabilities match the teacher at every level, and another matches the teacher only at the levels that can fix the item ranking. The cost is training time growing from 27 to 49 minutes, and all datasets are small public academic sets.

What problem is it actually solving?

SID-based generative recommendation encodes each item as several levels of discrete tokens. The model generates them level by level, and an item's score is the sum of its per-level log-probabilities. Models such as TIGER work well, but beam search expands many prefixes at every level, which is expensive online. The common fix is distillation: keep the teacher's history encoder, replace level-by-level decoding with a few lightweight prediction heads, and have a small student imitate the teacher.

Figure 1 of the paper gives a game recommendation example. Minecraft and Mario Kart 8 share the first two SID tokens (113, 242), so at the second level they belong to the same candidate group. Mario Odyssey starts with (587, 35) and Animal Crossing with (587, 77). Grouping by the first two tokens, a group's probability is the sum of its items' probabilities. The teacher orders the four groups (113, 242) first, then (113, 140), (587, 35) and (587, 77), yet the student ranks (587, 77) first.

There are two problems here. First, group order can only be learned indirectly from item probabilities, and when student capacity is limited, item probabilities can roughly match while group order is wrong. Second, an item's score is the sum of per-level contributions, so matching the teacher at one level can push the final ranking in the wrong direction. The authors measured both:

These two numbers show that item-level distillation misses a lot of information on SID models, and that aligning all levels equally is not the best approach either.

The overall pipeline

MuSID-KD changes only training, not serving. For each training request the teacher first generates up to 64 valid candidate items and caches their scores and per-level token probabilities. The student rescores the same candidates, and three losses are computed:

HCDA (hierarchical candidate distribution alignment) and RIGD (ranking-impact-guided level-wise distillation) are the authors' names for the two new losses. The total loss is a weighted sum of the three:
$$\mathcal{L}={\mathcal{L}}_{\mathrm{CE}}+{\lambda }_H{\mathcal{L}}_{\mathrm{HCDA}}+{\lambda }_R{\mathcal{L}}_{\mathrm{RIGD}}$$
${\lambda }_H$ and ${\lambda }_R$ control the strength of the two distillation losses. Online, the student uses its history encoder, prefix-conditioned prediction heads and constrained beam search as usual, with no new modules.

HCDA: make group probabilities match at every level

Why is this step needed? If the student's item distribution equaled the teacher's exactly, every level's group distribution would match too. But a small model can only approximate, and a small error in item probabilities can flip group order. HCDA normalizes complete SID sequence scores over the candidate set and sums the resulting item probabilities by prefix; these group probabilities differ from the cumulative prefix scores used by online beam search.

HCDA is computed in three steps. First, on the teacher's candidate set, normalize each item's score into a probability with temperature τ. Then sum probabilities by prefix: at level d, items sharing the first d tokens form one group, and a group's probability is the sum of its items' probabilities. Finally, compute the KL divergence between teacher and student at each level and add them with equal weights:
$${\mathcal{L}}_{\mathrm{HCDA}}={\mathbb{E}}_u\left[\sum _{d=1}^L{D}_{\mathrm{KL}}\right({\pi }_T^{(d)}\Vert {\pi }_S^{(d)}\left)\right]$$
${\pi }_T^{(d)}$ and ${\pi }_S^{(d)}$ are the teacher's and student's group distributions at level $d$; the last level is the item distribution itself. A hand-checkable illustration: the teacher gives Minecraft 0.30, Mario Kart 8 0.10 and Zelda 0.25, and the student gives 0.20, 0.15 and 0.30. At the item level the differences are modest, but at the second level group 242 has 0.40 for the teacher, above the 0.25 of group 140; for the student it has 0.35, still above 0.30. If the student drifts a little more, the group order flips, and HCDA penalizes that drift directly.

RIGD: match the teacher only where it helps

Why not align every level? Take the item pair from Figure 1 again. The teacher prefers Minecraft over Mario Odyssey with a score gap of +0.020; the student's gap is −0.059, the wrong order. Split across three levels, the teacher's contributions are −0.074, +0.022 and +0.072, and the student's are +0.090, −0.082 and −0.067. So the student's error comes from levels 2 and 3; at level 1 the student actually favors Minecraft more than the teacher does.

Try replacing just one level with the teacher's value, leaving the others alone, and see what the student's gap becomes:

Matching the teacher at level 1 pushes the pair further apart, so that level should not get supervision for this pair. RIGD turns this "try replacing one level" idea into weights.

It works in four steps. First, pick item pairs: one within the teacher's top 10, the other ranked 11 to 40, because pairs crossing the recommendation cutoff matter most. Next, find the first level where the two diverge; earlier levels have identical tokens and contribute nothing. Then, for each eligible level, replace the student's contribution with the teacher's and compare a margin-based softplus ranking loss before and after; the drop is that level's gain, and negative drops count as zero. Finally, normalize the gains into weights and use a weighted Huber loss to pull the student's per-level contributions toward the teacher's. RIGD supervises contributions only from the first divergence through the penultimate position; the final position disambiguates item identity and is excluded from RIGD supervision, but its contribution remains in the score gap.

An analogy is ranking by total score across several subjects. If a student's total puts two people in the wrong order relative to the answer key, the teacher does not make him revise every subject. The teacher tries one eligible subject at a time and weights it by the reduction in pairwise ranking surrogate loss from switching to the reference score, with zero weight if the loss does not decrease.

Experiments: better on three students and three datasets

The main experiments use Amazon Beauty, Toys and Video Games. Each request looks at up to 20 previous interactions, evaluated with time-ordered leave-one-out. The teacher is a fixed TIGER, and there are three student readouts: attention plus MLP, GRU and temporal-convolution CNN, all sharing the teacher's frozen history encoder, with each student's readout followed by four prefix-conditioned MLP heads. Baselines include three classic ranking distillation methods from recommendation, two general distillation methods for large models, and two SID-specific methods, SID-MLP and SmartGR. Basic distillation means the SID-MLP approach: along the SID prefix of the user's real next item, make the student's next-token distribution match the teacher's at every level, plus cross-entropy. This differs from the diagnosis earlier, which aligned item probabilities over the teacher's candidate set. The table shows NDCG@10:

MuSID-KD ranks first in all 36 comparisons (three datasets, three students, four metrics). The gain is largest for the CNN student: on Toys its NDCG@10 is 18.84% above basic distillation, nearly matching the teacher. The weaker the student, the more it loses from the missing hierarchical information, and the more it gains once that is restored. On Steam-100K, MuSID-KD improves NDCG@10 over basic distillation by 2.15% on average across the three student architectures. Switching the SID from RQ-VAE to residual K-means, it still ranks first in all 24 settings, so the effect does not depend on one tokenizer.

The ablations attribute the gains to each part. Dropping intermediate group alignment and keeping only the item distribution plus RIGD lowers NDCG@10 by 4.21% on average. Deliberately scrambling the correspondence between prefixes and probabilities before HCDA also hurts, which shows that the real group structure matters, not just extra KL terms. The RIGD ablations use the residual: at one level, the teacher's score-gap contribution minus the student's, so for the pair above the level-1 residual is −0.074 − 0.090 = −0.164. Equal weights per level, weights by absolute residual, and equal weights only on levels with positive residuals all fall short of gain-based weights, because an absolute value says how big the gap is but not whether fixing it improves the order.

The diagnostics agree: for the CNN student on Beauty, the KL of the second-level group distribution falls from 0.0663 to 0.0386, and the share of cases where the student disagrees with the teacher's top group drops from 31.58% to 25.23%.

The limitations as I see them

First, the data is small. The main experiments look at only 20 interactions per request, and absolute NDCG@10 is only 0.03 to 0.05, so a 1% to 5% relative gain often amounts to only about a tenth of a percentage point in absolute terms (0.0356 to 0.0369 for the MLP student on Beauty). Whether a gap this size holds on industrial-scale, long-sequence data is the first thing to check.

Second, training cost nearly doubles. For the MLP student on Beauty, 200 epochs take 27 minutes with basic distillation and 49 minutes with MuSID-KD, and peak GPU memory rises from 7.4 GB to 9.2 to 9.5 GB. Most of the overhead is building item pairs and per-level replacement, which grows with larger candidate sets and cutoffs.

Third, the hyperparameters are not intuitive. The default RIGD weight ${\lambda }_R$ is 31.42, and the authors ran sensitivity analysis only on Beauty, so other datasets will need re-tuning.

How a recommendation or search team could try it

1. Start with diagnosis and no training change: take an existing distilled student and its teacher, and compute the intermediate-level group misordering rate on a batch of requests. If it is near one half as in the paper, hierarchical information really is being lost.
2. Add HCDA alone: it only sums existing candidate item probabilities by prefix and adds a few KL terms, so it is the cheapest to implement. Check whether NDCG@10 improves, especially for the lightest student.
3. If that works, add RIGD, starting from the paper's top 10 versus ranks 11 to 40 and adjusting to the real production cutoff. Watch training time and memory, since the number of item pairs grows with the cutoff.
4. Finally compare online throughput and confirm that the student serves as many requests per second as with basic distillation. In the paper the throughput of all three students is essentially unchanged, because distillation targets are computed only during training.

Verdict

The hierarchical structure of the SID is itself knowledge to distill, so distillation losses should follow that structure level by level instead of looking only at item probabilities at the last level.