[论文解读] Top-KAST: Top-K Always Sparse Training

Siddhant M. Jayakumar, Razvan Pascanu+3 more|arXiv (Cornell University)|Jun 7, 2021

Advanced Neural Network Applications被引用 21

一句话总结

Top-KAST 是一种完全稀疏的训练方法，通过在前向传播中选择激活的前K个权重大小，在反向传播中使用更大的参数集进行梯度更新，同时引入探索损失以防止收敛到次优掩码。该方法在 ImageNet 和 enwik8 上实现了最先进性能，在高达 80% 稀疏度下与稠密模型表现相当，同时实现了高效、可扩展的训练，无需计算稠密参数或梯度。

ABSTRACT

Sparse neural networks are becoming increasingly important as the field seeks to improve the performance of existing models by scaling them up, while simultaneously trying to reduce power consumption and computational footprint. Unfortunately, most existing methods for inducing performant sparse models still entail the instantiation of dense parameters, or dense gradients in the backward-pass, during training. For very large models this requirement can be prohibitive. In this work we propose Top-KAST, a method that preserves constant sparsity throughout training (in both the forward and backward-passes). We demonstrate the efficacy of our approach by showing that it performs comparably to or better than previous works when training models on the established ImageNet benchmark, whilst fully maintaining sparsity. In addition to our ImageNet results, we also demonstrate our approach in the domain of language modeling where the current best performing architectures tend to have tens of billions of parameters and scaling up does not yet seem to have saturated performance. Sparse versions of these architectures can be run with significantly fewer resources, making them more widely accessible and applicable. Furthermore, in addition to being effective, our approach is straightforward and can easily be implemented in a wide range of existing machine learning frameworks with only a few additional lines of code. We therefore hope that our contribution will help enable the broader community to explore the potential held by massive models, without incurring massive computational cost.

研究动机与目标

• 开发一种完全稀疏的训练方法，避免在训练过程中生成稠密参数或梯度计算，从而实现可扩展且高效的模型训练。
• 通过在整个训练过程中保持稀疏性，降低大规模稠密模型训练的计算和环境成本。
• 通过保持恒定稀疏性，降低硬件和能耗需求，从而实现大规模模型的普惠化。
• 在低资源环境下，有效实现如 Transformer 等大型架构的稀疏化，特别是在自然语言处理领域。
• 提供一种简单、框架兼容的方法，可轻松集成到现有深度学习库中。

提出的方法

• Top-KAST 在前向传播中根据参数大小选择前K个参数，以在推理过程中保持稀疏性。
• 在反向传播中，梯度作用于更大的参数集合 B 上，其中 B 包含前K个集合 A，以实现更有效的权重更新。
• 引入辅助探索损失，以鼓励稀疏掩码的动态适应，防止固定在次优权重子集上。
• 该方法完全避免了稠密参数或梯度计算，确保前向和反向传播中均保持恒定稀疏性。
• 稀疏度水平由用户定义的超参数控制，包括前K个参数和探索集合的大小，可根据资源约束灵活调整。
• 该方法设计用于轻松集成到现有深度学习框架中，仅需极少代码修改。

实验结果

研究问题

• RQ1 是否能够实现一种完全稀疏的训练方法，在从不生成稠密参数或梯度的情况下，保持与稠密模型相当的性能？
• RQ2 在 ImageNet 和 enwik8 等标准基准上，Top-KAST 相较于现有的稀疏到稀疏和稠密到稀疏训练方法表现如何？
• RQ3 Top-KAST 是否能有效对大型 Transformer 架构进行稀疏化，同时在语言建模任务中实现最小的性能损失？
• RQ4 探索损失的引入是否能提升泛化能力，并防止收敛到性能较差的稀疏子网络？
• RQ5 在实际应用中，Top-KAST 在多大程度上能够扩展模型规模，同时降低计算和内存开销？

主要发现

• Top-KAST 在 ImageNet 基准上实现了多种 FLOP 预算下的最先进性能，优于现有稀疏到稀疏方法，并与 RigL 表现相当。
• 在 enwik8 语言建模任务中，Top-KAST 训练的 Transformer-XL 在 5500 万参数下达到 1.00 比特每字符（bpc），优于参数更少的稠密模型（2.77 亿参数，0.99 bpc）。
• 在 ImageNet 上，Top-KAST 在高达 80% 稀疏度下保持高性能，性能下降可忽略不计，仅在超过 90% 稀疏度后开始出现漂移。
• Top-KAST 能够训练出在高达 80% 稀疏度下与稠密模型性能相当的稀疏模型，展现出强大的效率和可扩展性。
• 该方法在视觉和语言建模任务中均有效，显示出在不同模型类型和领域中的广泛适用性。
• 该方法可在现有框架中轻松实现，仅需增加少量额外代码行，具备广泛采用潜力。