[论文解读] Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models

Tri Dao, Beidi Chen+5 more|arXiv (Cornell University)|Nov 30, 2021

Advanced Neural Network Applications参考文献 90被引用 7

一句话总结

Pixelated Butterfly 提出了一种简单且硬件高效的稀疏训练方法，通过结合平坦块状蝴蝶矩阵与低秩组件，对神经网络层进行稀疏化。该方法在 ImageNet 和 WikiText-103 上实现比 Vision Transformer、GPT-2 和 MLP-Mixer 等稠密模型快达 2.5 倍的训练速度，且准确率无下降，通过实现块对齐的并行计算，在保持与稠密模型相近逼近能力的同时显著提升效率。

ABSTRACT

Overparameterized neural networks generalize well but are expensive to train. Ideally, one would like to reduce their computational cost while retaining their generalization benefits. Sparse model training is a simple and promising approach to achieve this, but there remain challenges as existing methods struggle with accuracy loss, slow training runtime, or difficulty in sparsifying all model components. The core problem is that searching for a sparsity mask over a discrete set of sparse matrices is difficult and expensive. To address this, our main insight is to optimize over a continuous superset of sparse matrices with a fixed structure known as products of butterfly matrices. As butterfly matrices are not hardware efficient, we propose simple variants of butterfly (block and flat) to take advantage of modern hardware. Our method (Pixelated Butterfly) uses a simple fixed sparsity pattern based on flat block butterfly and low-rank matrices to sparsify most network layers (e.g., attention, MLP). We empirically validate that Pixelated Butterfly is 3x faster than butterfly and speeds up training to achieve favorable accuracy--efficiency tradeoffs. On the ImageNet classification and WikiText-103 language modeling tasks, our sparse models train up to 2.5x faster than the dense MLP-Mixer, Vision Transformer, and GPT-2 medium with no drop in accuracy.

研究动机与目标

• 解决现有稀疏训练方法依赖动态稀疏掩码或非结构化稀疏性所导致的低效与高开销问题。
• 通过将稀疏模式与块优化的内存访问对齐，克服稀疏矩阵运算中的硬件低效问题。
• 设计一种适用于多种神经网络层（尤其是注意力机制和 MLP 组件）的统一静态稀疏模式。
• 在不牺牲模型准确率的前提下，实现高训练加速，尤其适用于过参数化的模型。
• 通过简化实现并降低计算开销，推动稀疏训练在实际中广泛采用。

提出的方法

• 提出一种在稀疏矩阵超集上的连续优化框架，利用蝴蝶矩阵的乘积形式，以简化稀疏掩码的搜索过程。
• 引入平坦蝴蝶矩阵，作为标准蝴蝶矩阵的一阶近似，通过将序列乘积替换为求和操作，提升并行化能力。
• 设计具有固定、硬件对齐稀疏模式的块蝴蝶矩阵，以提升 GPU 上的内存访问效率。
• 将平坦块蝴蝶矩阵与低秩组件结合，以同时捕捉全局与局部特征表示。
• 在整个训练过程中使用静态稀疏模式，避免动态掩码更新，从而降低训练开销。
• 将该方法应用于 Transformer 和 MLP-Mixer 模型中的注意力层与 MLP 层的稀疏化，实现高效且可扩展的训练。

实验结果

研究问题

• RQ1 能否设计一种静态、硬件对齐的稀疏模式，在稀疏神经网络中同时实现高训练效率与强模型准确率？
• RQ2 平坦块蝴蝶矩阵在不同神经网络架构下的近似误差与稠密模型相比如何？
• RQ3 在不造成准确率下降的前提下，采用块对齐稀疏模式的稀疏训练在实际运行时间上相比稠密训练能提升多少？
• RQ4 统一的稀疏模式是否能有效泛化到注意力机制和 MLP 层等多样化的神经网络组件？
• RQ5 蝴蝶结构与低秩组件的结合是否在近似能力和效率上优于非结构化或标准蝴蝶稀疏性？

主要发现

• Pixelated Butterfly 在 ImageNet 和 WikiText-103 上实现比 Vision Transformer、GPT-2 medium 和 MLP-Mixer 等稠密模型快达 2.5 倍的训练速度，且准确率无下降。
• 与标准蝴蝶矩阵训练相比，该方法速度快 3 倍，主要得益于更高的硬件利用率和更强的并行化能力。
• 平坦蝴蝶矩阵通过将序列乘积替换为求和操作，实现了高效且可并行化的矩阵乘法。
• 块蝴蝶矩阵与硬件内存块对齐，显著提升了内存访问效率并降低了延迟。
• 平坦块蝴蝶矩阵与低秩组件的结合在近似准确率和训练速度上持续优于其他稀疏模式。
• 该方法在过参数化模型上保持了强大的泛化性能，支持稠密网络中观察到的双下降现象。