RoFiGR: Coarse Routing, Fine Refinement via Level-Adaptive Distillation in Multimodal Generative Recommendation

Teams building generative retrieval, especially with SIDs built from image and text content, should read this. It puts its effort on the first one or two codes of the SID, because every later, finer choice happens inside the branch those codes pick. A generative recommender produces an item's SID one code at a time, and the first codes set the broad direction. One wrong branch early on and no amount of later precision can recover the relevant items. This is sharper in multimodal recommendation, where SIDs are usually built from image and text semantics, so the first codes may not reflect which items the same users tend to buy. RoFiGR aligns the first two levels of code vectors with frozen collaborative embeddings when it builds the SIDs. During generation, a sequential recommender acts as a teacher and tells the student model, at each step, how much probability each branch that leads to a real item should get. The first levels get a fixed, larger distillation weight, and the rest goes to whichever levels the teacher favors the correct branch more than the student does. On three Amazon datasets, RoFiGR beats the compared baselines on all five ranking metrics. Averaged over the 15 dataset and metric pairs, it is about 9.4% above the strongest baseline. Part of that number comes from decoding more streams, and we take it apart below.

What problem is it solving?

Start with an illustrative scenario. A music store builds 4-level SIDs from image and text semantics. Level 1 splits products into 64 broad groups, level 2 subdivides them, and levels 3 and 4 identify the exact product inside a fine group. A user recently bought an electric guitar and an effects pedal, and the next thing they will actually buy is a guitar cable.

By semantics, the cable and the guitar sit in two different groups. Having seen this history, the model will likely pick "string instruments" at the first step, and then however accurate the next three codes are, the cable is already out. This is what the paper calls a routing error: the early codes decide which items are still reachable, and the beam budget is limited.

The paper builds on MACRec, a multimodal generative recommender with separate text and image trees, and points out that its training treats all four levels the same, with one cross-entropy. But the four levels differ in difficulty and role. Level 1 has no prefix and must choose among dozens of groups; level 2 narrows within a group; levels 3 and 4 only identify items in a small branch. RoFiGR adds collaborative information in two places: when SIDs are built, and during generation.

Overall structure: three changes

The pipeline has these stages:

The key division of labor: collaborative signals group the first two levels, and the teacher tells the student at each step which branches deserve more probability. Each piece is below.

Why align the first two levels with collaborative embeddings?

MACRec quantizes SIDs from content, so its early groups reflect "looks alike, reads alike". RoFiGR adds a contrastive loss: sum the two code vectors an item picks at levels 1 and 2, pass the sum through a linear projection, and require it to be closest to the teacher's collaborative embedding for the same item and far from those of other items in the batch. The collaborative embedding is the sum of the teacher's normalized ID, text, and image representations of the item, so it carries "who is bought with whom".

The constraint applies only to the first two levels, and the last two still handle image and text details. The reason is simple: the first two levels decide the branch and should reflect buying behavior, while the last two tell items apart inside the branch, where content details help more. The paper also tried aligning all four levels, and it did worse.

The effect is visible on the discrete tree. The paper measures how often behavioral neighbors land in the same level-1 branch. Behavioral neighbors are item pairs that often co-occur within 1 to 3 positions in training histories, with at most 20 per item. On the text tree of the Instruments dataset, the rate is 13.89% without alignment and 31.87% with it. The effective number of level-1 branches stays above 62 (out of 64), with no empty branch. Effective branches are the exponential of the entropy of items over branches: exactly 64 when items spread evenly, smaller as they crowd into a few. This shows alignment groups behaviorally similar items without squeezing everything into a few groups.

How does the teacher tell the student each branch's probability?

Plain next-token training only tells the student which child branch is correct and says nothing about how the others compare. RoFiGR uses the teacher to fill this in. The teacher is a sequential recommender that gives every item a score. In training, given the correct prefix, the exponentiated scores of all items under each legal child branch are summed and normalized across legal children, giving the teacher's probability for each branch.

For example, using the paper's illustration: a level has three child branches. Items under the first have probabilities 0.30, 0.21, and 0.06; under the second, 0.25 and 0.08; under the third, 0.10. Summed, the teacher's branch probabilities are 0.57, 0.33, and 0.10. The student's output at this level is also a softmax over just these three branches, and the loss is the KL between the two distributions. Only legal branches that lead to real items take part, so the student is never asked to put probability on an empty branch.

Why allocate weights per level? With the same distillation weight on all four levels, the hardest level gets no extra help. RoFiGR's weight has two parts. The first is a fixed anchor: 0.10 for level 1, 0.075 for level 2, and 0 for the last two, so routing gets teacher signal from the start of training. The second is a floating budget of 0.015 shared by "demand": each level's demand is smoothed with an exponential moving average, raised to the power 1.5, and the 0.015 is split by each level's share.

Demand works like this: on the correct child branch, take how much higher the teacher's probability is than the student's, and how much higher it is than a catalog-size guess. Take the log of both gaps and keep the smaller; below 0 counts as 0. A worked example (illustrative): at level 1 the teacher gives the correct branch 0.30, the student only 0.10, and the branch holds 0.05 of the items. The two log gaps are about 1.10 and 1.79, so demand is 1.10. At level 3 the teacher gives 0.5 and the student already gives 0.6, so demand is 0. The floating budget then goes almost entirely to level 1, so its total weight is about 0.10 plus 0.015. A tiny constant sits in the denominator, so the amount actually handed out is 0.015 times "total demand divided by (total demand plus that constant)". When all demands are near 0, the floating part is also near 0 and only the anchors remain; as demand grows, the total approaches 0.015. The floating part is off for the first 1% of training steps so the moving average can build up. Keeping the smaller gap avoids spending weight where the teacher merely repeats the catalog distribution.

