通过中间层分层压缩token,显著提升文档重排序效率。
Layer-wise Token Compression for Efficient Document Reranking

- 在Transformer中间层动态聚合token,实现自适应压缩
- 文档重排序推理速度提升116%,精度几乎不变
- 适合需要高效长文本重排序的场景,如大模型列表重排
基于Transformer的文档交叉编码器重排序器是现代信息检索系统的核心组件。尽管效果优异,但其在推理时处理长序列导致计算开销高。已有方法如令牌压缩(通过初始嵌入层聚合令牌组)可减少有效令牌数以加速计算,但在交叉编码器中效果有限。本文提出分层令牌压缩(LTC),在中间Transformer层进行自适应令牌池化。在MS MARCO段落与文档重排序任务上的大量消融实验表明,中层压缩在保持排名质量的同时,使段落重排序的每秒查询率(QPS)最高提升25%,文档重排序提升116%。我们还将LTC扩展至列表式大模型重排序,同样实现高效长上下文重排序,且提速更显著。更意外的是,用短段落训练的压缩模型在长文档任务上表现优于未压缩模型,说明压缩可能作为有益正则化,促进长度无关表示学习。
原文摘要 · Abstract (English)
Transformer-based document cross-encoder rerankers are a central component of modern information retrieval systems. Despite their success, these models suffer from high computational costs due to processing long query-document sequences at inference time. A known approach to improve efficiency is token compression, which consists of aggregating groups of tokens together in the initial embedding layer, reducing the effective number of tokens, and making the computation faster. While token compression has proven to be successful for bi-encoder retrievers, we empirically observed that this approach may be ineffective for cross-encoder rerankers. In this paper, we propose Layer-wise Token Compression (LTC), which applies adaptive token pooling at intermediate transformer layers. Through extensive ablation studies on MS MARCO passage and document ranking tasks, we demonstrate that compression at middle layers preserves ranking quality while increasing inference QPS by up to 25% for passage ranking and up to 116% for document ranking. We also extend LTC to listwise LLM rerankers and show that the same approach can be easily applied to long-context listwise reranking, where the QPS improvements are even greater. More surprisingly, when applying rerankers trained on short passages to long-document ranking tasks, models trained with compression outperform their uncompressed counterparts, suggesting that compression may act as a beneficial regularizer that encourages length-invariant representations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。