让大输出空间分类模型实现高效训练,可在普通硬件上跑百万标签任务。
Navigating Extremes: Dynamic Sparsity in Large Output Spaces
- 用稀疏进化训练动态保持权重稀疏,避免传统剪枝的内存浪费。
- 通过中间层或辅助目标恢复训练收敛,使稀疏分类器性能接近密集模型。
- 适用于标签数百万、分布极不均衡的大规模分类场景,适合资源受限者。
近年来,动态稀疏训练(DST)作为训练后剪枝的替代方案,旨在实现更高效的模型训练。理论上,DST在整个训练过程中保持稀疏性,可显著降低内存占用。然而,当前DST实践因稀疏矩阵乘法在GPU上效率远低于密集乘法而受限,多数实现通过掩码模拟稀疏。本文利用半结构化稀疏训练的新进展,将DST应用于大规模输出空间的分类任务,该场景中内存效率至关重要。当标签空间达数百万时,仅分类层就需数GB内存。将密集层替换为固定扇入稀疏层并使用稀疏进化训练(SET),虽能节省内存,却严重损害训练收敛,尤其在最大标签空间下。我们发现,稀疏分类器向密集文本编码器传递梯度能力差,导致输入表示学习困难。通过引入中间层或添加辅助训练目标,基本恢复了密集模型的泛化性能。结果表明,DST在高度偏斜标签分布的挑战性场景中具有实际应用价值,实现了在消费级硬件上端到端训练百万级标签的可行性。
原文摘要 · Abstract (English)
In recent years, Dynamic Sparse Training (DST) has emerged as an alternative to post-training pruning for generating efficient models. In principle, DST allows for a more memory efficient training process, as it maintains sparsity throughout the entire training run. However, current DST implementations fail to capitalize on this in practice. Because sparse matrix multiplication is much less efficient than dense matrix multiplication on GPUs, most implementations simulate sparsity by masking weights. In this paper, we leverage recent advances in semi-structured sparse training to apply DST in the domain of classification with large output spaces, where memory-efficiency is paramount. With a label space of possibly millions of candidates, the classification layer alone will consume several gigabytes of memory. Switching from a dense to a fixed fan-in sparse layer updated with sparse evolutionary training (SET); however, severely hampers training convergence, especially at the largest label spaces. We find that poor gradient flow from the sparse classifier to the dense text encoder make it difficult to learn good input representations. By employing an intermediate layer or adding an auxiliary training objective, we recover most of the generalisation performance of the dense model. Overall, we demonstrate the applicability and practical benefits of DST in a challenging domain -- characterized by a highly skewed label distribution that differs substantially from typical DST benchmark datasets -- which enables end-to-end training with millions of labels on commodity hardware.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。