PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient GPU Retrieval for Semantic Search

Dhritiman Das, Chujie Zheng, Ronak Kaoshik, Pratik Dixit, Vishal Shah, Yanbo Li, Jiahao Xu, Manika Agarwal, Chinmay Naik, Lingyu Zhang, Chetan Bhole, Chirag Bhanuprasad Mehta, Meng Zheng, Puneet Singh Ahluwalia, Shirisha Singh, Ping Jin, Manas Apte, Gokulraj Mohanasundaram, Tugrul Bingol, Raghavan Muthuregunathan, Fedor Borisyuk

arXiv:2608.28968v1cs.AIcs.LG

TL;DR

Semantic People Search needs first-stage retrieval that satisfies multi-faceted relevance constraints under a bottleneck-oriented GR policy. The paper introduces category-aligned segment and tagged multi-vector retrieval with a two-stage FP8/FP16 GPU architecture, achieving improved offline and live relevance while retaining near-full-FP16 recall at scale.

  • Problem

    Semantic People Search must satisfy every required attribute, match the LLM judge’s category-level relevance definition, and raise the candidate-recall ceiling imposed by first-stage retrieval.

  • Method

    The framework uses category-supervised segments with GR aggregation within tagged document slots, maximizes across slots, and serves them through FP8 coarse retrieval followed by FP16 re-ranking.

  • Results

    The live A/B test increased exploratory Precision@10 from 63.7% to 79.0% and navigational Precision@1 from 65.5% to 74.7%, while the GPU path recovered 99.6–99.8% of full-FP16 recall.

  • Takeaways & Limitations

    The design aligns representation learning, serving-time scoring, and the product GR policy while providing a practical high-throughput GPU retrieval path.

  • Takeaways & Limitations

    Stage 1 uses only slot 0, so documents strong in other slots can be excluded at finite depth; the candidate-loss gap is validated empirically rather than closed analytically.

Abstract

from arXiv · show

Semantic Search on LinkedIn must retrieve relevant profiles from a corpus of hundreds of millions in response to natural-language queries such as "a fintech founder in Berlin who worked in payments." The deployed relevance policy is bottleneck-oriented: every active non-negotiable facet must be satisfied, and a pre-existing LLM Graded Relevance (GR) judge operationalizes this through a fixed min/median aggregation over facet grades. Cosine similarity instead averages evidence, letting a strong match on one facet mask failure on another, capping the recall of the first-stage (L0) retriever. We present a policy-aligned retrieval framework: embeddings are partitioned into eight category-supervised segments whose scores follow the same min/median rule at serving time; for multi-vector retrieval, this segment score is computed independently per tagged document slot and maximized across slots. A lightweight single-slot Stage-1 scorer generates high-recall candidates, while scale-invariant relative-norm gating keeps category activation consistent across training, evaluation, and serving. On 21K held-out queries, this representation improves offline relevance over a matched-capacity baseline, with gains broadly distributed across facet combinations. We serve this framework with a two-stage GPU architecture: an FP8 coarse ranker scores the full corpus, increasing per-shard capacity by 71% and Stage-1 matmul throughput by 36%, then an FP16 stage exactly re-ranks an oversampled candidate set, recovering 99.6-99.8% of full-FP16 recall at over 500 QPS per shard replica. In a member-randomized A/B test, exploratory-query Precision@10 under the unchanged GR judge rises from 63.7% to 79.0% and navigational Precision@1 from 65.5% to 74.7%, with a blinded human evaluation independently confirming the Precision@10 gain.

1 Introduction

The paper studies efficient first-stage GPU retrieval for Semantic People Search, where natural-language queries describe profiles through multiple constraints. It introduces a policy-aligned representation and deployable GPU architecture, with offline evaluation, scaling experiments, and live validation.

  • Motivation: Semantic People Search retrieves relevant LinkedIn profiles for expressive queries containing spatial, temporal, and categorical constraints.The L0 stage filters hundreds of millions of profiles to a few thousand candidates before downstream re-ranking.
  • Contributions: A GR-aligned segment multi-vector representation applies the frozen policy aggregation within tagged document slots and maximizes the score across slots.Relative-norm gating and a lightweight single-slot Stage-1 scorer support parity-consistent candidate generation.
  • Evaluation: The paper evaluates the framework through offline experiments, GPU scaling measurements, and live member-randomized A/B testing under the unchanged GR judge.The section overview assigns retrieval-framework analysis to §4, offline evaluation to §5, and system and deployment evaluation to §6.

