[論文レビュー] White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?

Yaodong Yu, Sam Buchanan+8 more|arXiv (Cornell University)|Nov 22, 2023

Image and Signal Denoising Methods被引用数 4

ひとこと要約

本稿では、スパースレートリダクションに基づく白箱型トランスフォーマー構造であるCRATEを提案する。このアーキテクチャは、表現の圧縮とスパarsityを同時に最適化する。マルチヘッド自己注意機構をコーディングレートの勾配降下法として解釈し、MLPをスパースコーディングのプロキシマル勾配降下法としてモデル化することで、CRATEは視覚および言語タスクにおいて最先端のトランスフォーマーと同等の性能を達成する。これは、圧縮が効果的な表現学習の中心的役割を果たしていることを示している。

ABSTRACT

In this paper, we contend that a natural objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a low-dimensional Gaussian mixture supported on incoherent subspaces. The goodness of such a representation can be evaluated by a principled measure, called sparse rate reduction, that simultaneously maximizes the intrinsic information gain and extrinsic sparsity of the learned representation. From this perspective, popular deep network architectures, including transformers, can be viewed as realizing iterative schemes to optimize this measure. Particularly, we derive a transformer block from alternating optimization on parts of this objective: the multi-head self-attention operator compresses the representation by implementing an approximate gradient descent step on the coding rate of the features, and the subsequent multi-layer perceptron sparsifies the features. This leads to a family of white-box transformer-like deep network architectures, named CRATE, which are mathematically fully interpretable. We show, by way of a novel connection between denoising and compression, that the inverse to the aforementioned compressive encoding can be realized by the same class of CRATE architectures. Thus, the so-derived white-box architectures are universal to both encoders and decoders. Experiments show that these networks, despite their simplicity, indeed learn to compress and sparsify representations of large-scale real-world image and text datasets, and achieve performance very close to highly engineered transformer-based models: ViT, MAE, DINO, BERT, and GPT2. We believe the proposed computational framework demonstrates great potential in bridging the gap between theory and practice of deep learning, from a unified perspective of data compression. Code is available at: https://ma-lab-berkeley.github.io/CRATE .

研究の動機と目的

• データ圧縮とスパarsityに基づく、原理的かつ理論的根拠に基づいた表現学習のフレームワークを確立すること。
• スパースレートリダクションという統一的目的関数における反復的最適化の観点から、トランスフォーマーの成功を説明すること。
• 第一原理から出発して、完全に解釈可能で数学的に根拠を持つ深層ネットワークアーキテクチャ（CRATE）を導出すること。
• 可逆的かつ構造的なノイズ除去と拡散を用いて、エンコーダとデコーダの機能を1つのアーキテクチャに統合すること。
• 視覚および言語タスクにおいて、CRATEが意味的で構造的な表現を学習することを実証的に検証すること。

提案手法

• 表現学習を、内在的情報ゲインと外在的スパarsityを組み合わせた測度であるスパースレートリダクションの最適化として形式化する。
• 入力トークンのコーディングレートに対する近似勾配降下ステップとして、マルチヘッド自己注意機構を導出する。
• 特徴のスパースコーディングのためのプロキシマル勾配降下ステップとして、マルチレイヤーパーセプトロンをモデル化する。
• 数学的解釈可能性を保証するために、交互に配置された注意とスパースコーディングブロックのスタックとしてCRATEを構築する。
• 同じCRATEブロックを用いた構造的ノイズ除去拡散により、圧縮プロセスを逆転させることで可逆的デコードを実現する。
• モジュラーなコンponents（MSSA：マルチヘッド自己注意、ISTA：スパarsityのための反復的ス hrinkageスレッディング）を用いて、PyTorchでCRATEを実装する。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーの挙動は、情報理論とスパarsityに根ざした統一的最適化目的関数を通じて説明可能か？
• RQ2 原理的な圧縮目的から直接導出することで、深層ネットワークを白箱的かつ解釈可能なシステムとして設計できるか？
• RQ3 スパースレートリダクションから導出されたネットワークは、高度に最適化されたトランスフォーマー・モデルと同等の性能を達成できるか？
• RQ4 同じアーキテクチャが、構造的ノイズ除去を用いた可逆的で構造的なデノイジングにより、エンコーダとデコーダの両方として機能できるか？
• RQ5 CRATEにおける注意マップと学習された特徴は、オブジェクトレベルのセグメンテーションのような、自己組織的構造的意味的特徴を示すか？

主な発見

• CRATEは、ViT、MAE、DINO、BERT、GPT2といった最先端モデルと同等の性能を、画像およびテキストベンチマークで1〜2%の差で達成する。
• CRATEのレイヤーにおける注意マップは、オブジェクトの境界や顕著な特徴に注目が集中するなど、自己組織的意味的セグメンテーションを示す。
• レイヤーごとの分析により、CRATEが階層的かつ構造的な表現を学習していることが確認され、初期レイヤーが低レベルのパターンを捉え、深いレイヤーが高レベルの意味論を捉えている。
• CRATEエンコーダの逆関数は構造的ノイズ除去拡散により実現可能であり、同じアーキテクチャがエンコーディングとデコーディングの両方をサポートしていることを示している。
• アブレーションスタディにおいてもモデルの性能が安定しており、スパースレートリダクション目的関数の堅牢性が確認された。
• 学習された表現は単純で構造的であり、特徴が非一貫部分空間に集中しており、理論的圧縮目標と整合的である。