A Study on Token Pruning for ColBERT

TL;DR

This paper investigates token pruning techniques for the ColBERT model to reduce index size, demonstrating up to 30% pruning with minimal performance loss on MS MARCO passage data.

Contribution

It introduces token pruning methods for ColBERT, comparing heuristics and attention-based selection, to effectively reduce index size while maintaining accuracy.

Findings

Up to 30% index pruning without performance loss on MS MARCO passage data.

Token pruning faces challenges when applied to MS MARCO document collection.

Attention-based token selection can be effective for index compression.

Abstract

The ColBERT model has recently been proposed as an effective BERT based ranker. By adopting a late interaction mechanism, a major advantage of ColBERT is that document representations can be precomputed in advance. However, the big downside of the model is the index size, which scales linearly with the number of tokens in the collection. In this paper, we study various designs for ColBERT models in order to attack this problem. While compression techniques have been explored to reduce the index size, in this paper we study token pruning techniques for ColBERT. We compare simple heuristics, as well as a single layer of attention mechanism to select the tokens to keep at indexing time. Our experiments show that ColBERT indexes can be pruned up to 30\% on the MS MARCO passage collection without a significant drop in performance. Finally, we experiment on MS MARCO documents, which reveal several challenges for such mechanism.