2 Related Work

Prior work spans shared-vector semantic retrieval, GPU serving and quantization, multi-vector late interaction, and multi-objective ranking. This paper differs by aligning slot-level multi-vector scoring with an explicit multi-category relevance policy rather than relying on a single scalar similarity.

  • Embedding-Based Retrieval and Learned Similarity: Bi-encoders retrieve in a shared vector space, while industrial systems commonly use fixed cosine or dot-product scoring.Other approaches learn similarity functions or observe that contrastive cosine does not capture absolute relevance.
  • GPU Serving and Quantization: GPU retrieval systems use architectures such as LiNR, Bloom filtering, approximate nearest-neighbor methods, and billion-scale GPU search to address serving cost and index footprint.Approximate methods reduce index footprint at the cost of recall, while this system extends LinkedIn’s GPU-native LiNR architecture.
  • Multi-Vector Retrieval and Adaptive Dimensioning: ColBERT-style systems use token-level late interaction, whereas this framework uses multiple document vectors at the slot level.A shared encoder with distinct prefix tags produces independent representations, and the final score applies GR within each slot before taking the maximum.
  • Multi-Objective Ranking, Distillation, and LLM Judging: The framework distills an LLM judge’s explicit multi-category grading policy into segment representations rather than compressing the judge into one scalar.Its training combines contrastive objectives with retrieval-geometry regularizers.

3 Problem Definition

The problem is to retrieve profiles satisfying every required attribute while matching the product’s category-level relevance definition and improving the first-stage recall ceiling. The proposed setup replaces flat similarity with GR-aligned segment and multi-vector scoring within a two-stage GPU retrieval pipeline.

  • Challenges: Semantic People Search must satisfy every required attribute, match the LLM judge’s category-level relevance definition, and raise the first-stage candidate-recall ceiling.These are identified as challenges C1, C2, and C3.
  • Retrieval Framework: The framework partitions embeddings into category-specific subspaces and applies inherited min/median aggregation within each slot, followed by a maximum across slots.The encoder is learned, while serving applies the aggregation rules without additional parameters.
  • Serving Pipeline: The system uses FP8 coarse retrieval over the full corpus and FP16 scoring over an oversampled candidate set.This two-stage retrieve-then-rank pipeline preserves the same candidate-retrieval interface while replacing flat similarity with GR-aligned scoring.
  • Model Configurations: The baseline is a shared 4B-parameter Siamese bi-encoder with 512-dimensional query and profile embeddings, while the treatment uses three native 256-dimensional document slots.The treatment is trained and partitioned directly at 256 dimensions rather than obtained by truncating a trained 512-dimensional segment representation.
  • Relevance Policy: The GR judge grades query-profile pairs from 0–4 across eight categories and defines relevance as a final grade of at least 3.The final grade takes the minimum over six non-negotiable categories, the median over two negotiable categories, and then the minimum of those group scores.

4 Modeling Architecture

