Source-linked AI summary
Poly-encoders: Transformer Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring
Samuel Humeau, Kurt Shuster, Marie-Anne Lachaux, Jason Weston
TL;DR
Multi-sentence candidate scoring must balance accuracy with inference speed, since Cross-encoders are accurate but costly and Bi-encoders are faster but less expressive. The paper introduces Poly-encoders and task-matched pre-training, finding a favorable accuracy-speed trade-off across four tasks.
Problem
Multi-sentence scoring requires high prediction quality while remaining fast enough to score many candidates.
Method
The paper introduces Poly-encoders, which use learned global attention features, and compares Bi-, Cross-, and Poly-encoders with multiple pre-training strategies.
Results
Poly-encoders are more accurate than Bi-encoders and faster than Cross-encoders, while task-similar pre-training improves results across evaluated tasks.
Takeaways & Limitations
Poly-encoders provide an improved accuracy-speed trade-off and support fast real-time inference by retaining precomputable candidate representations.
Takeaways & Limitations
Cross-encoders cannot precompute candidate embeddings and are impractical for large candidate sets or real-time inference.
Abstract
from arXiv · showhide
The use of deep pre-trained bidirectional transformers has led to remarkable progress in a number of applications (Devlin et al., 2018). For tasks that make pairwise comparisons between sequences, matching a given input with a corresponding label, two approaches are common: Cross-encoders performing full self-attention over the pair and Bi-encoders encoding the pair separately. The former often performs better, but is too slow for practical use. In this work, we develop a new transformer architecture, the Poly-encoder, that learns global rather than token level self-attention features. We perform a detailed comparison of all three approaches, including what pre-training and fine-tuning strategies work best. We show our models achieve state-of-the-art results on three existing tasks; that Poly-encoders are faster than Cross-encoders and more accurate than Bi-encoders; and that the best results are obtained by pre-training on large datasets similar to the downstream tasks.
1 Introduction
Multi-sentence scoring tasks require balancing prediction quality with speed, motivating architectures that improve Cross-encoder efficiency without sacrificing accuracy. The paper introduces Poly-encoders and task-matched pre-training, evaluating them across dialogue and information-retrieval datasets.
- Multi-sentence scoring evaluates candidate labels against an input context, with prediction quality and speed as joint requirements.Scoring many candidates can be prohibitively slow.
- Bi-encoders encode inputs separately, whereas Cross-encoders apply full self-attention over the input-candidate pair.Cross-encoders generally attain higher accuracy, but their computation is expensive.
- The Poly-encoder introduces learned attention over global features, improving performance over Bi-encoders while providing large speed gains over Cross-encoders.
- Task-similar abundant pre-training data brings significant gains across the architecture choices and downstream tasks evaluated.
- The study compares architectures and pre-training strategies on four dialogue and information-retrieval datasets, obtaining a new state-of-the-art on all four.The paper also provides practical implementations for real-time use.
2 Related Work
Related work distinguishes Bi-encoders, which separately map inputs and candidates into a shared feature space, from richer Cross-encoders that model arbitrary interactions between concatenated sequences. Cross-encoders perform better in prior comparisons but incur substantial computational cost.
- Bi-encoders separately map an input and candidate into a common feature space and score their similarity.Typical scoring uses a dot product, cosine, or parameterized non-linearity.
- Cross-encoders concatenate the input and candidate, allowing a nonlinear function to model unrestricted dependencies between them.Transformer Cross-encoders apply self-attention at every layer, enabling word-level cross-sequence interactions.
- Prior BERT-based comparisons found that Cross-encoders perform better than Bi-encoders on dialogue and action tasks.
- Cross-encoder performance gains come at a steep computational cost.
3 Tasks
The paper evaluates sentence selection in dialogue and article search in information retrieval across four datasets. These benchmarks range from small candidate-choice evaluations to retrieval against 10,000 or more documents.
- The evaluation covers sentence selection in dialogue and article search in information retrieval, including ConvAI2, DSTC7, Ubuntu V2, and Wikipedia Article Search.
- Dialogue tasks: ConvAI2 asks models to select the correct response from 20 choices using dialogue history and persona information.The final system would retrieve from over 100k training utterances, but common evaluation avoids this for speed.
- Dialogue tasks: DSTC7 contains Ubuntu technical-support conversations, while Ubuntu V2 is a similar but larger corpus with many existing results.
- Information retrieval: Wikipedia Article Search uses a sentence as a query and ranks its source article against 10,000 other articles.The setup is intended to mimic web search for relevant documents.
- Table 1 summarizes the four datasets and their statistics.
4 Methods
The paper compares Bi-, Cross-, and Poly-encoders built on pre-trained transformers, alongside alternative pre-training and fine-tuning strategies for multi-sentence scoring. The Poly-encoder combines cached candidate representations with richer context-candidate interactions while reducing the cost of full Cross-encoder attention.
- Transformers and Pre-training Strategies: The models use BERT-base-sized transformers, while pre-training compares Wikipedia/Toronto Books data with Reddit data more similar to dialogue tasks.The architecture has 12 layers, 12 attention heads, and hidden size 768; Reddit pre-training uses 174 million [INPUT, LABEL] examples versus 150 million from Wikipedia and Toronto Books.
- Transformers and Pre-training Strategies: Pre-training combines masked language modeling with next-sentence prediction for Wikipedia/Toronto Books or next-utterance prediction for Reddit.The candidate is correct half the time and randomly sampled half the time, with alternating batches of the language-modeling and prediction objectives.
- Bi-encoder: Bi-encoders encode context and candidate separately, reduce each sequence to one vector, and score candidates by their dot product.Candidate representations can be cached, and in-batch labels can be reused as negatives during training.
- Cross-encoder: Cross-encoders concatenate context and candidate in one transformer, enabling candidate-sensitive self-attention but preventing candidate precomputation.This richer interaction mechanism permits selecting useful input features per candidate, while each candidate requires a full forward pass at inference.
- Poly-encoder: Poly-encoders cache one vector per candidate, represent the context with m global features, and apply candidate-query attention over those features before dot-product scoring.The global features are extracted using learned context codes; because m < N and cross-attention occurs only at the top layer, inference is faster than full Cross-encoder self-attention.
- Poly-encoder: The Poly-encoder’s final candidate score is the dot product yctxt · ycandi, retaining Bi-encoder-style scoring after the additional attention mechanism.The architecture therefore supports a precomputed response cache while adding richer interactions before scoring.
5 Experiments
Experiments compare Bi-, Poly-, and Cross-encoders across dialogue and information-retrieval tasks, evaluating prediction quality, pre-training, and speed. Poly-encoders improve over Bi-encoders while remaining substantially faster than Cross-encoders, and task-specific pre-training improves results.
- Experimental setup: Experiments measure Recall@k and mean reciprocal rank across four dialogue and information-retrieval tasks.Recall is reported as R@k/C, where C is the number of candidates.
- Bi-encoders and Cross-encoders: Cross-encoders outperform Bi-encoders on the three dialogue tasks after BERT fine-tuning, including 84.8% R@1 versus 81.7% on ConvAI2.On DSTC7, the corresponding scores are 67.4% and 66.8% R@1.
- Bi-encoders and Cross-encoders: BERT fine-tuning is omitted for Wikipedia IR because test-set contamination cannot be guaranteed and Cross-encoders are too slow for 10k candidates.This limits direct comparison with BERT-based fine-tuning on that evaluation setup.
- Poly-encoders: 68.9% R1 on DSTC7 is achieved by a BERT-pretrained Poly-encoder with 360 intermediate context codes, exceeding both the Cross-encoder's 67.4% and Bi-encoder's 66.8%.Poly-encoders outperform Bi-encoders on all tasks, with larger code sizes generally producing larger improvements.
- Domain-specific Pre-training: Reddit pre-training yields further state-of-the-art performance over BERT across all three dialogue tasks and all three architectures.Comparable Toronto Books + Wikipedia pre-training results indicate that the pre-training dataset, rather than another training detail, drives the difference.
- Inference Speed: At 100k candidates, Poly-encoders are 5-6x slower than Bi-encoders but remain tractable, whereas Cross-encoders are 2 orders of magnitude slower and intractable for real-time inference.Poly-encoders are also 3-4x faster to train than Cross-encoders and similar in training time to Bi-encoders.
6 Conclusion
The paper introduces Poly-encoders and related pre-training strategies for candidate-selection tasks. Poly-encoders improve the accuracy-speed trade-off by combining candidate-representation precomputation with context attention, while task-related pre-training improves performance.
- Conclusion: Poly-encoders attend over the context using the label candidate while preserving precomputable candidate representations for fast real-time inference.The method is presented as applicable beyond dialogue to tasks that score sets of candidates.
- Conclusion: Poly-encoders are more accurate than Bi-encoders and far faster than Cross-encoders, which are impractical for real-time use.The paper characterizes this as an improved trade-off between accuracy and speed.
- Conclusion: Pre-training strategies more closely related to downstream tasks bring strong improvements, with Reddit pre-training outperforming BERT across three dialogue datasets and architectures.The conclusion reports this result for all three model architectures tested.
A Training Time
Table 6 reports model training time in hours across the three datasets and four model types considered.
- The comparison covers three datasets and four model types.
- Training time is reported in hours.
B Reduction layer in Bi-encoder
The study compares several reduction functions for compressing Bi-encoder transformer outputs into a single representation on ConvAI2.
- The compared reductions use the first BERT output, averages of the first 16 or 64 outputs, or all outputs except the first.
- The evaluation metric is Recall@1/20.
- Table 7 evaluates Bi-encoder reductions on the ConvAI2 validation set.
C Alternative Choices for Context Vectors
The Poly-encoder can construct context vectors using several choices, including learnt codes, selected transformer outputs, or combinations of outputs, with performance and inference time compared across variants.
- Alternative Choices for Context Vectors: Learnt codes extract context representations by attending over all transformer outputs.
- Alternative Choices for Context Vectors: The First-m variant uses the first m transformer outputs, while alternatives use the last m outputs or concatenate them with the first output.
- Alternative Choices for Context Vectors: The variants are evaluated on ConvAI2 and DSTC7 Track 1, with bold values marking the highest-performing variant for each number of codes.
- Alternative Choices for Context Vectors: Inference time is reported for predicting the next utterance from N possible candidates.
- Alternative Choices for Context Vectors: Figure 2 depicts the Bi-encoder, Cross-encoder, and two Poly-encoder configurations: first m vectors and learnt codes.
- Alternative Choices for Context Vectors: Table 10 reports validation and test performances for Bi-, Poly-, and Cross-encoders on ConvAI2, DSTC7 Track 1, and Ubuntu v2.