Source-linked AI summary
Scaling the Vocabulary of Non-autoregressive Models for Efficient Generative Retrieval
Ravisri Valluri, Akash Kumar Mohankumar, Kushal Dave, Amit Singh, Jian Jiao, Manik Varma, Gaurav Sinha
TL;DR
Autoregressive generative retrieval faces high inference latency and cost, while standard non-autoregressive models can lose retrieval performance because they do not capture target-token dependencies. PIXAR adds phrase tokens to the NAR target vocabulary and optimizes inference for the expanded vocabulary. It outperforms standard NAR models on MS MARCO and NQ while retaining latency benefits over autoregressive models.
Problem
Method
PIXAR expands the NAR target vocabulary with document-identifier phrases and uses a shortlist of tokens to make inference with the large vocabulary more efficient.
Source: §1 · p. 2, §4.2 · p. 5, §5.3 · p. 9Results
PIXAR outperforms standard NAR retrieval on MS MARCO and NQ, with substantial gains while retaining latency speedups over autoregressive models.
Source: §5.2 · p. 7Takeaways & Limitations
The results position PIXAR as a candidate for latency-sensitive applications such as real-time search and recommendation systems.
Source: §6 · p. 9NAR models assume conditional independence among target tokens, an assumption that prior studies indicate rarely holds for real-world data.
Source: §3 · p. 3
Abstract
from arXiv · showhide
Generative Retrieval introduces a new approach to Information Retrieval by reframing it as a constrained generation task, leveraging recent advancements in Autoregressive (AR) language models. However, AR-based Generative Retrieval methods suffer from high inference latency and cost compared to traditional dense retrieval techniques, limiting their practical applicability. This paper investigates fully Non-autoregressive (NAR) language models as a more efficient alternative for generative retrieval. While standard NAR models alleviate latency and cost concerns, they exhibit a significant drop in retrieval performance (compared to AR models) due to their inability to capture dependencies between target tokens. To address this, we question the conventional choice of limiting the target token space to solely words or sub-words. We propose PIXAR, a novel approach that expands the target vocabulary of NAR models to include multi-word entities and common phrases (up to 5 million tokens), thereby reducing token dependencies. PIXAR employs inference optimization strategies to maintain low inference latency despite the significantly larger vocabulary. Our results demonstrate that PIXAR achieves a relative improvement of 31.0% in MRR@10 on MS MARCO and 23.2% in Hits@5 on Natural Questions compared to standard NAR models with similar latency and cost. Furthermore, online A/B experiments on a large commercial search engine show that PIXAR increases ad clicks by 5.08% and revenue by 4.02%.
1 Introduction
PIXAR addresses the latency of autoregressive generative retrieval with a non-autoregressive model that adds phrase tokens to reduce token dependencies. Its reported benchmark and commercial-search results show gains over standard NAR retrieval while retaining similar inference latency.
1 Introduction
- Autoregressive generative retrieval has high inference complexity because it generates tokens sequentially, while NAR models generate identifier tokens simultaneously but capture their dependencies less effectively.Source: §1 · pp. 1–2
- The paper asks whether phrase tokens can improve NAR retrieval accuracy and how inference latency can be reduced for the enlarged vocabulary without sacrificing accuracy.Source: §1 · p. 2
1.1 Our Contributions
- PIXAR expands the NAR target vocabulary with phrases from document identifiers and uses training and inference optimizations to mitigate the larger vocabulary’s computational burden.The paper presents its architecture in Figure 1.Source: §1.1 · p. 2, §1 · p. 2
- 24.0% higher MRR@10 on MS MARCO and 23.2% higher Hits@5 on NQ than standard NAR models, with similar inference latency.These are the reported benchmark improvements for PIXAR.Source: §1.1 · p. 2
- 5.08% more ad clicks and 4.02% higher revenue were observed in A/B testing on a large commercial search engine.Source: §1.1 · p. 2
2 Related Work
Prior generative retrieval methods use varied document representations but often rely on costly autoregressive generation. PIXAR instead uses phrase tokens for NAR retrieval and a query-level shortlist to address the computational cost of large-vocabulary softmax.
Generative retrieval
- Generative retrieval represents documents as text, numeric IDs, titles, substrings, pseudo-queries, or combinations, but reliance on autoregressive models brings high inference latency and computational cost.Source: §2 · p. 2
Non-autoregressive Models
- NAR methods generate tokens independently and in parallel, but can mix tokens from multiple valid outputs; PIXAR instead predicts phrases to reduce the number of independent predictions.Source: §2 · pp. 2–3
Efficient Softmax
- Softmax over a large vocabulary is a computational bottleneck, especially because NAR models perform multiple softmax operations in one forward pass.Source: §2 · p. 3
- PIXAR uses a dedicated shortlist embedding to narrow target tokens for an entire query, reducing latency while maintaining strong retrieval performance.Source: §2 · p. 3
Large Vocabulary
- Unlike non-parametric language models that predict from open-vocabulary n-grams and phrases using dense embeddings, this work learns classifier weights for an extended target vocabulary within NAR.Source: §2 · p. 3
3 Preliminaries
The paper uses document identifiers as generation targets, specifically pseudo-queries generated from documents. NAR models generate those identifiers in parallel under a conditional-independence assumption that can limit retrieval performance.
3 Preliminaries
- The paper uses pseudo-queries generated from documents as document identifiers, or docids, within the generative-retrieval framework.Source: §3 · p. 3
- NAR models generate all docid tokens in parallel by assuming conditional independence among target tokens, enabling efficient inference.Source: §3 · p. 3
- Prior studies indicate that target-token independence rarely holds for real-world data, and NAR models can consequently lose retrieval performance relative to autoregressive models.Source: §3 · p. 3
4 Proposed Work: PIXAR
PIXAR expands NAR document-identifier vocabularies with phrases and uses query-conditioned token shortlists to manage inference over that larger vocabulary.
4.2 PIXAR Pipeline
- PIXAR selects k query-matched token subsets, unions them into a shortlist, reranks candidates at each output position, and generates docids with trie-constrained beam search.The shortlist should contain tokens for every position in the output docid.Source: §4 · p. 3, §4.2 · p. 5
4.1 Vocabulary Construction
- The target vocabulary is designed to encode docids compactly, retain sufficiently frequent tokens for training, and include common phrases while respecting word boundaries.The two-stage construction filters substring candidates, then iteratively tokenizes and scores test vocabularies to remove tokens that compress the docid set poorly.Source: §4.1 · p. 4
4.2 PIXAR Pipeline
- PIXAR training combines token prediction, shortlist-coverage, and self-normalization losses, then trains cluster vectors to favor docid tokens for the most-aligned vector.The shortlist-coverage term encourages x0(q) to select subsets containing the output docid's tokens.Source: §4.2 · pp. 5–6
- Self-normalization enables faster probability estimates for large vocabularies by avoiding the denominator sum over the entire target vocabulary.The estimates are used for position-wise token scoring.Source: §4.2 · p. 5
- With |V| scaled to 5 million, tuned m, r, and k preserve high retrieval quality with negligible inference-time impact.The shortlist contains at most rk tokens, and selecting clusters can be efficient when m is much smaller than |V|.Source: §4.2 · p. 6
5 Experiments & Results
The experiments benchmark PIXAR against retrieval baselines, test its components, and evaluate sponsored search. PIXAR improves retrieval over standard NAR models while retaining efficient inference, with positive sponsored-search outcomes.
5 Experiments & Results
- Experiments benchmark PIXAR against AR and NAR retrieval methods, ablate model components and inference methods, and test sponsored search.Source: §5 · p. 6
5.2 Results
- PIXAR beats CLOVERv2 across MS MARCO metrics and MINDER by 22.5% in MRR at 10; on NQ it gains around 13-23% over CLOVERv2, shrinking LTRGR’s Recall@100 gap from 16.3% to 5.1%.On MS MARCO, PIXAR also trails LTRGR by 1.5 absolute points in MRR@10, a 5.8% relative difference.Source: Table 1 · p. 7, §5.2 · p. 7
5.1 Experimental Setup
- Evaluation spans MS MARCO and NQ passage retrieval plus SponsoredSearch-1B; latency is measured at batch size 1 on an Nvidia T4.Source: §5.1 · pp. 6–7
5.3 Ablations
- At 128K vocabulary size, the tokenizer ablation compares DeBERTa BPE, custom sub-word BPE, Unigram, and phrase-based tokenization on MS MARCO.Source: §5.3 · p. 8, Table 2 · p. 8
- Increasing vocabulary from 1 million to 5 million tokens improves MS MARCO Recall@5 by 7.7%, from 35.7 to 38.5, with gains persisting beyond 1 million.The scaling study uses the phrase-based tokenizer and full softmax across MS MARCO and NQ.Source: §5.3 · p. 8, Table 3 · p. 8
- PIXAR’s inference pipeline is 10.6x faster than full softmax and 2.1x faster than Fast Vocabulary Projection, with MRR@10 within 0.82% and Recall@100 within 0.85% of full softmax.Its latency is 21% higher than CLOVERv2 despite a vocabulary 39x larger.Source: Table 4 · p. 8, §5.3 · pp. 8–9
5.4 Further Analysis
- PIXAR’s phrase tokens represent multi-word entities and shorten target sequences: mean length falls from 10.98 to 4.05 and the 99th percentile from 18 to 9 versus CLOVERv2.Despite fewer tokens, PIXAR tends to generate longer outputs in words, addressing NAR models’ tendency to produce short outputs.Source: Table 5 · p. 9, §5.4 · p. 9
5.5 Application to Sponsored search
- On SponsoredSearch-1B, PIXAR raises P@100 from 23.5% to 29.1%; commercial A/B testing reports 5.08% more ad clicks and 4.02% more revenue.Advertiser bid keywords serve as ad docids in this sponsored-search setting.Source: §5.5 · p. 9
6 Conclusion
PIXAR uses phrase-level tokens in an expanded target vocabulary for non-autoregressive generative retrieval. The paper reports that it bridges the performance gap with state-of-the-art autoregressive methods while retaining non-autoregressive efficiency.
6 Conclusion
- PIXAR is a non-autoregressive retrieval approach that uses phrase-level tokens in an expanded target vocabulary.The paper describes it as a novel approach to generative retrieval.Source: §6 · p. 9
- PIXAR bridges the performance gap with state-of-the-art autoregressive methods while retaining the efficiency of non-autoregressive models.Source: §6 · p. 9
- PIXAR’s speed advantage makes it a promising candidate for latency-sensitive applications such as real-time search and recommendation systems.Source: §6 · p. 9
A.1 Model details
PIXAR and CLOVERv2 use a pretrained DeBERTa encoder initialization. PIXAR’s language modeling head is initialized from scratch.
A.1 Model details
- PIXAR and CLOVERv2 are initialized with the pretrained DeBERTa encoder.The specified checkpoint is microsoft/deberta-v3-base.Source: §A.1 · p. 13
- CLOVERv2 uses the provided 128K DeBERTa vocabulary for both input and target tokens.Source: §A.1 · p. 13
- PIXAR’s language modeling head is initialized from scratch.Source: §A.1 · p. 13
A.2 Document identifiers
The experiments use MINDER pseudo queries as document identifiers. Training data is augmented with pseudo-query pairs associated with the same passage.
A.2 Document identifiers
- MINDER pseudo queries serve as document identifiers in the experiments.Source: §A.2 · p. 13
- The identifier sets contain around 80 million unique pseudo queries for Natural Questions passages and about 170 million for MS MARCO passages.Source: §A.2 · p. 13
- Training data is augmented by sampling up to 20 pseudo queries per passage and mapping them to other pseudo queries for that passage.The sampled pseudo queries are added as questions.Source: §A.2 · p. 13
A.3 Training details
All models share the stated optimizer, learning-rate schedule, warmup, batch size, and loss-scaling settings. Training duration differs by dataset.
A.3 Training details
- All models use a learning rate of 5 × 10^-5, 1000 warmup steps, and an effective batch size of 6400.Source: §A.3 · p. 13
- The Adam optimizer is paired with a linear-decay learning-rate scheduler; λ2 is 0.25 and λ3 is 1.0.λ2 scales shortlist loss and λ3 scales self-normalization loss.Source: §A.3 · p. 13
- Models train for 5 epochs on MSMARCO and 10 epochs on Natural Questions.Source: §A.3 · p. 13
A.4 Compute
The appendix specifies compute hardware and training time, shortlisting settings, phrase-vocabulary construction parameters, and sequence-length comparisons across vocabulary sizes.
A.4 Compute
- Models with 5M-token vocabularies were trained on 8 Nvidia H100 GPUs, other vocabulary sizes on 16 AMD Mi200 GPUs, and inference on an NVIDIA Tesla T4 GPU.Training took 1–2 days depending on vocabulary size.Source: §A.4 · p. 13
A.5 Shortlisting module
- The shortlisting module uses m=4096, r=20000, and k=5.Source: §A.5 · p. 13
A.6 Vocabulary construction
- 5M-token PIXAR vocabularies are constructed separately for MS MARCO and NQ, with candidate phrases occurring at least 20 times and strict mode preventing minor phrase variants from receiving multiple tokens.TokenMonster binaries build vocabularies from each dataset’s full document-identifier set.Source: §A.6 · p. 14
A.7 Sequence Lengths and Target Sentence Lengths
- Tokenizing document identifiers with larger vocabularies nearly halves required tokens versus baseline; phrase-based PIXAR targets are longer and more relevant while using fewer tokens.Table 6 reports mean and 99th-percentile sequence lengths across vocabulary sizes.Source: Table 6 · p. 14, §A.7 · p. 14