The architecture aligns retrieval with the bottleneck-oriented GR policy by learning category-specific segments and applying the same GR aggregation across training, evaluation, and serving. A lightweight slot-0 Stage 1 generates candidates, while Stage 2 applies exact GR-then-max scoring across tagged slots.

  • GR-Aligned Segment Representation: The GR-aligned scorer takes the minimum over active non-negotiable segments, aggregates negotiable segments, then applies the final group minimum.This bottleneck score prevents strong matches on some facets from compensating for failure on another non-negotiable facet.
  • Two-Stage Retrieval: Stage 1 approximates the non-negotiable minimum over slot 0 to generate candidates efficiently, while the final multi-slot score maximizes the complete GR score across slots.The minimum is an upper bound on the full GR-aligned score and is therefore better aligned with Stage 2 than a flat sum.
  • Two-Stage Retrieval: At M=8000, flat 256-d Stage 1 recovers 99.8% of final FP16 top-1000 recall, while the non-negotiable minimum reaches 100% recall by M≈1500.The experiment supports compensating for the flat approximation through higher over-fetch depth.
  • Multi-Vector Retrieval: Slot 0 is the sole corpus-wide embedding for Stage 1, whereas Stage 2 gathers and scores all K tagged slots for each forwarded candidate.For K=3, the document stores three embeddings total rather than one global vector plus three additional slots.
  • Multi-Vector Retrieval: Distinct prefix tags generate independent slot representations, and diversity, balancing, and slot-0 recall losses encourage specialization without sacrificing globally informative candidate geometry.The shared backbone re-encodes the same profile separately for each tagged slot.
  • GR-Aligned Segment Representation: The model partitions each slot embedding into eight GR-supervised segments, using 64 dimensions per segment at 512-d and 32 at 256-d.Each segment is trained with category-specific supervision and independently L2-normalized for cosine scoring.
  • Training and Serving Consistency: Relative-norm gating preserves category activation behavior, while identical GR-then-max scoring across training, offline evaluation, and Stage 2 avoids scorer drift.The combined score is used as the document logit in the primary contrastive objective and in serving.

5 Evaluation Design and Offline Retriever Quality

The evaluation measures offline relevance, candidate recall, and capacity–quality trade-offs on 21,237 uniformly sampled queries, while testing whether policy-aligned segmentation improves retrieval across facet combinations.

  • Evaluation design: The offline protocol evaluates 21,237 category-uniform queries over the full corpus using each model’s own scoring function.The evaluation set is independently constructed from training data, and metrics are macro-averaged per query.
  • Evaluation design: Reported relevance metrics include P@1, P@10, RS-NDCG@10, Capped R@10, and PMR@10, with lower PMR@10 indicating fewer clearly poor matches.RS-NDCG@10 is retrieved-set based, while Capped R@10 uses relevant profiles within each model’s top-100.
  • Candidate recall: At M=8000, flat 256-dimensional dot-product Stage 1 recovers 99.8% of the final FP16 top-1000, while the segment minimum reaches 100% by M=1500.Candidate recall measures how much of the final FP16 top-1000 survives Stage 1; these results support the selected over-fetch depth without guaranteeing recall.
  • Offline relevance: The standalone 512×1 segment model raises RS-NDCG@10 from 0.8704 to 0.8820 and P@10 from 0.5880 to 0.6377 versus single-vector L0.PMR@10 decreases from 0.2987 to 0.2809 in the same comparison.
  • Offline relevance: At matched 512×3 capacity, GR-aligned segmentation improves all five metrics, while the 256×3 setting raises P@10 from 0.6379 to 0.6462 and lowers PMR@10 from 0.2806 to 0.2717.Across five independent runs, both three-slot models significantly outperform single-vector L0 under paired tests (p< 0.05).
  • Capacity–quality trade-offs: The 512×4 plain multi-vector model is best offline, but 256×3 is the best observed trade-off because larger 512-dimensional configurations exceed the co-resident serving memory budget.The configurations do not form a complete factorial grid, and fixed-candidate checks do not test end-to-end Stage-1 parity.
  • Offline relevance: The combined model improves 19 of 21 facet combinations over tagged-slot MaxSim, while the standalone segment model improves 20 of 21 combinations.The reported macro gains are +0.0514 and +0.2574, respectively, with regressions limited to specified combinations.

6 Deployable GPU Retrieval

