[Paper Review] Poly-encoders: Transformer Architectures and Pre-training Strategies for Fast and Accurate Multi-sentence Scoring

Samuel Humeau, Kurt Shuster+2 more|arXiv (Cornell University)|Apr 22, 2019

Topic Modeling34 references154 citations

TL;DR

Introduces Poly-encoder, a transformer architecture that combines caching benefits of Bi-encoders with the rich interactions of Cross-encoders, enabling faster inference and higher accuracy for multi-sentence scoring tasks; shows pre-training on task-like data (notably Reddit) yields state-of-the-art results across dialogue and IR datasets.

ABSTRACT

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.

Motivation & Objective

• Advance multi-sentence scoring by balancing accuracy and speed in candidate retrieval tasks.
• Compare Bi-encoder, Cross-encoder, and the proposed Poly-encoder architectures under various pre-training strategies.
• Demonstrate the impact of task-relevant pre-training data on downstream performance.
• Show scalability and practical inference speed trade-offs for real-time systems.

Proposed method

• Propose three architectures: Bi-encoder (separate context and label encodings), Cross-encoder (joint encoding of context and label with full self-attention), and Poly-encoder (context attends to a small set of learned global context codes while label is encoded separately).
• Develop pre-training regimes: BERT-like pre-training on Wikipedia/Toronto Books; from-scratch pre-training on Reddit (for dialogue relevance) with MLM and next-sentence/utterance prediction tasks.
• Fine-tune each architecture on four tasks (ConvAI2, DSTC7, Ubuntu V2, Wikipedia Article Search) with task-appropriate negatives and evaluation metrics.
• Investigate hyperparameters, including batch size, negatives per batch, and which transformer layers to fine-tune, to maximize performance.
• Assess inference speed across CPU/GPU and compare training efficiency among Bi-, Poly-, and Cross-encoders.
• Evaluate the effect of different numbers of context codes (m) in Poly-encoders on performance and speed.

Experimental results

Research questions

• RQ1 Can Poly-encoders achieve a favorable balance of accuracy and speed compared to Bi-encoders and Cross-encoders for multi-sentence scoring tasks?
• RQ2 How do different pre-training data (Reddit vs. Wikipedia/Toronto Books) influence downstream performance across dialogue and IR datasets?
• RQ3 What is the impact of the number of context codes (m) in Poly-encoders on retrieval accuracy and inference time?
• RQ4 What are the practical training and inference time trade-offs for deploying these architectures in real-time systems?

Key findings

• Poly-encoders consistently outperform Bi-encoders in accuracy across evaluated tasks.
• Poly-encoders achieve substantially faster inference than Cross-encoders, making them suitable for real-time retrieval and dialogue systems.
• Pre-training on Reddit (dialogue-like data) yields state-of-the-art results across all three dialogue tasks and the IR task when combined with Poly-encoders or Bi-encoders.
• Using larger numbers of context codes (higher m) generally improves Poly-encoder performance, with diminishing returns weighed against speed as m grows.
• Inference times show Cross-encoders are orders of magnitude slower, while Poly-encoders offer a favorable accuracy-speed trade-off suitable for production deployments.
• Pre-training on Reddit leads to improvements over BERT-based initializations across all architectures and tasks tested.