Source-linked AI summary
SetMIR: Multi-Interest Retrieval as Set Prediction
Xiaodong Liu, Congfei Zhang, Hsiang-wei Chao, Siman Wang, Xiao Bai, Tong Zhao, Jingxiao Ma, Wen Zhang, Zhe Liu, Shantanu Aggarwal, Di Huang, William Leach, Yunzhi Zhou, Yajun Wang, Jinchao Li, Yu Zhang
TL;DR
Multi-interest retrieval must represent diverse user behavior without collapsing interests or issuing unnecessary ANN queries. SetMIR treats retrieval as set prediction with learnable queries, Hungarian matching, and presence-based adaptive dispatch. On Snap DPA, it outperforms four learned multi-interest retrievers while using fewer queries and improves production retrieval-source outcomes.
Problem
Existing multi-interest retrieval methods can suffer interest collapse from argmax assignment and static dispatch from issuing ANN searches for all K embeddings.
Method
SetMIR uses a fixed bank of K learnable queries, Hungarian one-to-one matching with absence targets, and presence gating with query-level NMS for serving.
Results
SetMIR outperforms four learned multi-interest retrievers on every metric while issuing 33% fewer ANN queries per request, and lifts production CVR by 3.11% over the existing retrieval mix.
Takeaways & Limitations
SetMIR provides a request-dependent set of active ANN queries while retaining a single-vector item index and improving offline and production retrieval outcomes.
Abstract
from arXiv · showhide
Embedding-based retrieval is at the core of industrial recommender systems, but a single user embedding is often too limited to capture a user's diverse interests. Multi-interest retrieval addresses this by using multiple user embeddings, yet existing methods still suffer from two issues: interest collapse, where different embeddings learn the same interest, and static dispatch, where serving uses a fixed retrieval budget even when some embeddings are unnecessary. We propose SetMIR, which treats multi-interest retrieval as a set prediction problem. SetMIR encodes a user's behavior history with a transformer and uses K learnable queries to decode a set of user interests, each producing a retrieval embedding and a presence score. During training, Hungarian matching assigns targets to queries one-to-one, so matched queries learn distinct interests and the presence head learns which queries are active. At serving time, SetMIR uses presence scores and query-level Non-Maximum Suppression (NMS) to issue only active, non-redundant ANN queries. On Snap's Dynamic Product Ads (DPA) data, SetMIR outperforms four learned multi-interest retrievers on every metric while issuing 33% fewer ANN queries per request. Deployed as a new retrieval source in the DPA production stack, SetMIR lifts overall CVR by 3.1%, while lifting CTR by 44% and CVR by 51% over the item-to-item retrieval source with the same item embeddings, ANN index, and retrieval quota.
1 Introduction
SetMIR reframes multi-interest retrieval as set prediction to address interest collapse during training and static dispatch during serving. Learnable queries, one-to-one matching, presence gating, and query-level NMS produce a variable active query set.
- A single user embedding can miss narrower or short-lived interests when behavior spans diverse product categories.
- Existing multi-interest methods face interest collapse because argmax assignment can concentrate targets on a few queries and leave others without gradients.
- Static dispatch issues ANN searches from all K embeddings, wasting retrieval budget on redundant or low-quality queries without per-query quality signals.
- SetMIR uses K learnable queries and Hungarian matching to assign engaged targets one-to-one, giving matched queries retrieval gradients and unmatched queries explicit absence targets.
- Presence gating followed by query-level NMS activates a request-dependent subset of queries for adaptive, non-redundant ANN retrieval.
2 Related Work
Related work largely constructs multiple user embeddings but commonly retains argmax assignment and fixed ANN dispatch. SetMIR adapts learnable-query set prediction to retrieval while preserving single-vector item representations and existing ANN indices.
- Prior multi-interest recommenders construct multiple user embeddings through capsule routing, self-attention, prototypes, virtual-kernel experts, clustering, or learnable transformer queries.
- Many existing methods use argmax target assignment, allowing multiple targets to share a query while unmatched queries receive no training signal.
- Most multi-interest retrievers statically issue a fixed number of ANN queries regardless of how many interests are active.
- Set prediction produces a variable-size output set from a fixed bank of learnable queries, with Hungarian assignment and explicit absence targets.
- Unlike ColBERT-style retrieval, SetMIR keeps each item as one vector and places multiplicity only on the user side, reusing single-vector ANN indices.
3 Method
SetMIR frames multi-interest retrieval as set prediction: learnable queries produce distinct interest embeddings and presence scores, with Hungarian matching during training and adaptive query selection at serving.
- Problem formulation: SetMIR predicts a variable-size interest set from K learnable queries, each producing an interest embedding and presence score.The target-set size varies across users and is unknown at inference time.
- Architecture: A transformer encoder processes behavior history, while a query decoder uses self-attention and cross-attention to produce per-query outputs.Query self-attention allows queries to exchange information before reading the encoded history, encouraging specialization.
- Set-prediction training: A presence head estimates whether each query is active, while a margin diversity loss penalizes overly similar active query pairs.The diversity penalty applies only when active-query cosine similarity exceeds margin m = 0.3.
- Set-prediction training: Hungarian matching assigns queries to targets one-to-one, giving matched queries retrieval supervision and unmatched queries explicit absence targets.The assignment combines embedding similarity with the predicted presence logit.
- Inference: At serving, presence gating and query-level NMS reduce the fixed query bank to active, non-redundant ANN queries whose results are merged by maximum score.Each surviving query issues an ANN search, with per-query depth set approximately to preserve the overall retrieval depth.
4.1 Experimental Setup
SetMIR is evaluated on Snap DPA engagement data using held-out users and standardized comparisons against learned multi-interest retrievers under matched retrieval conditions.
- Dataset: Snap DPA provides engagement logs across views, swipe-ups, add-to-carts, and purchases from a catalog containing hundreds of millions of products.Targets use high-intent events from the last 3-day window, deduplicated by product id and capped at W_max = 15.
- Comparison protocol: Table 1 compares SetMIR with four learned multi-interest methods using the same data, frozen item embeddings, K = 10, optimizer, batch size, and retrieval budget.Baseline methods dispatch all K queries, while each metric is normalized so SetMIR equals 100%.
- Evaluation pool: Offline evaluation uses a held-out test set of around 1M users and a target pool of around 5.3M products.A separate validation split selects the presence and NMS thresholds and checkpoints.
- Metrics: Metrics include recall at N in {1, 5, 10, 100}, HR@100, NDCG@100, and MRR.These metrics evaluate retrieval quality at multiple list depths and ranking criteria.
4.2 Offline Experiments
SetMIR outperforms four learned multi-interest retrievers across offline metrics while reducing ANN queries through adaptive dispatch. Ablations show that Hungarian assignment is the dominant contributor to retrieval quality, while presence gating, diversity, and query-count choices provide additional benefits.
- Comparison with learned multi-interest baselines: SetMIR is strongest on every metric, while issuing 6.70 ANN queries per request versus 10 for the baselines, a 33% smaller budget.ComiRec-SA reaches 67.4% of SetMIR’s R@1, 80.0% of its MRR, and 99.9% of its R@100.
- Comparison with learned multi-interest baselines: 52.0% of SetMIR’s R@10 is reached by KuaiFormer, despite its comparable transformer architecture, capacity, data, and item embeddings.The comparison isolates differences in query supervision, with the SetMIR argmax variant reported separately.
- Component ablation: 36.5% of R@10 and 43.6% of MRR are lost when Hungarian matching is replaced by per-target argmax assignment.Mean presence-active queries fall from 6.70 to 2.16 of 10, illustrating interest collapse when targets compete for the same query.
- Component ablation: 91.4% and 93.2% of SetMIR’s R@10 remain without presence and diversity losses, respectively, showing smaller but separate effects.Without diversity, mean pairwise query cosine rises from 0.298 to 0.321; without presence supervision, thresholding admits 8.23 rather than 6.70 queries and still loses recall.
- K sweep: 155% of single-query R@10 and 153% of single-query R@100 are reached with K=10, while K=7 already reaches 99.8% of K=10’s R@10.The largest gain occurs from K=1 to K=3, and average active queries grow sublinearly: 4.72 at K=5, 6.70 at K=10, and 7.59 at K=15.
- Inference-time gating and dispatch reduction: Presence gating and NMS reduce static 10-query dispatch to 6.70 queries while retaining 99.2% of static R@10.Gating removes 23% of calls at 99.4% of R@10, and NMS removes a further 1.03 calls at 99.2% of R@10.
4.3 Online Experiment
SetMIR was evaluated in a single live A/B experiment in Snap’s DPA production stack, with the downstream ranker held fixed. The experiment reports both group-level lifts over the existing stack and source-level comparisons against item-to-item retrieval under matched conditions.
- Experiment setup: A single live A/B experiment added SetMIR as a retrieval source in Snap’s DPA production stack while holding the downstream ranker fixed.The experiment ran on a representative live-traffic slice for about one week; exact traffic fractions, sample sizes, and significance thresholds were omitted per Snap’s policy.
- Experiment setup: SetMIR used its default inference policy and an ANN index built from the same pretrained content-embedding model.
- Source comparison: SetMIR and item-to-item retrieval differed only in how query vectors were produced, with frozen item embeddings, ANN index, and per-source retrieval quota.Because item-to-item already dispatches several seeds per request, the comparison involved two ways of selecting a small set of query vectors.
- Results: +44% CTR and +51% CVR were achieved by SetMIR over item-to-item retrieval in source-level attribution within the treatment arm.The downstream ranker was held fixed across both analyses.
4.4 Qualitative Analysis
A qualitative example shows how SetMIR separates a diverse behavior history into specialized retrieval interests, whereas a single-query baseline concentrates on one category.
- Qualitative analysis: SetMIR’s five active queries specialize in distinct interest clusters and collectively cover most of a user’s varied interest mix.The history spans shirts and tops, pants, dresses, shoes, skirts, handbags, and accessories.
- Qualitative analysis: The single-query K=1 baseline returns only pants despite the user’s history spanning seven product categories.
5 Conclusion
SetMIR combines set prediction with adaptive query selection to address interest collapse during training and static dispatch during serving. On Snap DPA, it improves retrieval quality while reducing ANN-query volume, and it delivers strong online gains against item-to-item retrieval.
- Conclusion: SetMIR uses K learnable interest queries, Hungarian matching against engaged targets, and a presence head supervised by matched versus unmatched queries.
- Conclusion: Presence gating and query-level NMS convert the fixed query bank into a request-dependent set of active ANN queries served by a single-vector item index.
- Conclusion: 33% fewer ANN queries per request were issued while SetMIR outperformed four learned multi-interest retrievers on every metric.
- Conclusion: +3.11% CVR was achieved over the existing retrieval mix, alongside +44% CTR and +51% CVR over the item-to-item source.