The deployable system combines full-corpus FP8 coarse retrieval with exact FP16 reranking to scale semantic search while preserving recall. A brokered multi-shard architecture and targeted optimizations support live evaluation, where relevance improved under unchanged judging.

  • 6.2 Quantization: 71% higher document capacity per shard and ≈36% faster matmul follow FP16-to-FP8 quantization on Hopper-class hardware.Capacity rose from 35M to 60M documents per shard.
  • 6.2 Two-Stage Retrieval: FP8 coarse retrieval can reorder near-neighbors, so the system combines full-corpus FP8 filtering with FP16 reranking of oversampled candidates.Single-pass FP8 top-1000 recall was 89–93% across 15M–50M-document tests.
  • 6.2.1 Architecture: Stage 1 scans only slot 0 with a flat 256-dimensional FP8 dot product, applies attribute filters, and selects M=8000 candidates for K=1000 results.This scorer is an empirically high-recall generator rather than a recall-preserving bound for the final multi-slot maximum.
  • 6.2.1 Architecture: Stage 2 gathers all FP16 slots, applies the GR-aligned multi-slot score, and returns the top K candidates; transfer and exact scoring scale with rerank depth M rather than corpus size N.The two-pass path retains 99.8% recall in the selected configuration.
  • 6.2 Trade-offs: The two-pass design recovers 99.6–99.8% of exact-FP16 recall while supporting 60M documents, but multi-vector host-RAM storage caps practical shards at 35M documents.The reported 1.36× figure is Stage-1 matmul throughput, not end-to-end throughput.
  • 6.3 Performance: At 400 QPS, increasing rerank depth from 2,000 to 20,000 raises P95 latency from 72.8 to 158.7 ms, exposing the latency–recall trade-off.Under a 150 ms P95 target and 0.5% gRPC error limit, sustainable capacity ranges from 600 to 300 QPS.
  • 6.4 Live-Traffic Evaluation: Exploratory P@10 rises from 63.7% to 79.0%, while navigational P@1 rises from 65.5% to 74.7% in the member-randomized live test.A blinded human evaluation independently confirms the significant Human P@10 gain.
  • 6.6 Lessons Learned: The practical lessons favor FP8 for coarse scanning, larger over-fetch depth for the cheap scorer, and memory-aware multi-vector configurations for serving.FastTopK, Bloom filtering, and derived features reduce overhead without changing the relevance scorer or A/B treatment.

7 Conclusion

The paper concludes with an end-to-end retrieval design aligned to the product GR policy, combining category-aware representations with a two-stage GPU serving path. Offline gains translate into online metric lifts, while the FP8/FP16 system demonstrates practical scale.

  • Conclusion: The end-to-end design aligns learned representations and GPU serving with the product GR definition.Category-supervised segments preserve required facets, tagged slots capture complementary profile aspects, and slot 0 supports candidate generation.
  • Conclusion: The live A/B test confirms that offline gains translate into online metric lifts.A blinded 50-query evaluation independently confirms the P@10 improvement, while the corresponding P@1 gain is directional rather than significant.
  • Conclusion: The FP8/FP16 path increases per-shard capacity by 71%, recovers 99.6–99.8% of full-FP16 recall, and sustains over 500 QPS per GPU shard replica.These measurements support the system’s at-scale practicality.
  • Conclusion: Stage 1 uses a flat FP8 slot-0 dot product as a serving-time approximation, while Stage 2 applies the complete GR-aligned scorer to all slots.The two stages preserve the efficiency–fidelity division of the deployed pipeline.
  • Conclusion: The regularization configuration combines diversity, balancing, and slot-0 recall terms to support distinct, usable slots and stand-alone candidate generation.The evaluated coefficients are α=0.01, β=0.01, and δ=0.01.

A.4 Training-objective component analysis

The training-objective analysis isolates slot-specialization regularizers under a fixed 256 × 3 architecture and shared training setup. Diversity improves top-rank precision, while combining diversity and balancing produces the strongest overall ablation profile.

  • Experimental Setup: Every ablation row keeps the 256 × 3 architecture, tagged slots, data, optimization budget, evaluation protocol, and shared base objectives fixed.Only diversity and balancing regularizers are ablated.
  • Ablation Results: Diversity regularization raises P@1 from 0.7558 to 0.7730 and P@10 from 0.6244 to 0.6462 relative to the shared-base configuration.The passage identifies diversity as providing most of the top-rank precision gain.
  • Ablation Results: Balancing alone achieves the highest Capped R@10 at 0.8785 but slightly worsens PMR@10 to 0.2809 versus 0.278.The metrics capture different aspects of retrieval quality and poor-match behavior.
  • Ablation Results: Combining diversity and balancing preserves P@10 at 0.6462, lowers PMR@10 to 0.2717, and yields the highest RS-NDCG@10 of 0.8835 and P@1 of 0.7731.Lower PMR@10 is better.
  • Metric Definitions: RS-NDCG@10 measures only ordering quality within each model’s retrieved top-10, not corpus-level retrieval quality or comparison against a shared judged pool.Its ideal list is constructed from the same model’s retrieved set.

