ChatGPT の脳内を可視化

トランスフォーマーのアテンションとは

ChatGPTやBingAIなどの人気システムを支えるモデルとして、トランスフォーマーは世界を席巻しています。トランスフォーマーのニューラルネットワークアーキテクチャは、NLPやコンピュータビジョンの場面で使用され、さまざまなタスクで大幅な性能向上を達成しています。

トランスフォーマーの成功の鍵は、自己注視メカニズムにあり、これにより、シーケンスの要素間の豊かで文脈的な関係を学習することができます。例えば、"the sky is blue "という文章では、"sky "と "blue "の間には高い注意が、"the "と "blue "の間には低い注意が向けられると考えられます。
日本語で言う所の「て」「に」「を」「は」ですね。

https://catherinesyeh.github.io/attn-docs/　より引用

注目度を計算するために、まず各入力要素（例：文中の単語や画像のパッチ）を対応するクエリとキーベクトルに変換する。高いレベルでは、2つの単語や画像パッチの間の注目度は、対応するクエリとキーの間の内積関数として見ることができます。このように、注意はしばしば二部グラフ表現（言語の場合）で可視化され、各辺の不透明度がクエリとキーのペア間の注意の強さを示します。

アテンションの「グローバルビュー」の提案がGPTの突破口

注意のパターンをスケールで分析・合成するというこの課題に対処するために、トランスフォーマーの注意のグローバルビューのアイデアが重要な鍵となります。このグローバルビューは、新しい可視化技術を設計し、それを応用して、トランスフォーマーモデルのアテンションを探索するためのインタラクティブなツールを構築することによって作成されます。

AttentionViz の動作中のスナップショット
https://github.com/catherinesyeh/attention-viz　でソースDLしてビルドできます

技法
トランスフォーマーの各アテンションヘッドについて、入力シーケンスのセットを対応するクエリーベクトルとキーベクトルに変換し、高次元空間におけるジョイントエンベッディングを作成します。この結合埋め込み空間では、距離が注目度の妥当な代用品であることが判明しました。t-SNEやUMAPなどの手法を用いて、この埋め込みを2次元または3次元で可視化し、注目パターンの「グローバル」なビューを提供します。

このジョイントエンベッディング技術を使って、自己の注意パターンをスケールで可視化するインタラクティブなツール、AttentionVizが開発されました。AttentionVizは、言語と視覚の両方の変換器に対して、複数のレベルで注意の探索を可能にします。現在、BERT（言語）、GPT-2（言語）、ViT（視覚）をサポートしています。AttentionVizへの入力例を以下に示します（実際には、多くの文章や画像を用いて、各共同クエリーキー埋め込みを形成します）。

MatrixView

AttentionVizから得られた知見の例
AttentionVizにより、言語と視覚のトランスフォーマーにおける自己注意について、いくつかの興味深い知見を発見しました。以下にいくつかの例を紹介します。

色相・明度の特化
ViTでは、色や明るさに基づくパターンに特化した視覚的注意の頭部があるかどうかが気になりました。これをテストするために、合成色/明るさのグラデーション画像でデータセットを作成し、結果のクエリとキートークンをAttentionVizにロードしました。

その結果、白黒画像のトークンを明るさに基づいて整列させるヘッド（レイヤー0ヘッド10）と、カラフルなパッチを色相に基づいて整列させるヘッド（レイヤー1ヘッド11）が発見されました。このデータセットには、あらゆる方向の色と明るさのグラデーション画像が含まれており、元画像の位置に関係なく、類似したパッチが共同埋め込み空間内でクラスタリングされていることがわかります。画像ビューの注意ヒートマップは、これらの発見を裏付けています。トークンは、同じ色や明るさの他のトークンに最も注意を払います。

グローバルな注意の痕跡
BERT を探索している間、ユニークで識別可能な形状を持ついくつかの注意の頭部を観察した。例えば、初期のモデル層では、いくつかの螺旋状のプロット（例えば、層3の頭部9）が見受けられます。トークン位置で色分けすると、螺旋の外側から内側に向かうにつれてトークン位置が高くなるという、位置的な傾向が明らかになりました。センテンスビューでは、「次のトークン」に注目するパターンがあることが確認できます。

らせん状のプロット

同様に、「小さな塊」を持つプロットも位置のパターンを符号化していることに気づいた（例：レイヤー2ヘッド0）。これは、各トークンをその位置mod 5で色分けすることで確認でき、より離散的な位置の配色を形成し、文の位置の小さなオフセットに基づくクエリとキーペア間の関係を視覚化するのに役立つと思われます。「スパイラル」と「クランプ」の主な違いは、トークンが1つ離れた位置の他者に選択的に注意を払うか、複数の異なる位置にいる可能性があるかという点にあると思われます。

BERTで誘導ヘッド

AttentionVizを誘導ヘッドの探索に使用できるかどうかも興味深い視点です。高度なレベルでは、誘導ヘッドは言語変換器が文脈内学習を行うのを助けるために、繰り返されるシーケンスに対して前方一致とコピーを実行します。誘導ヘッドはGPT-2のような一方向のモデルでのみ研究されてきましたが、AttentionVizにより、双方向の注意メカニズムを使用するBERTにおいても、潜在的な誘導ヘッドの動作を発見しました。その1つであるレイヤー8ヘッド2は、トークンA（例：-）が、その前に出現したトークンB（例：8または10）に注意を払うという、標準的なコピー動作を示しているように見えます。以下の例では、各 A が複数の B に注目できることも示している。BERT は双方向性であるため、双方向のコピーを実行することもできる。

もう一つの頭部、レイヤー9の頭部9は、「逆」誘導頭部となる可能性があるようです。この場合、トークンAは、別のオカレンスでその後に来たトークンBに注意を払うことになります。BERTのような双方向変換器における誘導ヘッドと文脈内学習の可能性を示唆します。

GPTに限らず、その基盤となるCNNですら
「なぜこうすると上手く機能するのか？」
を開発当事者ですら説明できていないのが現状です。
つまり人類は原理よりも先に結果を先に手に入れてしまったのです。
上手く機能してしまった結果から逆算して原理を追求する際に、補助するAttentionVizような「可視化ツール」の重要度が急騰してます。

AttentionViz はGitHubでソース公開されており誰でも可視化に参加できます。