针对大规模多标签分类的输出层瓶颈,提出高效稀疏训练方法。
HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces
- 设计分组共享的固定输入稀疏结构,语义相关标签共用特征子集。
- 前向加速达4.4倍,反向传播提速25倍,接近稠密模型性能。
- 无需额外损失函数,适合长尾分布的大规模标签任务。
极端多标签分类(XMC)需在百万级标签空间中建模,导致输出层成为内存与计算瓶颈。现有稀疏方法常因不规则内存访问、硬件利用率低或依赖附加组件,在长尾场景下难以实现预期加速。本文提出分组共享固定输入稀疏结构:语义相关的标签共享稀疏输入模式,同时保持独立权重。该设计引入任务对齐归纳偏置,提升特征复用率,降低索引内存开销,并通过定制CUDA核实现高效GPU执行。作为替代辅助目标的方案,利用XMC的长尾特性,将输出层分解为高频标签的小型稠密头与剩余标签的分组共享稀疏尾,保留有效梯度路径的同时维持稀疏优势。微基准测试显示,该方法将算术减少转化为实际速度提升,前向传播最高提速4.4倍,反向传播最高提速25倍,且性能仅比同FLOPs稠密瓶颈低几个百分点。在多个大规模XMC基准上,精度表现优于或等同于现有稀疏基线,显著缩小与稠密模型的差距。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。