[論文レビュー] 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）を提唱する。
• 接続スキームを、どのブロックが自己注意モジュールに接続されるかを示すバイナリーベクトルとしてモデル化し、3つの基準（高速な推論、パラメータ数の削減、精度の維持）を満たすスキームを「チケット」と定義する。
• ポリシー勾配法を用いた強化学習コントローラーを設計し、最適な接続スキームを探索する。複合報酬関数を用い、スパarsity、検証精度、およびランダムネットワーク差異化（RND）を用いた好奇心駆動型探索を統合する。
• サンプルされた接続スキームを再利用することで、コントローラーの学習効率を向上させるために、近接ポリシー最適化（PPO）を組み込む。
• サブネットワークの性能を推定するスーパーネットを採用し、スーパーネット推定精度と独立した検証精度のピアソン相関係数（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}} $ はスパarsityを促進し、$ g_{\text{val}} $ は高い検証精度に報酬を与え、$ g_{\text{rnd}} $ は新規なスキームの探索を促進する。

実験結果

リサーチクエスチョン

• RQ1 スパースな自己注意接続スキームは、パラメータ数と推論時間を削減しながら、完全な自己注意ネットワークと同等またはそれ以上の精度を達成できるか？
• RQ2 完全な自己注意ネットワーク内に、顕著に少ない接続数で精度を維持するサブネットワーク（「ロットリーチ」）が存在するか？
• RQ3 提案された強化学習ベースの探索手法は、さまざまなデータセットおよびバックボーンアーキテクチャにおいて、このような最適なスパース接続スキームを効果的に発見できるか？
• RQ4 発見された接続スキームは、元の学習分布とは異なる下流のビジョンタスクに一般化可能か？

主な発見

• 実験的結果から、特定のスパースな自己注意接続スキームが完全な自己注意ネットワークを上回る精度を達成することが示された。これは、接続数が多いほど性能が向上するという仮定に反する。
• 提案された強化学習手法により、CIFAR-100およびImageNetでパラメータ数の増加を抑え、推論を高速化しながらも、精度を維持または向上させる接続スキームを効果的に発見できた。
• CIFAR-100では、スーパーネット推定精度と独立した検証精度の相関係数が0.71であった。これにより、探索プロセスにおけるスーパーネットベースの報酬信号の信頼性が裏付けられた。
• RNDの好奇心報酬の導入により、探索が促進され、コントローラーの学習における早期収束のリスクが軽減された。
• 発見されたチケットは、群衆数え上げやセマンティックセグメンテーションを含む他のビジョンタスクにも一般化可能であり、移行性が確認された。
• ResNet50を用いたImageNet2012では、完全な自己注意ネットワークと比較して顕著に計算コストとパラメータ数を削減しながら、競争力のある精度を達成した。