Memory Layer: Train the In-Model Cache for Recommendation Models

2026-07-27

AI Summary

This work addresses the inconsistency in item representations between training and serving phases in recommender systems caused by divergent caching mechanisms, which degrades model performance and increases operational overhead. The authors propose a Memory Layer that embeds a key-value item embedding cache directly within the model architecture. By integrating this layer with joint training and an always-on fallback embedding mechanism, they establish the first end-to-end self-contained update pipeline that writes during training and reads during serving. This approach unifies the representation pathway across both phases, ensuring full coverage and high freshness. Deployed in Instagram Reels, the method achieves 100% prediction coverage, improves embedding freshness from minutes to under 20 seconds, reduces the normalized entropy gap between training and serving by 86%, doubles recall for newly added content, boosts cold-start engagement by 5–6%, and lowers training-to-deployment costs by 30%.

Abstract

Early ranking stages in recommendation systems precompute item embeddings and cache them in-model for scoring within strict latency constraints. Because this cache exists only at serving time, outside the training loop, training and serving use different item representations, a structural discrepancy that limits quality and adds operational fragility. We show that co-designing the training and serving paths removes this representation discrepancy at its source. We introduce the memory layer, an in-model key-value embedding cache co-trained with the model: the item tower writes embeddings during training and the model reads them at serving, one source of truth for item representations by construction. Always-on embeddings cover items not yet cached, so every item receives a prediction, and the design consolidates three separate trainer-to-predictor update paths into a single self-contained pipeline. Deployed in production on Instagram Reels, the memory layer raises prediction coverage from 96% to 100%, improves embedding freshness from $O(5\text{ min})$ to $O(20\text{ s})$, and narrows the training-serving Normalized Entropy (NE) gap by up to 86%, yielding over $2\times$ recall for the freshest content and a 5-6% cold start engagement lift. Because embeddings are produced during training, the system needs no separate bulk-evaluation or publish-time recomputation, cutting training-and-publish computational cost by 30% at neutral serving computational cost.

Problem

• training-serving discrepancy
• in-model cache
• item embeddings
• recommendation systems
• representation consistency

Innovation

• memory layer
• in-model cache
• training-serving consistency
• embedding freshness
• recommendation systems