TransRetrieval: Scaling Up Transformer-Based Retrieval for Industrial Recommendation

Zhifei Zheng; Yunfei Liu; Bin Liu; Qiren Zhu; Hanbing Liu; Ziru Xu; Han Zhu; Jian Xu; Qi Qi; Bo Zheng

Source-linked AI summary

TransRetrieval: Scaling Up Transformer-Based Retrieval for Industrial Recommendation

Zhifei Zheng, Yunfei Liu, Bin Liu, Qiren Zhu, Hanbing Liu, Ziru Xu, Han Zhu, Jian Xu, Qi Qi, Bo Zheng

arXiv:2608.25528v1Aug 26, 2026cs.IR

TL;DR

Recommendation retrieval has struggled to show the predictable scaling seen in other domains because heterogeneous features cause Transformer token-norm divergence. TransRetrieval addresses this with feature aggregation, target compression, and domain embeddings, achieving log-linear scaling and a 2.53% platform revenue lift.

Problem
Recommendation retrieval has limited evidence of predictable scaling, as conventional models show diminishing or negative returns when scaled up.
Method
TransRetrieval combines weighted average aggregation, target token compression, and position-style domain embeddings to support efficient cross-domain Transformer retrieval.
Results
TransRetrieval delivers log-linear scaling and state-of-the-art recall on industrial and KuaiRand datasets, with a 2.53% platform revenue lift online.
Takeaways & Limitations
The three designs unlock scaling laws for recommendation retrieval across computational budgets and domains while maintaining the production baseline’s end-to-end latency.
Takeaways & Limitations
Baselines were implemented within the authors’ framework with shared preprocessing, while retaining each baseline’s original feature-processing design.

Abstract

from arXiv · show

Applying scaling laws to recommendation retrieval is hindered by feature heterogeneity: naively stacking Transformer layers yields diminishing returns because heterogeneous fields produce severe token-norm divergence. We present TransRetrieval, a Transformer-based retrieval framework that scales with both computational budget and cross-domain data. The key enabler is (1) weighted average aggregation, which restores the homogeneous-token assumption Transformers rely on. Building on this, we introduce (2) target token compression that cuts per-candidate FLOPs by 85% while preserving cross-attention expressiveness, and (3) position-style domain embeddings that unify multiple domains at negligible additional cost, turning cross-domain data into a scaling asset. On a 40-billion-interaction industrial dataset and the public KuaiRand benchmark, scaling compute from 0.1 to 2 MFLOPs per target yields +19.3/+22.2 pt Recall@2000, confirming robust log-linear scaling. In online A/B tests, TransRetrieval lifts platform revenue by 2.53% under the same end-to-end latency constraint as the production baseline.

1 Introduction

TransRetrieval addresses the failure of recommendation retrieval models to scale predictably by adapting Transformers to heterogeneous features, massive candidate pools, and strict latency budgets. Its design enables log-linear scaling and delivers substantial offline and online gains in industrial deployment and public benchmarking.

Motivation: TransRetrieval targets retrieval over candidate pools of up to 10^8 items under latency budgets measured in tens of milliseconds, a setting where simply scaling traditional recommendation models produces diminishing or negative returns.Prior scaling efforts largely target ranking after retrieval and often flatten or pre-group heterogeneous features.
Motivation and approach: Weighted average aggregation resolves heterogeneous feature scale misalignment by decoupling embedding norms from feature cardinality and restoring Transformers’ homogeneous-token assumption.Recommendation inputs combine user profiles, behavioral sequences, and item attributes whose mapped token scales can otherwise structurally distort the embedding space.
Approach: The framework combines weighted aggregation with target token compression and position-style domain embeddings to scale across four business domains using standard operators without extra infrastructure.These designs address feature scale misalignment, computation-versus-latency constraints, and cross-domain deployment.
Results: 2.53% platform revenue lift is achieved in online A/B tests, while scaling compute from 0.1 to 2 MFLOPs per target improves Recall@2000 by +19.3 points on the industrial dataset and +22.2 points on KuaiRand.These results confirm robust log-linear scaling in retrieval under the deployment setting.

2 Related Work

Prior work identifies scaling inefficiency and heterogeneous features as central recommendation challenges, while existing approaches bypass rather than resolve this barrier. Related research also advances model-based retrieval, Transformer efficiency, and multi-domain parameter sharing, but leaves limitations in scoring capacity and cross-domain transfer.

Scaling in recommendation and the heterogeneous feature barrier: Wukong [24] attributes recommendation scaling inefficiency to embedding-table expansion that increases parameters without improving feature interaction, while HSTU [23] sacrifices feature-type distinctions by flattening fields into one sequence.These approaches bypass the heterogeneity barrier rather than resolving it.
Model-based retrieval and Transformer efficiency: NANN [3] uses a DNN scoring function to guide HNSW beam search, evaluating only a small corpus fraction, whereas BAR [12] adds bidding awareness but retains a small scoring network.BAR [12] serves as the production baseline and targets multi-stage consistency.
Multi-domain recommendation and cross-domain data scaling: MMoE and PLE share experts across domains, but undertrained gates limit sparse-domain access to shared capacity, while STAR [17] lacks demonstrated cross-domain transfer for data-scarce domains.STAR [17] combines shared and domain-specific components for unified serving.

