# UniDEC : Unified Dual Encoder and Classifier Training for Extreme Multi-Label Classification

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

## Key facts

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

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2010.08191/) (2020)
- [Hard Negative Mixing for Contrastive Learning](https://mlsanity.com/paper/2010.01028/) (2020)
- [Supervised Contrastive Learning for Pre-trained Language Model Fine-tuning](https://mlsanity.com/paper/2011.01403/) (2020)
- [Unified Contrastive Learning in Image-Text-Label Space](https://mlsanity.com/paper/2204.03610/) (2022)
- [AttentionXML: Label Tree-based Attention-Aware Deep Model for High-Performance Extreme Multi-Label Text Classification](https://mlsanity.com/paper/1811.01727/) (2018)
- [DiSMEC - Distributed Sparse Machines for Extreme Multi-label Classification](https://mlsanity.com/paper/1609.02521/) (2016)
- [Fast Multi-Resolution Transformer Fine-tuning for Extreme Multi-label Text Classification](https://mlsanity.com/paper/2110.00685/) (2021)
- [Bonsai -- Diverse and Shallow Trees for Extreme Multi-label Classification](https://mlsanity.com/paper/1904.08249/) (2019)
- [LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification](https://mlsanity.com/paper/2101.03305/) (2021)
- [A no-regret generalization of hierarchical softmax to extreme multi-label classification](https://mlsanity.com/paper/1810.11671/) (2018)
- [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)
- [DECAF: Deep Extreme Classification with Label Features](https://mlsanity.com/paper/2108.00368/) (2021)
- [NGAME: Negative Mining-aware Mini-batching for Extreme Classification](https://mlsanity.com/paper/2207.04452/) (2022)
- [Stochastic Negative Mining for Learning with Large Output Spaces](https://mlsanity.com/paper/1810.07076/) (2018)
- [On Missing Labels, Long-tails and Propensities in Extreme Multi-label Classification](https://mlsanity.com/paper/2207.13186/) (2022)
- [CascadeXML: Rethinking Transformers for End-to-end Multi-resolution Training in Extreme Multi-label Classification](https://mlsanity.com/paper/2211.00640/) (2022)
- [InceptionXML: A Lightweight Framework with Synchronized Negative Sampling for Short Text Extreme Classification](https://mlsanity.com/paper/2109.07319/) (2021)
- [Propensity-scored Probabilistic Label Trees](https://mlsanity.com/paper/2110.10803/) (2021)
- [Dual-Encoders for Extreme Multi-Label Classification](https://mlsanity.com/paper/2310.10636/) (2023)
- [Unbiased Loss Functions for Multilabel Classification with Missing Labels](https://mlsanity.com/paper/2109.11282/) (2021)
- [Speeding-up One-vs-All Training for Extreme Classification via Smart Initialization](https://mlsanity.com/paper/2109.13122/) (2021)

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

- [BroadGen: A Framework for Generating Effective and Efficient Advertiser Broad Match Keyphrase Recommendations](https://mlsanity.com/paper/2505.19164/) (2025)

## Abstract (arXiv)

Extreme Multi-label Classification (XMC) involves predicting a subset of relevant labels from an extremely large label space, given an input query and labels with textual features. Models developed for this problem have conventionally made use of dual encoder (DE) to embed the queries and label texts and one-vs-all (OvA) classifiers to rerank the shortlisted labels by the DE. While such methods have shown empirical success, a major drawback is their computational cost, often requiring upto 16 GPUs to train on the largest public dataset. Such a high cost is a consequence of calculating the loss over the entire label space. While shortlisting strategies have been proposed for classifiers, we aim to study such methods for the DE framework. In this work, we develop UniDEC, a loss-independent, end-to-end trainable framework which trains the DE and classifier together in a unified manner with a multi-class loss, while reducing the computational cost by 4-16x. This is done via the proposed pick-some-label (PSL) reduction, which aims to compute the loss on only a subset of positive and negative labels. These labels are carefully chosen in-batch so as to maximise their supervisory signals. Not only does the proposed framework achieve state-of-the-art results on datasets with labels in the order of millions, it is also computationally and resource efficient in achieving this performance on a single GPU. Code is made available at https://github.com/the-catalyst/UniDEC.