C Full Category-Combination Results

Across category combinations, GR-aligned segmentation improves retrieval broadly rather than only on frequent query types, and matched-capacity multi-vector results show gains beyond vector count alone.

  • 20 of 21 GR category combinations improve over the matched single-vector L0 baseline, or 21 of 22 groups including navigational queries.The group-macro improvement is +0.2574 over GR combinations and +0.2639 including navigational queries.
  • The standalone segment model outperforms the matched single-vector L0 baseline in all seven single-category groups.
  • 19 of 21 GR category combinations improve with the combined segment multi-vector model over plain tagged-slot MaxSim, or 20 of 22 groups including navigational queries.The group-macro gain is +0.0514 over GR combinations and +0.0522 including navigational queries.
  • Matched-capacity results rule out additional vector count as the sole explanation for aggregate improvement across single-category and conjunctive multi-category queries.

D Bloom Filter False-Positive Analysis

The Bloom-filter analysis quantifies how false positives scale with large connection sets and shows that semantic similarity suppresses many operationally projected errors.

  • 0.1% FPP on a 600M-document corpus yields 600,000 false positives, potentially dominating the top-K pool when only 500,000 true positives exist.FPP is the chance that a non-matching candidate incorrectly tests positive.
  • The expected candidate volume needed to obtain 500 true matches is ⌈500/(1−ρ)⌉.The calculation uses the Bloom filter false-positive probability ρ.
  • Semantic embedding similarity acts as an implicit second filter, producing substantially fewer top-K false positives than conservative uniform-distribution projections.
  • A 2 MB bitmap keeps FPP below 1% for users with up to approximately 1.6M second-degree connections while adding minimal payload.
  • At a 150 ms P95 latency target and 0.5% gRPC error-rate limit, sustainable capacity declines from 600 QPS at S=2,000 to 300 QPS at S=20,000.Intermediate depths of 4,000, 8,000, 10,000, and 16,000 support 550, 500, 450, and 350 QPS, respectively.

F Fast TopK Implementation and Full Benchmarks

The system uses GPU-native filtering and top-K selection to reduce serving overhead while preserving the compact candidate-processing path.

  • Two FP16 passes identify the K-th-score threshold and collect exactly K selected value-index pairs on-device.Boundary values use an overflow buffer, and only the compact K-element result is sorted.
  • FP16 top-K selection requires two full input passes instead of four for FP32 in the common case.A slow-path fallback rescans the input when pathological ties overflow the boundary buffer.
  • Speedups reach 7.6× over torch.topk and 6.3× over the chunked baseline.
  • The dual-role kernel consolidates three requests into one, yielding a 3× efficiency gain at neutral quality.On a 35M-document shard, the attribute-filter-heavy mix sustains 500 QPS at 123.1 ms P95.
  • The dual-role kernel computes derived features alongside filtering and writes them only for candidates surviving the filter.

H Blinded Human Relevance Evaluation

A blinded evaluation on 50 paired queries independently supports the system’s Human P@10 improvement, while the Human P@1 increase remains directional at this sample size.

  • The evaluation pooled and randomized deduplicated query-document pairs without model identity or rank information.Human P@1 and P@10 were computed per query and averaged across 50 queries.
  • Human P@10 improves by 6.0 points, from 0.61 to 0.67, with statistical significance (p=0.004; 95% CI [+2.0, +10.0] points).P@10 is the primary confirmatory result of the blinded check.
  • Human P@1 improves by 8.0 points, from 0.68 to 0.76, but does not reach significance at n=50 (p=0.424; 95% CI [−2.1, +18.0] points).The result is directionally consistent with the online P@1 lift.
Loading 2608.28968v1…