[論文レビュー] White-Box Transformers via Sparse Rate Reduction

Yaodong Yu, Sam Buchanan+6 more|arXiv (Cornell University)|Jun 1, 2023

Advanced Neural Network Applications被引用数 23

ひとこと要約

本論文は、トランスフォーマー風の層を、疎なレート削減目的を最適化する展開ステップとして統合し、完全に解釈可能なホワイトボックスアーキテクチャ（CRATE）を生み出し、トークン表現を圧縮・稀疎化し、大規模ビジョンデータ上で設計済みのトランスフォーマーと競合する性能を発揮する。

ABSTRACT

In this paper, we contend that the objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a mixture of low-dimensional Gaussian distributions supported on incoherent subspaces. The quality of the final representation can be measured by a unified objective function called sparse rate reduction. From this perspective, popular deep networks such as transformers can be naturally viewed as realizing iterative schemes to optimize this objective incrementally. Particularly, we show that the standard transformer block can be derived from alternating optimization on complementary parts of this objective: the multi-head self-attention operator can be viewed as a gradient descent step to compress the token sets by minimizing their lossy coding rate, and the subsequent multi-layer perceptron can be viewed as attempting to sparsify the representation of the tokens. This leads to a family of white-box transformer-like deep network architectures which are mathematically fully interpretable. Despite their simplicity, experiments show that these networks indeed learn to optimize the designed objective: they compress and sparsify representations of large-scale real-world vision datasets such as ImageNet, and achieve performance very close to thoroughly engineered transformers such as ViT. Code is at \url{https://github.com/Ma-Lab-Berkeley/CRATE}.

研究の動機と目的

• データ分布を低次元部分空間の混合へ圧縮し、表現を疎にすることとして表現学習を動機づける。
• 失われる符号化レートとスパース性を組み合わせた統一の疎なレート削減目的を導入し、コンパクトなトークン表現を学習する。
• 注意機構とMLPブロックの数学的解釈可能性を提供するよう、トランスフォーマー風の層を展開された最適化ステップとして導出する。
• CRATE（Coding RAte TransformEr）を、データから学習された分布と辞書の層毎の確率モデルを備えたホワイトボックスアーキテクチャとして提案する。
• CRATEが大規模ビジョンデータで表現を圧縮・疎化することを学習でき、ViT風の性能に近づくことを示す。

提案手法

• 統一された目的を定義する：疎なレート削減を最大化し、レート削減とℓ0スパースペナルティを組み合わせ、Z = f(X)とする。
• 各層ごとに学習された基底U[K]を持つ低次元部分空間の混合としてトークン分布をモデル化する。
• サブスペース混合（MSSA）に対して符号化レートを最小化する勾配ステップとして、自己注意様の更新を導出する。
• Zの疎性を促進するため、学習された辞書Dに対するISTA様の更新で疎性を表現する。
• MSSAベースの圧縮を行い、その後ISTAベースの疎化を行う層を積み重ねてCRATEを構築し、層ごとのU[K]とDをエンドツーエンドで学習する。
• 再現用のコードリンクを提供: https://github.com/Ma-Lab-Berkeley/CRATE

実験結果

リサーチクエスチョン

• RQ1 スパース性を伴うレート削減目的は、コンパクトで解釈可能なトークン表現を生み出せるか？
• RQ2 展開型最適化から導出されたホワイトボックスなトランスフォーマー層は、大規模なビジョンタスクで競争力のある性能を達成するか？
• RQ3 自己注意とMLPブロックを、統一的な枠組みの中でノイズ除去/圧縮と疎な符号化のステップとして再解釈できるか？
• RQ4 層ごとに学習された部分空間基底および辞書が表現品質と転移性に与える影響は何か？
• RQ5 提案されたMSSAおよびISTAブロックが、訓練中の意図した最適化目的とどれだけ忠実に一致するか？

主な発見

• CRATE層は、部分空間の混合へ向けてトークン分布を圧縮し、表現を疎化する漸進的最適化を実装する。
• MSSA成分は、勾配ステップのような操作に対応し、サブスペースに対するレート削減デノイジングから導出され、自己注意に似ている。
• ISTAベースの疎化層は、学習済み辞書で疎性を促進し、レートベースの多様性の扱いやすい近似を可能にする。
• ImageNet-1Kの実験で、CRATEが表現を圧縮・疎化することを学習し、ViTのような設計済みトランスフォーマーに近い性能を示す。
• 層ごとの分析は、圧縮と疎化の両方が層を追うごとに改善されることを示し、目的指向の設計を裏付ける。