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

TL;DR

Extreme Multi-Label Classification with millions of labels is typically hindered by excessive computation when training large one-vs-all classifiers or fully computing losses over the label space. UniDEC presents a loss-independent, end-to-end framework that unifies dual-encoder and classifier training using Pick-Some-Labels (PSL), coupled with dynamic ANNS hard-negative mining, to compute losses on a small, informative subset of labels. The approach achieves state-of-the-art results on six public XMC datasets and a proprietary 450M-label dataset while operating on a single GPU by leveraging a compact label shortlist and in-batch supervision. In online sponsored search, offline gains transfer to real-world metrics, including modest improvements in IY, CY, CTR, and QC, demonstrating practical impact for large-scale retrieval and recommendation tasks. Overall, UniDEC advances computational efficiency without sacrificing accuracy by jointly optimizing DE and classifier heads under a unified, PSL-driven multitask objective.

Abstract

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.

Figures (2)

• Figure 1: The architecture for the UniDEC framework, denoting the the classifiers and DE trained in parallel, along with the loss functions used. The inference pipeline is shown in the rectangular box.
• Figure 2: (a) Visualizing UniDEC's batching strategy. Such a framework naturally leads to higher number of positives per query, enabling us to scale without increasing the batch size significantly. (b) Scatter plot showing the average number of positive labels per query, when we sample $\beta$ positives and $\eta$ hard negatives in the batch. Note that, even with $\beta = 3$ and $\eta = 0$, avg($|P_i^{\mathcal{B}}|) = 13.6$.