Optimal Brain Apoptosis
作者: Mingyuan Sun, Zheng Fang, Jiaxu Wang, Junjie Jiang, Delei Kong, Chenming Hu, Yuetong Fang, Renjing Xu
发布时间: 2025-02-27
来源: arxiv
研究方向: 深度学习模型压缩与剪枝
主要内容
本文提出了一种新的神经网络剪枝方法——Optimal Brain Apoptosis (OBA)，旨在通过直接计算Hessian-vector product来精确估计参数的重要性，从而在不显著降低模型性能的前提下，减少卷积神经网络（CNNs）和Transformer模型的参数量和计算复杂度。该方法基于Optimal Brain Damage (OBD)的思想，通过分解Hessian矩阵并计算每层之间的Hessian子矩阵，提出了一种高效的计算二阶泰勒展开的方法。
主要贡献
1. 提出了Optimal Brain Apoptosis (OBA)方法，通过直接计算Hessian-vector product来估计参数的重要性，避免了传统方法中对Hessian矩阵的近似。
2. 通过分解Hessian矩阵并分析层间Hessian子矩阵的非零条件，提出了一种高效计算二阶泰勒展开的方法。
3. 在多个数据集（如CIFAR10、CIFAR100和ImageNet）和模型（如VGG19、ResNet32、ResNet50和ViT-B/16）上验证了OBA的有效性。
4. 提出了Jacobian-Vector Product Forward Propagation (JVPF)方法，用于高效计算层间的Hessian-vector product。
研究方法
1. Hessian矩阵分解与层间Hessian子矩阵分析
2. Hessian-vector product直接计算
3. Jacobian-Vector Product Forward Propagation (JVPF)
4. 结构化剪枝与非结构化剪枝策略
实验结果
实验结果表明，OBA在多个数据集和模型上均表现出色。在ImageNet数据集上，OBA在ResNet50模型上实现了2倍的加速，且准确率仅下降0.53%。在ViT-B/16模型上，OBA实现了1.3倍的加速，准确率下降1.43%，优于其他方法。在CIFAR10和CIFAR100数据集上，OBA在ResNet32模型上的表现也优于其他剪枝方法，尤其是在FLOPs减少到6%时，准确率损失仅为0.79%。
未来工作
未来的研究可以探索将OBA应用于更复杂的网络结构，如RNN和State Space Models，这些模型处理时间序列数据时Hessian矩阵的计算更为复杂。此外，进一步优化计算效率和内存消耗也是未来的研究方向。