Saved in:
Bibliographic Details
Main Authors: Ullah, Nasib, Zhang, Jinbin, Randrianantenaina, Jean Lucien, Schultheis, Erik, Babbar, Rohit
Format: Preprint
Published: 2026
Subjects:
Online Access:https://arxiv.org/abs/2606.01117
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866914620797616128
author Ullah, Nasib
Zhang, Jinbin
Randrianantenaina, Jean Lucien
Schultheis, Erik
Babbar, Rohit
author_facet Ullah, Nasib
Zhang, Jinbin
Randrianantenaina, Jean Lucien
Schultheis, Erik
Babbar, Rohit
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2606_01117
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces
Ullah, Nasib
Zhang, Jinbin
Randrianantenaina, Jean Lucien
Schultheis, Erik
Babbar, Rohit
Machine Learning
Artificial Intelligence
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.
title HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2606.01117