3 TransRetrieval: Methodology

TransRetrieval stabilizes Transformer retrieval across heterogeneous features through weighted-average aggregation, compresses target tokens to enable deeper and wider models, and shares domain context without increasing attention complexity. Its architecture reuses user-side key/value caches across candidates and co-designs inference for production latency.

3.1 Architecture Overview: The architecture aggregates heterogeneous fields, compresses target features into one token, adds domain embeddings to all tokens, and reuses one user-side KV cache across candidates.A shared Pre-LN Transformer then scores each candidate against the cached user representation.
3.2 Weighted Average Aggregation: Weighted-average aggregation bounds token norms independently of field cardinality, correcting the up-to-10× user–target norm mismatch caused by weighted sums.Weighted sums accumulate scale with feature cardinality, causing high-norm tokens to dominate attention; weighted averages also avoid learned or estimated statistics.
3.3 Target Token Compression: Target compression maps concatenated target embeddings to one D-dimensional token using a 3-layer MLP, reducing per-candidate token cost and reallocating FLOPs to model width and depth.The compressor uses hidden sizes [8D, 4D, D], Parametric ReLU activations, and LayerNorm between layers.
3.4 Position-Style Domain Embeddings: Position-style domain embeddings inject domain identity by element-wise addition, preserving sequence length and attention complexity while sharing all Transformer parameters across domains.Unlike appended domain tokens, this design lets samples from every domain contribute gradients to the same attention and FFN parameters.
Production Inference: Inference co-design broadcasts one user KV cache without copying, reducing per-pass latency by 30%, while GPU-resident retrieval reduces retrieval latency by 89% versus CPU retrieval.The system also recomputes only mutated target representations asynchronously through a decoupled compressor sub-graph.

4 Experiments

The experiments evaluate TransRetrieval through four research questions covering baseline comparisons, scaling behavior, design-pillar contributions, and real-world deployment improvements.

Research questions: The experiments compare TransRetrieval with state-of-the-art baselines under similar computational budgets (RQ1; Sec. 4.2).
Research questions: They test whether TransRetrieval exhibits predictable scaling as model capacity increases (RQ2; Sec. 4.3).
Research questions: They measure how each of the three design pillars contributes to performance (RQ3; Sec. 4.4).
Research questions: They assess whether TransRetrieval delivers significant improvements in real-world deployment (RQ4; Sec. 4.5).

4.1 Experimental Protocol

Experiments evaluate TransRetrieval on industrial and public multi-domain recommendation datasets using Recall@2000. Baselines share preprocessing and tuned training settings, while dataset-specific negative sampling reflects their different scales.

Datasets: Experiments use two multi-domain recommendation datasets spanning display advertising and short-video recommendation.The Industrial dataset contains 40 billion interactions over 52 million advertised items across four anonymized domains; KuaiRand is publicly available.
Evaluation metric: Recall@2000 measures the proportion of ground-truth positive items appearing among the top-2000 retrieved candidates.The metric is abbreviated R@2000 when space is limited.
Baselines: Baselines are implemented within the same framework and preprocessing pipeline, while retaining their original feature-processing designs.The comparison includes Production Baseline, KuaiFormer, HSTU, and RankMixer.
Implementation: All models use one tuned hyperparameter set with AdamW, learning rate 1e-3, weight decay 3e-5, and the production baseline’s pairwise LTR loss.Experiments run on NVIDIA H20 GPU clusters using PyTorch and RecIS [28].
Implementation: Negative sampling uses 5 negatives per positive for Industrial and 200 for KuaiRand to reflect their different training-data scales.Industrial contains 40 B interactions versus 152 M for KuaiRand.

4.2 Overall Performance (RQ1)

TransRetrieval outperforms the production and scalable retrieval baselines across domains under comparable per-candidate compute budgets. Increasing its compute further improves Overall Recall@2000, establishing state-of-the-art performance.

Evaluation setup: Table 2 compares TransRetrieval configurations with the Production Baseline, HSTU, KuaiFormer, and RankMixer across both datasets and all domains.Configurations are denoted by embedding dimension and Transformer-layer count, while RankMixer uses token-mixing blocks.
Overall comparison: 0.603 vs. 0.576 Overall Recall@2000: TransRetrieval-64D3L surpasses the Production Baseline across all domains with 35% fewer FLOPs.The configurations use 0.45 and 0.69 MFLOPs′ per candidate, respectively.
Scaling performance: 0.657 Overall Recall@2000: TransRetrieval-128D5L improves by +5.4 pt over TransRetrieval-64D3L and sets a new state-of-the-art against HSTU and RankMixer.The larger configuration uses 1.91 MFLOPs′, while HSTU and RankMixer operate at similar budgets of approximately 1.7–2.1 MFLOPs′.

