Navigating Extremes: Dynamic Sparsity in Large Output Spaces

TL;DR

This work develops Spartex, a Dynamic Sparse Training (DST) framework tailored for Extreme Multi-label Classification (XMC) with millions of labels. It employs semi-structured, fixed fan-in sparsity for the final classification layer and introduces an auxiliary objective to stabilize gradient flow from a sparse classifier to a trainable encoder, enabling end-to-end training under GPU memory constraints. Empirical results on large-scale datasets show substantial memory reductions (about 3.4x on Amazon-3M) with only modest accuracy loss, and competitive performance against dense baselines and state-of-the-art XMC methods, especially for tail labels. The approach demonstrates practical DST viability for large output spaces on commodity hardware and provides CUDA kernel bindings to facilitate adoption in real-world XMC pipelines.

Abstract

In recent years, Dynamic Sparse Training (DST) has emerged as an alternative to post-training pruning for generating efficient models. In principle, DST allows for a more memory efficient training process, as it maintains sparsity throughout the entire training run. However, current DST implementations fail to capitalize on this in practice. Because sparse matrix multiplication is much less efficient than dense matrix multiplication on GPUs, most implementations simulate sparsity by masking weights. In this paper, we leverage recent advances in semi-structured sparse training to apply DST in the domain of classification with large output spaces, where memory-efficiency is paramount. With a label space of possibly millions of candidates, the classification layer alone will consume several gigabytes of memory. Switching from a dense to a fixed fan-in sparse layer updated with sparse evolutionary training (SET); however, severely hampers training convergence, especially at the largest label spaces. We find that poor gradient flow from the sparse classifier to the dense text encoder make it difficult to learn good input representations. By employing an intermediate layer or adding an auxiliary training objective, we recover most of the generalisation performance of the dense model. Overall, we demonstrate the applicability and practical benefits of DST in a challenging domain -- characterized by a highly skewed label distribution that differs substantially from typical DST benchmark datasets -- which enables end-to-end training with millions of labels on commodity hardware.

Figures (7)

• Figure 1: Model configurations and performance comparisons at various sparsity levels. The left panel illustrates our model configurations: 'S' represents a semi-structured fixed fan-in sparse layer, 'W' denotes an intermediate layer, and 'Aux' refers to an auxiliary head of meta-classifiers. These configurations help maintain performance as the label space size increases from 31K to 670K and beyond. The right panel demonstrates the comparative precision at 1 for our model against other methods across increasing levels of sparsity on the Amazon670K dataset.
• Figure 2: Gradient Flow of the encoder during training with and without Auxiliary Objective.
• Figure 3: left: Comparison of performance declines as the size of the label space increases, given a fixed sparsity. right: Performance of our model at different epochs, across various sparsity ratios.
• Figure 4: Effect of rewiring interval on final performance for Precision@1 (left) and propensity-scored Precision@1 (right) in the LF-AmazonTitles-131K dataset.
• Figure 5: Impact of intermediate layer size on overall and tail label performance. The plots show precision, propensity-scored precision, and macro precision across epochs for different intermediate layer sizes (1024, 2048, 4096, and 8192).