# Prototypical Extreme Multi-label Classification with a Dynamic Margin Loss

arXiv: https://arxiv.org/abs/2410.20401 · HTML: https://mlsanity.com/paper/2410.20401/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 56 references that are in the map.
- First posted on arXiv 27 Oct 2024.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [FaceNet: A Unified Embedding for Face Recognition and Clustering](https://mlsanity.com/paper/1503.03832/) (2015)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs](https://mlsanity.com/paper/1603.09320/) (2016)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [Pay Attention to MLPs](https://mlsanity.com/paper/2105.08050/) (2021)
- [Autoregressive Entity Retrieval](https://mlsanity.com/paper/2010.00904/) (2020)
- [GraphFormers: GNN-nested Transformers for Representation Learning on Textual Graph](https://mlsanity.com/paper/2105.02605/) (2021)
- [DiSMEC - Distributed Sparse Machines for Extreme Multi-label Classification](https://mlsanity.com/paper/1609.02521/) (2016)
- [LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification](https://mlsanity.com/paper/2101.03305/) (2021)
- [ECLARE: Extreme Classification with Label Graph Correlations](https://mlsanity.com/paper/2108.00261/) (2021)
- [DeepXML: A Deep Extreme Multi-Label Learning Framework Applied to Short Text Documents](https://mlsanity.com/paper/2111.06685/) (2021)
- [NGAME: Negative Mining-aware Mini-batching for Extreme Classification](https://mlsanity.com/paper/2207.04452/) (2022)
- [Pairwise Supervised Contrastive Learning of Sentence Representations](https://mlsanity.com/paper/2109.05424/) (2021)
- [ProtAugment: Unsupervised diverse short-texts paraphrasing for intent detection meta-learning](https://mlsanity.com/paper/2105.12995/) (2021)
- [More Robust Dense Retrieval with Contrastive Dual Learning](https://mlsanity.com/paper/2107.07773/) (2021)
- [ELIAS: End-to-End Learning to Index and Search in Large Output Spaces](https://mlsanity.com/paper/2210.08410/) (2022)
- [Towards Improved Proxy-based Deep Metric Learning via Data-Augmented Domain Adaptation](https://mlsanity.com/paper/2401.00617/) (2024)
- [PINA: Leveraging Side Information in eXtreme Multi-label Classification via Predicted Instance Neighborhood Aggregation](https://mlsanity.com/paper/2305.12349/) (2023)
- [Contrastive Learning with Prompt-derived Virtual Semantic Prototypes for Unsupervised Sentence Embedding](https://mlsanity.com/paper/2211.03348/) (2022)
- [AdaTriplet: Adaptive Gradient Triplet Loss with Automatic Margin Learning for Forensic Medical Image Matching](https://mlsanity.com/paper/2205.02849/) (2022)
- [A Multi-Granularity-Aware Aspect Learning Model for Multi-Aspect Dense Retrieval](https://mlsanity.com/paper/2312.02538/) (2023)

## Built on by (most cited first, 2 of 2)

- [Fine-grained auxiliary learning for real-world product recommendation](https://mlsanity.com/paper/2510.04551/) (2025)
- [Characterize Then Distill: Mechanistic Reasoning in Large Output Spaces](https://mlsanity.com/paper/2606.06840/) (2026)

## Abstract (arXiv)

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.