[论文解读] The Lottery Ticket Hypothesis for Self-attention in Convolutional Neural Network

Zhongzhan Huang, Senwei Liang+4 more|arXiv (Cornell University)|Jul 16, 2022

Advanced Neural Network Applications被引用 8

一句话总结

本文提出了自注意力机制的彩票理论（LTH4SA），证明了稀疏自注意力连接方案——即仅将卷积神经网络（CNN）中的部分模块连接至自注意力模块——可在减少参数量和推理时间的同时，达到或超越完整自注意力网络的准确率。该方法采用基于强化学习的搜索策略，识别出这些最优的稀疏连接模式，并在CIFAR-100和ImageNet上验证了其性能达到当前最先进水平。

ABSTRACT

Recently many plug-and-play self-attention modules (SAMs) are proposed to enhance the model generalization by exploiting the internal information of deep convolutional neural networks (CNNs). In general, previous works ignore where to plug in the SAMs since they connect the SAMs individually with each block of the entire CNN backbone for granted, leading to incremental computational cost and the number of parameters with the growth of network depth. However, we empirically find and verify some counterintuitive phenomena that: (a) Connecting the SAMs to all the blocks may not always bring the largest performance boost, and connecting to partial blocks would be even better; (b) Adding the SAMs to a CNN may not always bring a performance boost, and instead it may even harm the performance of the original CNN backbone. Therefore, we articulate and demonstrate the Lottery Ticket Hypothesis for Self-attention Networks: a full self-attention network contains a subnetwork with sparse self-attention connections that can (1) accelerate inference, (2) reduce extra parameter increment, and (3) maintain accuracy. In addition to the empirical evidence, this hypothesis is also supported by our theoretical evidence. Furthermore, we propose a simple yet effective reinforcement-learning-based method to search the ticket, i.e., the connection scheme that satisfies the three above-mentioned conditions. Extensive experiments on widely-used benchmark datasets and popular self-attention networks show the effectiveness of our method. Besides, our experiments illustrate that our searched ticket has the capacity of transferring to some vision tasks, e.g., crowd counting and segmentation.

研究动机与目标

• 探究在卷积神经网络中，稀疏自注意力连接是否能在降低计算成本的同时，保持或超越完整自注意力网络的准确率。
• 识别出自注意力模块与CNN模块块之间的连接方案，使模型在减少参数量和加快推理速度的同时保持准确率。
• 开发一种基于强化学习的方法，高效搜索此类最优稀疏连接方案。
• 验证所发现连接方案在不同视觉任务（包括人群计数和语义分割）中的泛化能力。

提出的方法

• 提出自注意力机制的彩票理论（LTH4SA），主张完整自注意力网络中存在一个具有稀疏连接的子网络，该子网络在减少参数量和推理时间的同时保持准确率。
• 将连接方案建模为一个二值向量，指示哪些模块块与自注意力模块相连，定义满足以下三个标准的方案为“彩票”：推理速度更快、参数量更少、准确率保持不变。
• 设计一个基于策略梯度的强化学习控制器，用于搜索最优连接方案，采用复合奖励函数，结合稀疏性、验证准确率以及通过随机网络蒸馏（RND）实现的探索性驱动机制。
• 引入近端策略优化（PPO）以重用采样的连接方案，提升控制器训练效率。
• 使用超网络（supernet）估算子网络性能，并通过超网络与独立验证准确率之间的皮尔逊相关系数（0.71）验证奖励信号的可靠性。
• 采用混合奖励函数：$ G(\mathbf{a}) = \lambda_1 \cdot g_{\text{spa}} + \lambda_2 \cdot g_{\text{val}} + \lambda_3 \cdot g_{\text{rnd}} $，其中 $ g_{\text{spa}} $ 促进稀疏性，$ g_{\text{val}} $ 奖励高验证准确率，$ g_{\text{rnd}} $ 鼓励探索新颖方案。

实验结果

研究问题

• RQ1 是否存在一种稀疏自注意力连接方案，能在减少参数量和推理时间的同时，实现比完整自注意力网络更高的准确率或相当的性能？
• RQ2 在完整自注意力网络中是否存在一个子网络，其连接数显著减少但准确率保持不变，从而支持‘彩票’存在的合理性？
• RQ3 所提出的基于强化学习的搜索方法是否能有效发现不同数据集和主干网络架构下的最优稀疏连接方案？
• RQ4 所发现的连接方案是否能泛化至原始训练分布之外的下游视觉任务？

主要发现

• 实证结果表明，某些稀疏自注意力连接方案的准确率甚至高于完整自注意力网络，这与‘连接越多性能越好’的假设相矛盾。
• 所提出的强化学习方法成功发现了能减少参数量增量并加速推理速度的连接方案，同时在CIFAR-100和ImageNet上保持或提升了准确率。
• 在CIFAR-100上，超网络估计值与独立验证准确率之间的相关系数为0.71，支持在搜索过程中使用基于超网络的奖励信号。
• 引入RND探索奖励后，控制器的探索能力得到提升，有效缓解了训练过程中的过早收敛问题。
• 所发现的‘彩票’方案可泛化至其他视觉任务，包括人群计数和语义分割，展现出良好的可迁移性。
• 在ImageNet2012上使用ResNet50时，该方法实现了具有竞争力的准确率，同时显著降低了计算和参数开销，相比完整自注意力网络具有明显优势。