Which eight tasks does the student learn?

The student is a T5-style encoder-decoder shared by eight tasks, each with its own soft prompt. Six come from MACRec: generate the next item's text code from a text-code history, and the next image code from an image-code history; generate the next code of one modality from a history in the other (one task per direction); and, without any history, translate an item's text code to its image code and back (one per direction). RoFiGR adds two: each history item carries both its text and image codes, and the task generates the next item's text code or image code. The six history tasks get distillation; the two translation tasks do not.

Serving uses four of these task heads, each decoding one stream: text history to text code, image history to image code, bimodal history to text code, and bimodal history to image code. Each stream runs prefix-tree-constrained beam search (beam 20) to produce a full 4-code SID. Each complete SID maps to exactly one item, so it converts straight back to an item ID. Each stream normalizes item probabilities over its own candidates, then the four streams' probabilities are averaged, with 0 for an item missing from a stream, and items are ranked by the average.

How does the history tell "which item, which level" apart?

Each history item has 4 text codes and 4 image codes, so 20 items flatten into a long token sequence. From token distance alone, the model cannot tell whether two adjacent tokens belong to one item, or to the end of one item and the start of the next. RoFiGR adds two biases to the attention in the encoder's first layer: one for how many items apart two tokens are, and one for which SID level the attended token comes from. Concretely, each attention head has two small learned tables. One gives a number by item offset, clipped to 19 items either way, for 39 entries. The other gives a number by the attended token's level (4 levels) and the item offset between the two tokens, grouped into 11 bins: each offset within 2 items gets its own bin, and farther ones merge into 3 to 4, 5 to 8, and 9 or more. Both numbers are added, together with T5's existing token-distance bias, to the attention score of that SID-token pair, so the model can tell item boundaries apart and knows which level the information comes from. Six heads times (39 plus 4 times 11) makes 498 parameters in total, zero-initialized, and the decoder is unchanged.

What the results show

The three datasets are Amazon reviews for Instruments, Arts, and Games, with 6,250 to 13,839 items. Each user's last interaction is the test target, ranking is over the full catalog, and results average five seeds. NDCG@10:

On all three datasets, the HR@10 and NDCG@10 gains over the strongest baseline remain significant under five-seed paired tests with multiple-comparison correction. Note, though, that the main table runs RoFiGR with four streams while the baselines keep their native two.

The paper's own breakdown says more (NDCG@10 times 100):

With the same eight tasks and the same two streams, RoFiGR beats MACRec by about 1.07 points on Instruments, and that reflects the whole set of RoFiGR changes under the same tasks and decoding: first-two-level alignment, the structure-aware attention biases, and branch distillation. Going from two streams to four adds about 0.63 more. But when each of the four streams decodes with beam 20 and, before fusion, at most 40 unique items are picked round-robin by each stream's own rank, then renormalized and averaged, the result lands almost exactly on the two-stream result, which shows that much of the four-stream gain depends on the larger candidate set. Yet at an equal budget of 80 beam slots, four streams still beat two on all three datasets, so candidate count and fusion are not fully separated.

Every part contributes in the ablation. On Instruments, removing the branch teacher, replacing it with uniform probability over legal branches, or removing first-two-level alignment costs the most, around 0.9 to 1.0 points of NDCG@10. Removing only the floating weight drops validation NDCG@10 by about 6%. On Arts, greedy level-1 accuracy on the text tree rises from 17.73% for MACRec to 22.53%, which shows the gain does land on routing.

Limitations

First, the datasets are small. The largest has just over ten thousand items and about forty thousand test users, and level 1 has only 64 branches. Industrial catalogs run to tens of millions, where early grouping looks very different, and the paper has no large-scale test of whether alignment still pulls co-purchased items into one prefix.

Second, the teacher cache blows up with scale. The distillation target needs precomputed scores over the whole catalog for every training history; on Games the full FP16 cache is already 6.360 GiB, growing with histories times items. The paper tried keeping only the teacher's top K items: larger K brings the cached distribution closer to the full teacher, but student ranking does not improve monotonically, and K=2000 shows a non-monotonic result.

Third, the main-table gain mixes in decoding budget. Four streams raise mean latency on Instruments from 31.8 ms to 56.7 ms, and at equal candidate counts four streams and two are nearly tied. Reading only the 9.4% overstates the method itself.

How a recommendation team could try it

1. First check whether routing is the bottleneck. On your current generative retrieval, measure how often the correct prefix survives in the beam at each level. Only if most losses happen at levels 1 and 2 is this idea worth trying.
2. The cheapest step is first-two-level alignment. Use your existing collaborative embeddings as targets and add a contrastive loss on the sum of the level-1 and level-2 code vectors during SID training. Check whether co-purchased items share prefixes more often and whether any branch goes empty.
3. If you have a reliable sequential recommender, add branch distillation next. Precompute teacher scores offline, first add KL with fixed weights on the first two levels only, then add the demand-based floating weight, and compare each step's increment.
4. Evaluate at a fixed total candidate count, for example capping every system to the same number, so decoding budget is not credited to the method.

Is it worth reading?

Most generative recommendation errors are already made in the first two levels, so spending collaborative signal there pays off more than spreading it evenly over four levels.