Multi-Head Encoding for Extreme Label Classification

TL;DR

This work introduces Multi-Head Encoding (MHE) to address Classifier Computational Overload in eXtreme Label Classification by decomposing extreme labels into $H$ local labels across multiple heads via a Kronecker product, enabling efficient training and inference. It develops three implementations—Multi-Head Product (MHP) for XSLC, Multi-Head Cascade (MHC) for XMLC, and Multi-Head Sampling (MHS) for model pretraining—each tailored to its task while maintaining a unified low-rank interpretation. The authors prove that CE-based low-rank approximations can closely match the vanilla one-hot encoder, reducing reliance on label preprocessing and demonstrating state-of-the-art performance with substantial speedups and memory savings across multiple benchmarks and even in neural machine translation. The findings highlight MHE as a scalable, flexible framework to tackle CCOP in diverse extreme classification problems and beyond.

Abstract

The number of categories of instances in the real world is normally huge, and each instance may contain multiple labels. To distinguish these massive labels utilizing machine learning, eXtreme Label Classification (XLC) has been established. However, as the number of categories increases, the number of parameters and nonlinear operations in the classifier also rises. This results in a Classifier Computational Overload Problem (CCOP). To address this, we propose a Multi-Head Encoding (MHE) mechanism, which replaces the vanilla classifier with a multi-head classifier. During the training process, MHE decomposes extreme labels into the product of multiple short local labels, with each head trained on these local labels. During testing, the predicted labels can be directly calculated from the local predictions of each head. This reduces the computational load geometrically. Then, according to the characteristics of different XLC tasks, e.g., single-label, multi-label, and model pretraining tasks, three MHE-based implementations, i.e., Multi-Head Product, Multi-Head Cascade, and Multi-Head Sampling, are proposed to more effectively cope with CCOP. Moreover, we theoretically demonstrate that MHE can achieve performance approximately equivalent to that of the vanilla classifier by generalizing the low-rank approximation problem from Frobenius-norm to Cross-Entropy. Experimental results show that the proposed methods achieve state-of-the-art performance while significantly streamlining the training and inference processes of XLC tasks. The source code has been made public at https://github.com/Anoise/MHE.

Figures (12)

• Figure 1: The deep neural networks which are composed of three parts: input, backbone and classifier. In multi-head encoding, labels are decomposed onto the outputs of the multi-head classifier during training, and the outputs are combined in testing to obtain the predicted labels.
• Figure 2: The training and testing process of the multi-head classifier using MHE. (a) During training, the global label is decomposed into multiple local labels. (b) During testing, MHE works in the inverse manner of training. (c) The equivalent form of MHE for testing combines the local predictions of the multi-head classifier to obtain the global prediction. Note that the operations marked by the red dotted frame are used for ease of understanding, which are not required.
• Figure 3: Three MHE-based training and testing processes for XLC tasks. The part indicated by the red dotted frame is for ease of understanding, which is not required in practice.
• Figure 4: The low-rank approximation ability of the classifier. (a) $G$ groups of multi-head classifier using MHE (G-MHE). (b) A bottleneck layer is added to the origin classifier to achieve the low-rank property of $\mathcal{W}$.
• Figure 5: Experiments with different loss functions and $R(\tilde{W})$. (a, b) The performance of two-layer linear networks on random samples generated from a Gaussian distribution. (c) The performance of ResNet-18 on the CIFAR-100 dataset.

Theorems & Definitions (10)

• Theorem 1
• Corollary 1
• Theorem 2
• Theorem 3
• Theorem 4
• proof
• proof
• proof
• proof
• proof