4.3 Scaling Behavior (RQ2)

TransRetrieval exhibits predictable log-linear scaling across Industrial and KuaiRand as backbone capacity increases. Its largest configuration achieves substantial recall gains and outperforms a much larger RankMixer model with far less compute.

4.3 Scaling Behavior (RQ2): Recall@2000 rises from 0.464 to 0.657 on Industrial (+19.3 pt) and from 0.251 to 0.473 on KuaiRand (+22.2 pt) from 32D1L to 128D-5L, with R2=0.82/0.88.Figure 2 shows clean, consistent log-linear scaling across both datasets.
4.3 Scaling Behavior (RQ2): TransRetrieval-128D5L outperforms RankMixer-768D-16T by 12.6 pt on Industrial while using 1.91 versus 150.99 MFLOPs.The comparison uses nearly two orders of magnitude less compute for TransRetrieval.
4.3 Scaling Behavior (RQ2): KuaiFormer shows near-flat scaling, with Recall@2000 increasing only from 0.485 to 0.495 when doubling dimension, indicating a two-tower inner-product bottleneck.Enlarging the encoder alone does not overcome this bottleneck.

4.4 Ablation Study (RQ3)

Ablations show that weighted average aggregation is essential for heterogeneous features, while target-token compression enables substantial cost savings with limited recall loss. Multi-domain modeling also requires explicit domain context under latency constraints.

Feature aggregation: Weighted Average achieves the best Recall@2000 across domains, outperforming alternative feature aggregation strategies in the same 64D-5L configuration.Weighted Sum, LayerNorm plus Weighted Sum, and Max Pooling score 0.608, 0.593, and 0.571 versus 0.618 for Weighted Average.
Target token compression: ∼85% lower MFLOPs′ costs compression from 8 target tokens to 1, with Recall@2000 falling only 2.1 points.Within 64D-3L, cost decreases from 3.04 to 0.45 MFLOPs′ while Recall@2000 decreases from 0.624 to 0.603.
Target token compression: 128D-5L with 1-token compression achieves higher Recall@2000 at lower cost than 64D-3L with 8 tokens.Recall@2000 rises from 0.624 to 0.657 while cost falls from 3.04 to 1.91 MFLOPs′, confirming compression-then-scale-up as the better operating point.
Multi-domain modeling: Mixing multi-domain data without domain tokens causes severe degradation, while one domain token provides limited improvement and adds computational overhead.Table 6 evaluates these strategies with 64D-3L and reports user-side FLOPs to quantify sequence-expansion overhead.

4.5 Online A/B Test Results (RQ4)

In a month-long Alibaba display-advertising A/B test on 5% of production traffic, TransRetrieval improved platform revenue and revenue per mille over the production baseline while matching its end-to-end latency and maintaining stable ecosystem metrics.

Online A/B Test Results (RQ4): TransRetrieval delivered these gains at identical end-to-end latency to the production baseline, serving 230 QPS with P99 latency under 40 ms.The test ran for one month on 5% of production traffic, with retrieval as the only changed model variable.
Online A/B Test Results (RQ4): 2.53% increase in Platform Revenue and 1.28% lift in Revenue Per Mille were achieved over the highly optimized Production Baseline.The daily gains remained consistent throughout the month, with no systematic decay.
Online A/B Test Results (RQ4): User experience metrics remained stable, while advertisers’ Return on Investment stayed comparable at +0.06%.The higher per-target computational cost reduced per-engine QPS to 230 versus 300, and horizontal scaling absorbed the difference.

5 Discussion and Conclusion

TransRetrieval’s three designs unlock log-linear retrieval scaling, state-of-the-art recall, and a 2.53% platform revenue lift in month-long online A/B tests. The discussion emphasizes conditioning inputs before redesign, reallocating savings from compression, and unifying domains for scaling.

Lessons learned: Weighted-average input conditioning, rather than architectural novelty, restored homogeneous token norms and enabled the scaling breakthrough.The fix is described as parameter-free and one-line.
Lessons learned: Compressing 8 target tokens to 1 frees 85% FLOPs, and reinvesting savings into 64D-3L→128D-5L yields +3.3 pt Recall at lower total cost.The lesson frames compression as reallocation rather than sacrifice.
Lessons learned: Domain unification is presented as a scaling asset alongside input conditioning and target-token compression.The three designs together unlock scaling laws for the retrieval stage.
Conclusion: 2.53% platform revenue lift accompanied TransRetrieval’s log-linear scaling, state-of-the-art recall, and R^2 =0.82/0.88 across industrial and KuaiRand benchmarks.The results use a 40-billion-interaction industrial dataset and the public KuaiRand benchmark, with month-long online A/B tests.