HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces

Extreme multi-label classification (XMC) involves learning models over large output spaces with millions of labels, making the output layer a memory-compute bottleneck. While sparsity-based methods reduce arithmetic complexity, they often fail to yield proportional speedups due to irregular memory access, poor hardware utilization, or reliance on auxiliary architectural components in long-tailed regimes. We introduce group-shared fixed fan-in sparsity, a semi-structured output-layer design in which semantically related labels share a sparse input pattern while retaining independent weights. This grouping introduces a task-aligned inductive bias---encouraging related labels to share feature subsets---while reducing index memory overhead, increasing feature reuse across labels, and enabling efficient GPU execution via custom CUDA kernels that leverage modern accelerator primitives. As an alternative to auxiliary objectives, we exploit the long-tailed structure of XMC by decomposing the output layer into a small dense head over frequent labels and a group-shared sparse tail over the remainder, providing an informative gradient pathway while preserving the memory benefits of sparsity. Through kernel-level microbenchmarking, we show that group-shared fixed fan-in translates arithmetic reductions into practical wall-clock gains, achieving up to $4.4\times$ speedup in the forward pass and up to $25\times$ speedup in backward passes over standard fixed fan-in sparsity, while operating within a few percent of a FLOPs-matched dense bottleneck. Across large-scale XMC benchmarks, our approach matches or improves precision@k over prior sparse baselines, while narrowing the performance gap to dense.

Many AI systems need to choose from a very large set of possible labels: for example, assigning product tags, recommending related items, or matching search queries to advertisements. When there are millions of labels, the final prediction layer of the model can become too large and slow to train, even if the rest of the model is manageable. Our work introduces HASTE, a training method that makes this large prediction layer sparse in a hardware-friendly way. The key idea is to group related labels, such as similar products or topics, and let labels in the same group reuse the same small subset of model features. This reduces memory use and makes the computation better aligned with modern GPUs, instead of relying on random sparse access patterns that are often inefficient. We also handle frequent and rare labels differently to make training more stable. Overall, HASTE helps train large-label models faster and with less memory while keeping prediction quality close to dense models.