Retrieval-augmented Multi-label Text Classification

TL;DR

This paper introduces a retrieval-augmented approach for multi-label text classification that enhances performance on infrequent labels by incorporating similar document retrieval into a Transformer-based model, especially benefiting low-resource and long-document scenarios.

Contribution

It proposes a novel retrieval-augmented architecture for multi-label classification that improves sample efficiency and tail label performance in skewed label distributions.

Findings

Retrieval augmentation significantly improves performance on infrequent labels.

The approach is especially effective in low-resource and long-document scenarios.

Experiments on legal and biomedical datasets validate the method's effectiveness.

Abstract

Multi-label text classification (MLC) is a challenging task in settings of large label sets, where label support follows a Zipfian distribution. In this paper, we address this problem through retrieval augmentation, aiming to improve the sample efficiency of classification models. Our approach closely follows the standard MLC architecture of a Transformer-based encoder paired with a set of classification heads. In our case, however, the input document representation is augmented through cross-attention to similar documents retrieved from the training set and represented in a task-specific manner. We evaluate this approach on four datasets from the legal and biomedical domains, all of which feature highly skewed label distributions. Our experiments show that retrieval augmentation substantially improves model performance on the long tail of infrequent labels especially so for lower-resource training scenarios and more challenging long-document data scenarios.