UniPinRec: Unifying Generative Retrieval and Ranking at Pinterest Scale

Abstract: Modern recommendation systems predominantly train retrieval and ranking as separate models despite both increasingly relying on large transformers encoding the same user behavior data, duplicating parameters, compute, and serving cost. Prior work unifies the model architecture but not the full pipeline: input formats, training procedures, and serving stacks remain fragmented across stages. We present UniPinRec, which achieves full-stack unification of retrieval and ranking at Pinterest: one input format, one model, one training stage, deployed within existing serving infrastructure. A shared transformer encodes the user action sequence into candidate-independent representations that branch into retrieval (ANN dot-product) and ranking (cross-attention) via task-specific heads. Three ideas make this work: (1) Masked Action Modeling (MAM) eliminates interleaving, enabling weight sharing without doubling context length; (2) Blended training examples pair action sequences with feedview impression slates to satisfy both objectives jointly; (3) Cross-stage KV cache sharing reuses user-history computation from retrieval for ranking, reducing total FLOPs versus serving two independent models. Deployed in the Pinterest core surfaces, UniPinRec delivers approximately +1% online engagement lift while cutting end-to-end serving latency by 11.1% and lifting QPS by 63.6%. To our knowledge, this is the first full-stack unification of retrieval and ranking, covering inputs, model, training and serving, deployed in a production recommendation system.

• The paper introduces UniPinRec, a unified transformer that combines generative retrieval and action-based ranking through Masked Action Modeling, joint training, and shared candidate representations.
• The model matches production retrieval quality with Recall@10 of 0.77659 and improves offline Save Hit@3 by 14.7% over the existing TransAct V2 and DCNv2 ranking system.
• Cross-process KV-cache sharing, optimized attention, and mixed-precision serving reduce ranking forward latency by 3.92×, while online tests show 11.1% lower end-to-end latency and 63.6% higher QPS.

Motivation and problem statement

Modern recommendation systems are organized as multi-stage funnels in which candidate generation retrieves thousands of items from a corpus of millions and a downstream ranker rescores a smaller candidate set with richer features. Although both stages increasingly rely on large transformer backbones trained over the same user behavior data, they are typically trained and served independently, duplicating parameters, compute, and serving cost. Prior attempts at unification address only part of the problem: HSTU shares a transformer paradigm but requires different input formats (interleaved for ranking, non-interleaved for retrieval) so stages cannot share computation; OnePiece applies unified multi-task training but was deployed as either a retrieval or ranking model, never both from one model; generative approaches such as OneRec and OneRanker bypass the funnel entirely via autoregressive decoding of semantic IDs, at the cost of composability with existing candidate sources and per-stage operational control. UniPinRec targets full-stack unification — one input format, one model, one training stage, and deployment within existing serving infrastructure — and reports production deployment at Pinterest (2606.00422).

Method

UniPinRec builds on PinRec, a production generative retrieval model in which a causal decoder-only transformer encodes a user's positively interacted item sequence (with pre-trained item, multimodal, and search embeddings projected into an L2-normalized space) and is trained with a sampled softmax next-item loss using frequency-corrected similarity. The paper's design principle is compositional: it reuses ANN indexing, cross-attention-style ranking, and KV caching rather than introducing new mechanisms, making the unified model a drop-in replacement.

Three design choices enable unification:

• Masked Action Modeling (MAM). Instead of interleaving action tokens between items (which doubles context length and breaks input-format compatibility with retrieval), actions are concatenated with each item embedding along the feature axis and randomly masked with probability $p_{\text{mask}}$ during training. Past-position masking uses a dedicated [MASK] token, and candidate positions are always masked, guaranteeing no information leakage under causal attention. The attention pattern follows M-FALCON: candidates attend to all past positions but not to each other, reducing attention cost from $\mathcal{O}((n+k)^2)$ to $\mathcal{O}(n^2 + nk)$ and enabling KV-cache sharing. The model is trained jointly with the retrieval next-item loss and a per-action-type binary cross-entropy ranking loss over masked positions.
• Blended training data. Each training example pairs a past action sequence with a future feedview impression slate containing engagement labels and impression negatives. Non-engaged feedviews are subsampled (10% retained) to amplify sparse engagement classes, while evaluation uses an unbiased randomized replay dataset. An in-trainer bucket join on Ray over Iceberg tables hash-bucketed by user ID eliminates offline fanout duplication and makes sequence length and sampling ratios tunable at training time.
• Cross-stage KV-cache sharing. Retrieval encodes the $n$-token history once; the ranking stage runs an incremental decode-only pass over $k$ candidates attending to the cached history. Because retrieval and ranking run as separate OS processes, naive KV serialization would incur GPU–host–GPU roundtrips; instead, a pre-allocated GPU memory pool is shared via CUDA IPC handles exported at boot, so the ranking process reads retrieval's KV state without CPU-mediated transfer. Faiss search_and_reconstruct returns candidate IDs together with their stored embeddings, sparing the ranking stage a separate embedder forward pass and guaranteeing identical representations across stages.

