Prototypical Extreme Multi-label Classification with a Dynamic Margin Loss (2025)

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.

## introduction extreme multi-label classification (xmc) is the task of predicting the most relevant subset of labels in an extremely large label space (potentially millions of labels) for a given query data point. xmc methods find applications in various realworld problems, including product recommendation in e-commerce, document tagging, and sponsored search. the label is typically endowed with a short description in such applications. for instance, in a product label space, "kerplunk!: stories" and "the good samaritan strikes again" are relevant labels for the query "the grasshopper trap". in this example, relevancy is defined by products that were seen or bought together, while other applications might define different relevancy relations. the complete example is presented in xmc algorithms typically learn an encoder and extreme classifiers either jointly @xcite or separately by following a modular approach @xcite @xcite . this modular approach allows for effortless utilization of label metadata in the form of text @xcite , images @xcite , or graphs @xcite , thereby enhancing encoder's accuracy. therefore, learning a robust encoder has become a cornerstone for xmc methods. ## related work extreme multi-label classification (xmc). deep xmc methods @xcite @xcite @xcite have demonstrated that learning task-specific embeddings can yield significantly more accurate results as compared to traditional methods (see appendix g for details). it is well established that jointly learning the encoder and the extreme classifiers can be expensive with a large number of labels (dahiya et al., 2021b). in particular, the brute-force approach employed by renée (jain et al., 2023) can lead to accurate results, however, it is expensive to scale beyond a million labels even with multiple gpus. negative sampling approaches, i.e., ance (xiong et al., 2021) and ngame (dahiya et al., 2023a) narrow down the training complexity by selecting a small subset of hard negative labels. moreover, several methods use a modular strategy (dahiya et al., 2021b; zhang et al., 2021b) that decouples the training of encoders and classifiers to improve scalability. eclare (mittal et al., 2021), pina (chien et al., 2023), and xr-transformers (zhang et al., 2021a) train the encoder by learning meta-classifiers, i.e., classifiers are learned for groups of labels. conversely, ngame (dahiya et al., 2023b) proposed to learn a shared deep encoder in a siamese fashion, where labels are represented via label embeddings. building on ngame, dexa (dahiya et al., 2023b) introduces auxiliary parameters to bridge the semantic gap, i.e., the lack of information to accurately represent labels when their descriptions are short. on the other hand, dexml (gupta et al., 2024) demonstrates that siamese deep encoders alone can theoretically yield state-of-the-art results using a brute-force approach with large batch sizes and complete label sets as negatives. deep metric learning (dml). contrastive learning using the vanilla contrastive loss @xcite , triplet loss @xcite or infonce losses @xcite is the core of dml across nlp and computer vision. in nlp, they are widely used to learn robust sentence representations @xcite or, more generally, multi-purpose embeddings with encoders @xcite or decoderonly architectures @xcite . as for xmc, dml has successfully been applied using deep encoders @xcite @xcite . however, this approaches are based on data-to-data comparisons, which in xmc becomes specially complex, as demonstrated by the great benefits of increasing batch-wise comparisons @xcite . however, the dml community has proposed alternatives to narrow down such complexity exploiting data-to-prototype @xcite or data-to-proxy (kim et al., 2020) relations. prototypes and proxies are both representative embeddings for groups of semantically similar instances @xcite . the difference is that proxies are usually learned as network parameters @xcite , while prototypes are computed averaging embeddings of related instances @xcite @xcite . therefore, a few set of proxies or prototypes can capture the global structure of an embedding space and replace data-to-data comparisons for reduced training complexity @xcite . adaptive triplet loss. the triplet loss @xcite is a widely used alternative for metric learning across many fields @xcite @xcite . this loss function ensures that anchor-positive label relations have higher similarity than anchor-negative label ones, while including a margin to force a sufficiently high similarity gap. despite methods in xmc using a fixed margin that lacks flexibility to address ambiguous situations @xcite , some works in the wider metric learning community have proposed dynamic margins to tackle this limitation. for example, @xcite propose a margin function for measuring more detailed similarities between taxonomy paths within the context of selfsupervised taxonomy expansion. in forensic medical image matching, the authors in @xcite propose an auto-margin strategy based on exploiting similarity statistics over time. similarly, @xcite work improves the triplet loss by adapting the margin based on the interpolation strength of mixed samples. decoder models. decoder-only models provide significant benefits over encoder-only ones in multiple natural language problems including entity retrieval @xcite , natural question answering @xcite or the mteb benchmark @xcite . @xcite and quest @xcite have applied decoder-only models in extreme classification. however, their perforour method exploits queryto-prototype (h q to z l ), query-to-label (h q to h l ) and label-to-query (h l to h q ) embedding relations for robust encoder training. the label prototype network g ϕ computes enhanced label embeddings z l (label prototypes) by aggregating label-text embeddings h l , label centroids c l and learnable free vectors v l . as a result, the query-to-prototype similarity score b ql (z) accurately predicts the relevant labels for a given query. ## method we propose prime, a novel method for xmc based on label prototypes rather than the usual text-based label embeddings from related works @xcite @xcite . prime computes data-toprototype relations and uses a dynamic margin in its triplet loss objectives, achieving state-of-the-art results on multiple benchmarks. the overview of prime's architecture is presented in figure 2 . we pose xmc as a maximum inner product search task between query and label embeddings to accurately predict the positive labels for every query q i . to accomplish this, we aim to learn a function f θ : (q, y) → r d , where θ denotes the parameters of a neural network model that encodes the given textual representation into a ddimensional sentence embedding. this function is used to encode every query q i and label l r into h i q and h r l . note that h l can be defined as h p and h n to distinguish between positive and negative text-based label embeddings, respectively. existing siamese training for xmc @xcite directly imposes a data-to-data triplet loss objective as follows: ## conclusion this paper proposes prime, a novel prototypical extreme multi-label classification method demonstrating that encoder-based models do not need to sacrifice efficiency to deliver state-of-the-art performance. prime's key contributions revolve around two key concepts, leveraging dynamic marginbased contrastive learning and label prototypes. the proposed label prototype network learns to efficiently aggregate information from text-based em