Nested Hierarchical Transformer:ViTに数行のコードを加えるだけでデータ効率を大幅に向上(1/3)

１．Nested Hierarchical Transformer:ViTに数行のコードを加えるだけでデータ効率を大幅に向上(1/3)まとめ

・Vision Transformer(ViT)とその亜種は視覚理解タスクで最近大きく注目されている
・オリジナルのViTは大量のデータを必要とするためデータ効率化手法が多く研究されている
・Nested Hierarchical Transformerは階層構造を再考しデータ効率性と解釈可能性を両立した

２．Nested Hierarchical Transformerとは？

以下、ai.googleblog.comより「Nested Hierarchical Transformer: Towards Accurate, Data-Efficient, and Interpretable Visual Understanding」の意訳です。元記事は2022年2月9日、Zizhao Zhangさんによる投稿です。あったまい～ですね～。

Hierarchicalでstackedな感じのアイキャッチ画像のクレジットはPhoto by Edvard Alexander Rølvaag on Unsplash

視覚理解では、Vision Transformer(ViT)とその亜種が最近大きく注目されています。ViTは画像分類、物体検出、映像理解など多くの主要な視覚アプリケーションで優れた性能を発揮します。

ViTの中核となる考え方は、画像を小さな断片(patch)に区分し、断片間の大域的な関係を学習するためにself-attentionレイヤーの力を利用することです。

しかし、断片同士の接続数は画像サイズに対して二次関数的に増加します。このような設計は、データ効率が悪いことが確認されています。オリジナルのViTは、事前学習用に数億枚の画像があれば畳み込みネットワークよりも性能が向上しますが、そのような大量データは必ずしも実用的ではなく、より少ないデータが与えられた場合は畳み込みネットワークと比較して性能は低下します。

そのため、畳み込み層の追加や局所的なself-attentionによる階層構造の構築など、視覚表現を効率的に学習できる、より適切なアーキテクチャの再設計を模索する動きが多く見られます。

階層構造の原理は視覚モデルの核となる考え方の一つで、下位層は高次元の画素空間でより局所的な物体構造を学習し、上位層は低次元の特徴空間でより抽象的で高度な知識を学習するというものです。

既存のViTベースの手法は、このような階層を実現するためにself-attention層の内部に様々な修正を加えることに焦点を当てていますが、これらは有望な性能向上をもたらす一方で、しばしば大幅なアーキテクチャの再設計を必要とします。また、これらのアプローチは解釈可能な設計ではないため、学習済みモデルの内部構造を説明することが困難です。

これらの課題を解決するために、論文「Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding」では、既存の階層構造駆動型設計を再考し、それらを大幅に簡略化する新規かつ使い勝手が良いアプローチを紹介します。

この研究の中心的なアイデアは、特徴表現学習と特徴抽象化(プーリング)部分を切り離したことです。

nested transformerレイヤーは、画像の断片の視覚的知識を個別に符号化し、その後、処理された情報を集約します。この処理を階層的に繰り返すことで、ピラミッド型のネットワーク構造を実現しています。

その結果、ImageNetで競争力のある結果を達成し、データ効率を図るベンチマークで優れた結果を得ることができました。このような設計により、より速い収束でデータ効率を有意に向上させ、貴重な解釈可能性の利点を提供できることを示しました。さらに、推論時に学習済みモデルの決定過程を解釈するための新しい手法であるGradCATを紹介します。

アーキテクチャの設計

このアーキテクチャは、オリジナルのViTのソースコードにPythonのコードを数行追加するだけで、簡単に実装することができます。

オリジナルのViTでは、入力画像を小さな断片に分割し、各断片の画素をあらかじめ定義された次元数を持つベクトルに投影します。そして、そのベクトル列を、同一のtransformer層を複数積層したViTアーキテクチャに供給します。

ViTの各層は画像全体の情報を処理しますが、本方式では、空間的に隣接するいくつかの画像断片を含む画像の領域(ブロック)のみを処理するために、積層したtransformerを使用します。このステップはブロックごとに独立しており、特徴表現学習が行われる場所でもあります。

最後に、ブロック集約と呼ばれる新しい計算レイヤーが、空間的に隣接するすべてのブロックを結合します。各ブロックを集約した後、空間的に隣接する4つのブロックに対応する特徴表現は、transformer層を積み重ねた別のモジュールに送られ、これらの4つのブロックが共同で処理されます。

この設計により、ネットワークはピラミッド型の階層構造となり、下位層は局所的特徴(質感など)に、上位層は大域的特徴(物体の形状など)に、ブロック集約により低次元化された状態に集中する事が可能となります。

ネットワークが画像を処理する様子を可視化したもの
入力画像が与えられると、ネットワークはまず画像をブロックに分割し、各ブロックに4つの画像断片を格納します。各ブロックの画像断片は、線形に射影されベクトルとなり、同一のtransformer層を積み重ねることで処理されます。次に、提案するブロック集約層は、各ブロックの情報を集約し、その空間サイズを1/4倍に縮小します。最上位階層でブロック数を1にし、その出力後に分類を行います。