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

AIPR assessment

This is a hard, competitive, well-optimized problem area, not a niche benchmark. The paper's strengths compound: a simple structural idea, measurable runtime gains, and multiple ablations all point in the same direction, which increases confidence that the effect is real. The main weaknesses also compound less severely than in many system papers because the authors do provide code, benchmark across several large datasets, and isolate kernel behavior, but the absence of multi-run statistics and t

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.

Score Breakdown

More from this week

• Towards Interactive Video World Modeling: Frontiers, Challenges, Benchmarks, and Future Trends
• On Thin Perfect Matchings up to Polylogarithmic Factors
• ViBE: Co-Optimizing Workload Skew and Hardware Variability for MoE Serving
• LeAP: Learnable Adaptive Permutation for Feature Selection in Heterogeneous and Sparse Recommender Systems
• PolySpeech-100: A Large-Scale Benchmark for Speech Understanding Across 100+ Languages and Dialects

More in Machine Learning

• STAR-KV: Low-Rank KV Cache Compression via Soft Thresholding for Adaptive Rank Control
• Verifiable and Confidential DNN Inference on Low-End Edge Devices
• LeAP: Learnable Adaptive Permutation for Feature Selection in Heterogeneous and Sparse Recommender Systems
• TriAxialKV: Toward Extreme Low-Precision KV-Cache Quantization for Agentic Inference Tasks
• DynMuon: A Dynamic Spectral Shaping View of Muon