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

2026-05-31

AI Summary

This work addresses the significant memory and computational bottlenecks in extreme multi-label classification (XMC) caused by massive label sets, where existing sparse methods struggle to achieve efficient hardware acceleration. The authors propose a group-shared fixed fan-in sparse architecture that clusters semantically related labels into groups, sharing a common sparse input pattern while maintaining independent weights. To handle the long-tailed label distribution, the model combines a dense head for frequent labels with a sparse tail for rare ones. By incorporating task-aligned inductive bias, the approach reduces indexing overhead and enhances feature reuse, and it is further accelerated via custom CUDA kernels for hardware-aware training. Experiments demonstrate up to 4.4× speedup in forward computation and 25× speedup in backpropagation, achieving accuracy on large-scale XMC benchmarks that matches or exceeds prior sparse methods while substantially narrowing the performance gap with dense models.

Abstract

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.

Problem

• Extreme Multi-label Classification
• Sparse Training
• Hardware Awareness
• Output Layer Bottleneck
• Long-tailed Distribution

Innovation

• group-shared sparsity
• hardware-aware sparse training
• extreme multi-label classification
• semi-structured sparsity
• CUDA kernel optimization