Leveraging Passage Embeddings for Efficient Listwise Reranking with Large Language Models

2024-06-21

arXiv.org

AI Summary

To address the high prefill computational overhead, elevated decoding latency, and low inference efficiency of large language models (LLMs) in listwise document re-ranking—caused by long-context inputs—this paper proposes an embedding-driven dynamic constrained decoding framework. Methodologically, it explicitly models contextual semantics by treating paragraph embeddings as learnable “special tokens” and dynamically restricts the decoding vocabulary to this token set during generation, thereby eliminating redundant text production. The approach integrates paragraph embedding compression, special-token representation learning, and listwise learning-to-rank loss optimization. Evaluated on multiple standard re-ranking benchmarks, the method substantially reduces prefill computation and decoding steps, achieving 2.1–3.8× faster inference while maintaining state-of-the-art ranking performance.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Learning Preferences or Rankings
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Recent studies have demonstrated the effectiveness of using large language language models (LLMs) in passage ranking. The listwise approaches, such as RankGPT, have become new state-of-the-art in this task. However, the efficiency of RankGPT models is limited by the maximum context length and relatively high latency of LLM inference. To address these issues, in this paper, we propose PE-Rank, leveraging the single passage embedding as a good context compression for efficient listwise passage reranking. By treating each passage as a special token, we can directly input passage embeddings into LLMs, thereby reducing input length. Additionally, we introduce an inference method that dynamically constrains the decoding space to these special tokens, accelerating the decoding process. For adapting the model to reranking, we employ listwise learning to rank loss for training. Evaluation results on multiple benchmarks demonstrate that PE-Rank significantly improves efficiency in both prefilling and decoding, while maintaining competitive ranking effectiveness. The Code is available at https://github.com/liuqi6777/pe_rank.

Problem

• Large Language Models
• Long Passage Processing
• Computational Speed

Innovation

• PE-Rank
• Paragraph Compression
• Speed Optimization