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

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

Problem Statement
Extreme Multi-Label Classification (XMC) involves training models over output spaces containing millions of labels. While sparsity-based methods have been proposed to reduce the arithmetic complexity of the output layer, they often fail to deliver proportional speedups in practice. The primary bottlenecks are irregular memory access patterns, poor hardware utilization on modern GPUs (which are optimized for dense, contiguous blocks), and the "memory wall" caused by uncoalesced reads. Furthermore, training in long-tailed regimes with sparse connectivity often leads to instability due to weak gradient signals, frequently necessitating auxiliary objectives that introduce their own computational overhead and potential gradient conflicts.

Methodology
The authors propose HASTE, a framework that combines a novel semi-structured sparsity pattern with a data-driven optimization strategy.

1. Group-Shared Fixed Fan-In Sparsity:
   Unlike standard fixed fan-in sparsity (where each label has an independent set of input features) or rigid block sparsity, HASTE introduces a "group-shared" approach.

   • Structure: Labels are partitioned into groups based on semantic similarity (using label embeddings derived from training instances). All labels within a group share the same sparse input feature indices (fan-in pattern), though they retain independent weight vectors.
   • Hardware Alignment: This design enables "gather-once and dense MMA" execution patterns. By grouping labels, the input features for a group can be fetched once into shared memory and reused by all labels in that group. This facilitates the use of Tensor Core primitives via custom CUDA kernels, converting arithmetic reductions into actual wall-clock speedups.
   • Memory Efficiency: This approach drastically reduces index memory overhead. Instead of storing indices for every label, indices are stored once per group, reducing overhead by a factor proportional to the group size.

2. Head–Tail Decomposition (Auxiliary-Free Optimization):
   To address training instability in long-tailed distributions without relying on auxiliary objectives, HASTE decomposes the label space:

   • Head: A small subset of frequent labels is modeled using a standard dense classifier.
   • Tail: The remaining, vast majority of infrequent labels are handled by the group-shared sparse layer.
   • Mechanism: Both branches share the same encoder but utilize separate lightweight feature projections. The dense head provides a consistent, strong gradient signal to the encoder during training, stabilizing the learning process for the sparse tail without the need for auxiliary supervision branches.

Key Contributions

• Novel Sparsity Pattern: The introduction of group-shared fixed fan-in sparsity, which balances the expressiveness of independent feature selection with the hardware efficiency of structured, shared access.
• Hardware-Aware Kernels: The development of custom CUDA kernels that leverage the group structure to perform dense matrix multiplications (GEMM) on gathered feature tiles, enabling efficient use of Tensor Cores.
• Auxiliary-Free Stability: A decomposition strategy that stabilizes training in long-tailed regimes by leveraging the inherent frequency distribution of labels, eliminating the need for auxiliary objectives and their associated overhead.
• Empirical Validation: Comprehensive evaluation across large-scale benchmarks (up to 8.6 million labels) demonstrating that HASTE narrows the performance gap with dense models while significantly outperforming prior sparse baselines in speed and memory efficiency.

Experimental Results

• Speedup: Microbenchmarks on NVIDIA A100 GPUs show that HASTE achieves up to 4.4× speedup in the forward pass and up to 25× speedup in backward passes compared to standard fixed fan-in sparsity (SPARTEX). It operates within a few percent of a FLOPs-matched dense bottleneck.
• Accuracy: On benchmarks including Amazon-670K, Amazon-3M, and LF-Paper2Keywords-8.6M, HASTE matches or improves upon Precision@k (P@k) and Propensity-Scored Precision@k (PSP@k) compared to state-of-the-art sparse baselines (e.g., SPARTEX, BLOCK SPARSE).
• Memory Efficiency: The method significantly reduces peak training memory (e.g., 3.2 GiB vs. 5.0 GiB for SPARTEX on AmazonTitles-670K) and reduces epoch time.
• Group Size Trade-off: Experiments indicate a trade-off where smaller group sizes (e.g., 16) yield slightly higher accuracy due to better representational flexibility, while larger groups (e.g., 64) improve kernel efficiency by increasing feature reuse.

Significance and Claims
The paper claims that HASTE successfully bridges the gap between the theoretical arithmetic benefits of sparsity and the practical realities of modern GPU hardware. By introducing structure into the sparsity pattern (group-sharing) and aligning it with the long-tailed nature of XMC data (head-tail split), the approach converts FLOPs reductions into tangible wall-clock gains. The authors assert that their method makes training models with millions of labels more accessible under fixed compute budgets by reducing both training time and memory footprint, while maintaining competitive accuracy. They note that while efficiency is improved, it does not inherently solve issues related to data bias or imperfect labeling in long-tailed datasets, and users should remain vigilant regarding tail-label performance.