【論文瞬読】Lottery Ticket Hypothesis: 構造的スパース性による新たなブレークスルー
こんにちは！株式会社AI Nestです。
最近、面白い論文を読んだのでその内容をシェアしたいと思います。題名は「Coarsening the Granularity: Towards Structurally Sparse Lottery Tickets」。ニューラルネットワークの枝刈り (pruning) に関する研究なんですが、lottery ticket hypothesis (LTH) という仮説に新しい発見をもたらした論文です。
タイトル：Coarsening the Granularity: Towards Structurally Sparse Lottery Tickets
URL：https://arxiv.org/abs/2202.04736
コード：https://github.com/VITA-Group/Structure-LTH
著者：Tianlong Chen, Xuxi Chen, Xiaolong Ma, Yanzhi Wang, Zhangyang Wang
Lottery Ticket Hypothesis (LTH) とは？
まずはLTHについて簡単に説明しますね。LTHは、Jonathan Frankle氏とMichael Carbin氏によって2018年に提唱された仮説です。大規模でオーバーパラメトライズされたニューラルネットワークの中に、初期化時点から存在する、コンパクトでスパースでありながら同等の性能を発揮できるサブネットワーク (winning ticket) が存在するというものです。つまり、ニューラルネットワークの学習はwinning ticketを見つける作業とも捉えられるわけです。
このLTHが注目されている理由の一つは、現代のディープラーニングモデルが巨大化し、計算コストや環境負荷が問題になっている中で、モデルを小さくする (つまり枝刈りする) ための理論的根拠を与えてくれるからです。で、そのwinning ticketを見つける方法として知られているのが、iterative magnitude pruning (IMP) という手法。重みの絶対値が小さいものから順に一定の割合で枝刈りしていくというシンプルな方法です。
IMPの問題点：非構造的スパース性
で、ここまではLTHのごく基本的な話なんですが、実はIMPで見つかるwinning ticketには一つ問題がありました。それは、非構造的 (unstructured) なスパース性しか持たないということ。ここで、スパース性について少し詳しく説明しますね。
スパース性 (枝刈りされた重みの割合) に対してプロットしたグラフ
スパース性というのは、ニューラルネットワークの重みのうち、ゼロである (つまり枝刈りされている) 割合のことを指します。で、このスパース性にも2種類あって、一つは非構造的スパース性。これは、枝刈りされた重みが不規則に分布しているような状態です。もう一つは構造的スパース性。これは、枝刈りされた重みがある規則性を持って分布しているような状態で、例えば、特定のニューロンやチャネル、あるいはそれらのグループ単位で重みがゼロになっているような状態です。
で、問題は、非構造的スパース性を持つモデルは、ハードウェアでの高速化が難しいということ。現代のGPUやASICは、レギュラーな計算に最適化されているので、不規則なスパース性を持つモデルでは、その計算力を十分に活用できないんです。一方、構造的スパース性を持つモデルなら、ハードウェアでの高速化が可能になります。
でも、従来のIMPでは、この構造的スパース性を持つwinning ticketが見つからないと考えられてきました。これは、IMPが重みの大きさのみに基づいて枝刈りを行うため、枝刈り後のスパース性が不規則になりがちだからです。
Refilling と Regrouping：構造的スパース性を獲得する手法
そこでこの論文では、IMPで得られた非構造的スパースネットワークに対して、refillingとregroupingという2つの手法を適用することで、構造的スパース性を持つwinning ticketが発見できることを示しました。
Refillingは、重要度の高いチャネルの枝刈りされた要素を復活させることでチャネル単位の構造的スパース性を獲得する手法です。具体的には、まずIMPでスパースなモデルを得た後、各チャネルの重要度を何らかの指標 (例えば、そのチャネルに属する重みの絶対値の和) で評価します。そして、重要度の高いチャネルについては、枝刈りされた重みを復活させ (つまり、ゼロでない値に戻し)、重要度の低いチャネルについては、完全に削除します。このようにして、チャネル単位の構造的スパース性が得られるわけです。
Regroupingは、ハイパーグラフ分割を用いることでさらに細かい単位のグループ構造的スパース性を獲得する手法です。ハイパーグラフは、通常のグラフを拡張したもので、一つの辺が2つ以上の頂点を結ぶことができるようなグラフです。このハイパーグラフ上で、頂点をニューロンや重み、辺をそれらの関係性 (例えば、同じチャネルに属しているか) に対応づけることで、ニューラルネットワークの構造をグラフ表現することができます。そして、このグラフを適切に分割することで、密に結合されたニューロンやチャネルのグループを見つけ出すことができるわけです。Regroupingでは、このグループ単位で枝刈りを行うことで、より細かい単位の構造的スパース性を実現しています。
提案手法の意義と実験結果
この発見は、LTHをより実用的なものにするための重要なステップだと思います。構造的スパース性を持つモデルは、非構造的スパースモデルに比べてハードウェア上でより効率的な計算が可能になるからです。特に、モバイルやエッジデバイスでのAI実行が重要になってきている昨今、モデルの計算効率は非常に重要な問題です。この論文の成果は、そうした要求に応える一つの解になり得ると思います。
スパース性と精度の関係を示したグラフ
また、実験でも、CIFAR-10/100、Tiny-ImageNet、ImageNetなどの複数のデータセットと、Wide-ResNet、ResNet、MobileNet、VGGなどの様々なネットワークアーキテクチャで提案手法の有効性が示されていて、手法の汎用性の高さがうかがえます。例えば、ImageNetで学習したResNet-50では、58%のスパース性 (つまり、58%の重みがゼロ) を持つ構造的スパースモデルが、dense (枝刈りしていない) モデルと同等の精度を達成しながら、60%以上の推論時間短縮を実現しています。これは、非常に印象的な結果だと思います。
感想と今後の展望
個人的には、Refillingではチャネルの重要度をどのように測るか、Regroupingではどのようにグループを形成するかといった部分に自由度がありそうで、さらなる改善の余地がありそうだなと思いました。例えば、チャネルの重要度の指標として、重みの大きさだけでなく、勾配の大きさや、そのチャネルを枝刈りした時の精度への影響なども考慮できるかもしれません。また、Regroupingでは、グラフ分割のアルゴリズムや、辺の重み付けの方法など、工夫の余地がありそうです。
また、この論文はLTHの理解を深め、応用の可能性を広げるものだと思います。構造的スパース性を活用したモデル圧縮・高速化技術の発展にもつながることが期待されます。今後は、lottery ticketを発見するための様々な手法の提案や、チャネル単位・グループ単位以外の構造的スパース性の探索なども行われるかもしれませんね。
加えて、LTHのより理論的な理解も進むことを期待しています。なぜ初期化時点から良いサブネットワークが存在するのか、どのような初期化や学習アルゴリズムがwinning ticketの発見に適しているのか、といった問いは、まだ十分に解明されていません。この論文の成果を足がかりに、そうした理論的な問いにも迫っていけたらいいなと思います。