Collaborative Memory Augmentation for Generative Recommendation
Enze Liu, Zhen Tian, Wayne Xin Zhao
KDD 2026 Research Track
cs.IR
2026-08-02
OMEGA compresses global user behavior into a memory bank, retrieves entries by sequence and target similarity, and fuses them through a gate, lifting TIGER 13-24% across metrics.
Generative recommendation (GR) turns next-item prediction into sequence-to-sequence generation: each item becomes a semantic ID, and the model autoregressively emits the target ID. Existing GR frameworks share a weakness: they model only a single user's history and compress all cross-user collaborative patterns into fixed model parameters. That is an information bottleneck, and when the current user's behavior is sparse the model cannot dynamically borrow signals from similar users. Borrowing the in-context-learning idea from LLMs, the authors argue GR should also be able to consult users with similar trajectories.
OMEGA is a model-agnostic add-on (validated on T5-style encoder-decoders, with TIGER and Pctx as backbones). Three stages:
First, latent context compression. C learnable query tokens (Q-Former style) compress a variable-length user sequence into a fixed-length representation, cutting storage from l×d to C×d; C=2 already matches the original GR. The compressed representations of all training sequences form a memory bank M.
Second, target-aware retrieval. Computing similarity directly on the compressed representations fails from representation collapse, so a lightweight sequential retriever (HSTU, 2 layers) produces the query representation. The retrieval score combines sequence-level and target-level similarity: s equals 1 minus α times the sequence cosine, plus α times the target-query cosine, which stops the system from dragging back irrelevant memories based on sequence lookalikes alone.
Third, context-aware integration. The top-K (K=10) memories first pass a gate driven by retrieval score (low-scoring memories are pushed toward zero) for denoising, then a cross-attention layer (causal mask removed) plus an FFN with residuals fuses them into the user context for the decoder. Training uses two-stage fine-tuning: freeze the backbone and align the attention module first, then optimize end-to-end jointly.
Three Amazon 2023 datasets (Instrument, Scientific, Game), leave-one-out evaluation:
| Backbone | Typical gain with OMEGA |
| TIGER | +13.5% to 23.7% (across metrics) |
| Pctx | +5.0% to 15.4% |
Pctx plus OMEGA reaches state of the art on most metrics, beating HSTU, LETTER, LIGER, ETEGRec, and ActionPiece. In ablations, random retrieval degrades to nearly the backbone, showing the model ignores irrelevant memory and retrieval quality is decisive; mean-pooled representations in place of compression also do markedly worse. This is not distillation: a TIGER plus HSTU ensemble still loses to OMEGA-HSTU. On efficiency, memory retrieval takes under 2% of inference time, inference is about 1% slower, and trainable parameters grow about 7%; using only 1% of the memory bank already approaches full-bank performance.
It gives generative recommendation back the cross-user collaborative signal it was missing, and does so with a retrieval-augmentation approach (like RAG, but for recommendation). For anyone working on GR, latent compression makes the memory bank storage-feasible, gated target-aware retrieval keeps the signal clean enough, and two-stage fine-tuning stops the add-on from corrupting the backbone. Each piece is reusable on its own. Being model-agnostic is a plus, since it does not lock you into one specific GR implementation.
All experiments are on relatively small academic datasets (Amazon subsets, MovieLens-1M); unlike STEPS or GRACE there is no validation at industrial billion-user scale, so whether the gains survive production is unknown. The paper has no explicit limitations section, and the authors mention only in the complexity analysis that the memory bank may need clustering-based pruning at very large scale. The fact that 1% of the bank is enough also hints that the memory may encode only coarse, generic patterns, which raises the question of how much real personalization it adds. The retriever is a separate HSTU, which means another small model and another training pipeline, so it is not free engineering-wise.