Mixed-precision FP8 training (E4M3 forward, E5M2 backward) with Transformer Engine fused kernels further improves efficiency, at a conceded cost of a 0.5% offline metric drop.

Offline results (RQ1)

On Board More Ideas data, UniPinRec achieves a save Hit@3 of 0.10096 versus 0.088008 for the production TransAct V2 + DCNv2 ranker — a +14.7% relative improvement despite the unified model handling both tasks — and outperforms HSTU (0.097326) at matched context length. On retrieval, Recall@10 of 0.77659 marginally exceeds the production PinRec model (0.77486), establishing that joint training does not degrade retrieval quality, a prerequisite for drop-in deployment. Two ablations are informative: the ranking-only variant (no item loss) underperforms the jointly trained model, showing the retrieval objective improves ranking; and a sequentially fine-tuned PinRec (with the item embedder frozen to preserve Faiss-index compatibility) underperforms joint training, indicating that architectural compatibility alone is insufficient and that joint optimization — including a trainable embedder — is necessary.

Ablations on the masking ratio show that any masking beats none, with $p_{\text{mask}} = 0.2$ optimal and diminishing returns at 0.3. Scaling studies over depth (2–24 layers) and sequence length (256–2048 tokens) show consistent, unsaturated gains, though sequence length incurs quadratic attention cost and is identified as the more severe latency bottleneck.

Serving efficiency (RQ2)

The KV-reuse design converts ranking prefill into decode, yielding roughly a 2.4× speedup on its own. Stacked with flex attention (~1.3×) and FP8 plus CUDA graph capture (~1.25×), the best configuration achieves a 3.92× reduction in single ranking-pass forward latency (25.72 ms → 6.56 ms on L40S with $n{=}992$, $k{=}656$, $L{=}12$). In online serving under controlled GPU budget, the bf16 decode/flex/compile configuration delivers a −11.1% end-to-end latency reduction and +63.6% QPS lift; an FP8 variant reaches +109.1% QPS but at +6.7% latency because Transformer Engine shape constraints (leading dimension a multiple of 8) force request accumulation that penalizes the $S{=}1$ autoregressive retrieval steps. The paper notes this FP8 latency trade-off is unresolved and that bf16 was used for the online experiments.

Online A/B results (RQ3)

Production experiments deploy UniPinRec as an L0+L1 service: the unified model overfetches candidates via ANN, applies lightweight action-head ranking, and returns the refined top-K candidates to the existing TransAct V2 downstream ranker. Because the downstream ranker and blender remain in place, online lifts are attenuated relative to offline gains. On Board More Ideas, UniPinRec delivers +0.95% surface saves and +0.08% site-wide saves over the PinRec baseline. On Notifications, it delivers +0.91% push opens, +3.84% notification-surface saves, and +0.09% weekly active users, with the dormant-user push-open lift of +1.72% roughly double the all-user lift, suggesting personalized notification recommendations are particularly effective for reactivating lapsed users. The incremental L0+L1 rollout preserves composability with existing candidate generators and enables independent A/B testing and rollback per stage.

Limitations and open questions

The paper is explicit about several constraints. The unified model currently replaces only retrieval plus lightweight (L1) ranking; the L2 production ranker remains in place, with unification deferred due to dependencies in score calibration, training-data logging, and operational tooling — so the reported online lifts reflect a partially unified funnel. FP8 inference improves QPS but degrades latency under current kernel shape constraints, and the 0.5% offline metric cost of quantization is accepted rather than eliminated. Scaling results show sequence-length gains are bottlenecked by quadratic attention cost, leaving sequence compression as an explicitly open direction. The paper also concedes the deployment is confined to two surfaces, with extension to Search and Ads untested.

Conclusion

UniPinRec demonstrates that retrieval and ranking can be unified across the entire stack — input format, model, single-stage joint training, and shared serving infrastructure — in a production recommendation system. The key enabling ideas are Masked Action Modeling for format-compatible action supervision, blended action-plus-impression training examples, and cross-process KV-cache sharing that makes ranking a marginal-cost decode step. The result is a model that matches production retrieval recall, exceeds a dedicated ranker offline by +14.8% Hit@3, and delivers consistent online engagement gains while reducing end-to-end latency by 11.1% and lifting QPS by 63.6%. The authors frame the architecture as analogous to retrieval-augmented generation — a query-formulating backbone, a non-parametric ANN index, and a cached-context scorer — and identify iterative retrieval and learned retrieval-depth control as open questions carried forward from that literature.