Prototypical Extreme Multi-label Classification with a Dynamic Margin Loss

TL;DR

PRIME introduces a novel prototypical contrastive learning approach for extreme multi-label classification, balancing efficiency and high performance by leveraging label prototypes and a shallow transformer encoder.

Contribution

The paper presents PRIME, a new XMC method that uses a data-to-prototype prediction framework with a shallow encoder and adaptive triplet loss, outperforming existing approaches.

Findings

Achieves state-of-the-art results on multiple benchmarks.

Balances efficiency and performance in large label spaces.

Uses a novel adaptive triplet loss for better training.

Abstract

Extreme Multi-label Classification (XMC) methods predict relevant labels for a given query in an extremely large label space. Recent works in XMC address this problem using deep encoders that project text descriptions to an embedding space suitable for recovering the closest labels. However, learning deep models can be computationally expensive in large output spaces, resulting in a trade-off between high performing brute-force approaches and efficient solutions. In this paper, we propose PRIME, a XMC method that employs a novel prototypical contrastive learning technique to reconcile efficiency and performance surpassing brute-force approaches. We frame XMC as a data-to-prototype prediction task where label prototypes aggregate information from related queries. More precisely, we use a shallow transformer encoder that we coin as Label Prototype Network, which enriches label representations by aggregating text-based embeddings, label centroids and learnable free vectors. We jointly train a deep encoder and the Label Prototype Network using an adaptive triplet loss objective that better adapts to the high granularity and ambiguity of extreme label spaces. PRIME achieves state-of-the-art results in several public benchmarks of different sizes and domains, while keeping the model efficient.