行列乗算の3Dインサイト: これがAIの思考法

行列乗算の実行プロセスを 3D で表示できれば、行列乗算を学ぶのはそれほど難しくないでしょう。

今日では、行列乗算は機械学習モデルの構成要素となり、さまざまな強力な AI テクノロジーの基礎となっています。行列乗算の実行方法を理解することは、AI とますますインテリジェント化が進むこの世界に対する理解を深める上で間違いなく役立ちます。

PyTorch ブログのこの投稿では、行列乗算と行列乗算の組み合わせを視覚化するツールである mm を紹介しています。

mm は 3 つの空間次元すべてを使用するため、従来の 2 次元図よりも認知オーバーヘッドが少なく、特に (ただし、それだけではありません) 視覚的および空間的思考が得意な人にとって、アイデアを直感的に表示して刺激するのに役立ちます。

また、行列乗算を組み合わせるための 3 つの次元と、トレーニング済みの重みをロードする機能を備えているため、mm は大きな複合式 (アテンションヘッドなど) を視覚化し、実際の動作を確認できます。

mm は完全にインタラクティブで、ブラウザまたはノートブックの iframe で実行され、完全な状態を URL に保存するため、リンクは共有可能なセッションになります (この記事のスクリーンショットとビデオには、ツール内の対応する視覚化を開くリンクがあります。詳細については、元のブログを参照してください)。このリファレンスガイドでは、利用可能なすべての機能について説明します。

ツールアドレス: https://bhosmer.github.io/mm/ref.html
オリジナルブログ: https://pytorch.org/blog/inside-the-matrix

この記事では、まず視覚化の方法を紹介し、いくつかの簡単な行列の乗算と式を視覚化することで直感を養い、次にいくつかの拡張された例について詳しく説明します。

はじめに: この視覚化が優れているのはなぜでしょうか?
ウォームアップ: アニメーション - 標準行列乗算分解の動作を確認する
ウォームアップ: 表現 - 基本的な表現の構成要素を簡単に見てみましょう
注意ヘッドの詳細: NanoGPT を使用して GPT-2 の注意ヘッドのペアの構造、値、計算動作を詳しく調べます。
注意の並列化: 最近の Blockwise Parallel Transformer 論文の例を使用して、注意ヘッドの並列化を視覚化します。
注意層のサイズ: 注意層全体を単一の構造として視覚化すると、注意層の MHA 半分と FFA 半分は一緒にどのように見えるでしょうか?自己回帰デコード処理中に画像はどのように変化するのでしょうか?
LoRA: この注意ヘッドアーキテクチャの詳細説明の視覚的説明

1 はじめに

mm の視覚化アプローチは、行列乗算が本質的に 3 次元の演算であるという前提に基づいています。

言い換えると：

実際、それは次のように説明できます。

このように行列乗算をキューブにラップすると、引数の形状、結果の形状、共有ディメンション間の正しい関係がすべて整います。

行列の乗算は、これで幾何学的に意味を持ちます。結果行列の各位置 i,j は、立方体の内部で深さの次元 k に沿って走るベクトルを固定します。ここで、L の i 番目の行から伸びる水平面は、R の j 番目の列から伸びる垂直面と交差します。このベクトルに沿って、左と右の引数からの (i, k) (k, j) 要素のペアが出会って乗算され、その積が k に沿って合計され、結果が結果の i、j の位置に配置されます。

行列の乗算は直感的に次のようになります。

1. 2つの直交行列を立方体の内部に投影します。

2. 各交差点の値のペアを乗算して積グリッドを取得します。

3. 3 番目の直交次元に沿って合計し、結果行列を生成します。

方向の場合、ツールは立方体の内部に結果マトリックスを指す矢印を表示します。青い矢印の羽は左のパラメータから、赤い矢印の羽は右のパラメータから来ています。このツールは、各マトリックスの行軸を示す白いインジケーターラインも表示しますが、このスクリーンショットではこれらのラインは薄くなっています。

レイアウト制約は単純かつ直接的です。

左の引数と結果は、共有の高さ（i）次元に沿って連続している必要があります。
右辺の引数と結果は、共有幅（j）次元に沿って連続している必要がある。
左と右の引数は、共有される次元 (左の幅 / 右の高さ) に沿って連続している必要があり、これが行列乗算の深さ (k) 次元になります。

この幾何学的表現は、行列乗算のすべての標準的な分解を視覚化するための強固な基盤を提供し、次に説明するように、複雑な行列乗算の非自明な組み合わせを探索するための直感的な基礎を提供します。

2 ウォームアップ: アニメーション

より複雑な例に進む前に、ツールに対する直感を構築するために、このスタイルの視覚化がどのようなものかを見てみましょう。

2a ドット積

まず、古典的なアルゴリズムを見てみましょう。結果の各要素は、対応する左の行と右の列のドット積を計算することによって計算されます。このアニメーションからわかるように、乗算された値のベクトルは立方体全体をスイープし、毎回対応する場所に合計結果を送信します。

ここで、L には 1 (青) または -1 (赤) で埋められた行ブロックがあり、R には同様に埋められた列ブロックがあります。ここで k は 24 なので、結果のマトリックス (L @ R) の青の値は 24、赤の値は -24 になります。

2b 行列ベクトル積

行列ベクトル積に分解すると、行列の乗算は垂直面（左の引数と右の引数の各列の積）のように見え、立方体の内部を水平に掃引すると、列が結果にプロットされます。

例が単純なものであっても、分解の中間値を調べるのは興味深いことです。

たとえば、ランダムに初期化されたパラメータを使用する場合、中央の行列ベクトル積の顕著な垂直パターンに注目してください。これは、各中間値が左側のパラメータの列スケールのコピーであるという事実を反映しています。

2c ベクトル行列積

行列の乗算をベクトル行列の積に分解すると、立方体の内部を通過しながら結果の上に行をプロットする水平面のように見えます。

ランダムに初期化されたパラメータに切り替えると、行列ベクトル積に似たパターンが表示されます。ただし、今回は水平であり、各中間ベクトル行列積が右側のパラメータの行スケールのコピーであるという事実に対応しています。

行列の乗算が引数のランクと構造をどのように表すかを考えるとき、計算中に次の 2 つのモードが同時に発生することを想像すると役立ちます。

以下は、ベクトル行列積を使用して直感を構築する別の例です。単位行列が 45 度の角度の鏡のように機能し、その引数と結果の両方を反映する様子を示しています。

2d 外積の和

3 番目の平面分解は、ベクトルの外積を点ごとに合計することによって、k 軸に沿った行列乗算の結果を計算することです。ここでは、外積平面が立方体を「後ろから前へ」掃引し、結果に蓄積されるのがわかります。

ランダムに初期化された行列を使用してこの分解を行うと、すべてのランク 1 外積が結果に追加されるため、結果の値だけでなくランクの累積も確認できます。

これは、近似する行列が低ランク行列である場合に「低ランク因数分解」（つまり、深度次元で小さなパラメータを使用して行列の乗算を構築することにより行列を近似すること）が最も効果的に機能する理由も直感的に説明しています。これは後ほど説明する LoRA です。

3 ウォームアップ: 表現

この視覚化方法を行列乗算の分解に拡張するにはどうすればよいでしょうか?前の例では、行列 L と R の単一の行列乗算 L @ R を視覚化しましたが、L および/または R 自体が行列乗算だった場合はどうなるでしょうか。

このアプローチは複合式に対して適切に拡張できることがわかりました。重要なルールは単純です。サブ式の (サブ) 行列乗算は別のキューブであり、親行列乗算と同じレイアウト制約が適用されます。サブ行列乗算の結果の面は、共有結合した電子と同様に、親行列乗算の対応するパラメーター面でもあります。

これらの制約の範囲内で、サブ行列乗算のさまざまな側面を必要に応じて調整できます。ここではツールのデフォルトのスキームが使用され、交互に凸型と凹型の立方体を生成します。このレイアウトは実際にはうまく機能し、スペースの使用を最大化しながら閉塞を最小限に抑えます。 (ただし、レイアウトは完全にカスタマイズ可能です。詳細については、mm ツールページを参照してください。)

このセクションでは、機械学習モデルの主要な構成要素のいくつかを視覚化します。これにより、読者はこの視覚表現に慣れ、そこから新たな直感を得ることができます。

3a 左結合式

以下では、それぞれ独自の形状と特徴を持つ、(A @ B) @ C 形式の 2 つの表現を紹介します。 (注: mm は行列の乗算が左結合であるという規則に従うため、(A @ B) @ C は単に A @ B @ C と記述できます。)

まず、A@B@C には特徴的な FFN 形状が与えられ、"隠れた次元" は "入力" または "出力" 次元よりも広くなります。 (具体的には、この例では、B の幅が A または C の幅よりも大きいことを意味します。)

単一行列乗算の例と同様に、浮動矢印は結果行列を指し、青い矢印は左の引数から、赤い矢印は右の引数から来ます。

B の幅が A または C の幅よりも小さい場合、A@B@C の視覚化には、オートエンコーダの形状に似たボトルネックが発生します。

凹面モジュールと凸面モジュールを交互に配置するパターンは、任意の長さのチェーンに拡張することもできます。たとえば、次の多層ボトルネックです。

3b 右結合式

次に、右結合式 A @ (B @ C) を視覚化します。

左結合式の水平展開（いわばルート式の左引数から開始）と同様に、右結合式チェーンはルート式の右引数から開始して垂直に展開されます。

右結合形式で形成された MLP を目にすることがあります。この形式では、入力が右側の列になり、重みレイヤーが右から左に実行されます。上記の 2 層 FFN の例の行列を使用すると (適切な転置後)、次のようになります。C が入力、B が最初の層、A が 2 番目の層になります。

また、羽根の色 (左側は青、右側は赤) の他に、左の引数と右の引数を区別する 2 つ目の視覚的な手がかりは、それらの方向です。左の引数の行は結果の行と同一平面にあり、同じ軸 (i) に沿って積み重ねられています。たとえば、上記の (B @ C) では、両方のヒントから、B が左側のパラメーターであることが分かります。

3c バイナリ式

ビジュアライザーが有用であるためには、単純な教育例に使用されるだけでなく、より複雑な表現にも簡単に適用できる必要があります。実際の使用例では、重要な構造コンポーネントはバイナリ式、つまり左辺と右辺の両方にサブ式がある行列乗算です。

最も単純な式 (A @ B) @ (C @ D) を視覚化したものを以下に示します。

3D に関するちょっとしたメモ: パーティション分割と並列処理

このトピックの完全な説明はこの記事の範囲を超えていますが、その実用性については、後ほどアテンションヘッドのセクションで説明します。しかし、ここでウォーミングアップとして、2 つの簡単な例を見て、このスタイルの視覚化によって、単純な幾何学的分割を通じて、並列化された複合式についての推論が非常に直感的になる方法を確認しましょう。

最初の例では、典型的な「データ並列」パーティション分割を、上記の左結合のマルチレベルボトルネックの例に適用します。 i に沿って分割し、最初の左側のパラメータ (「バッチ」) とすべての中間結果 (「アクティベーション」) をセグメント化しますが、後続のパラメータ (「重み」) はセグメント化しません。このジオメトリにより、式のどの参加者がセグメント化され、どの参加者がそのまま残っているかが明らかになります。

2 番目の例は、明確な幾何学的サポートがなければ直感的に理解するのが困難です。これは、左の部分式を j 軸に沿って、右の部分式を i 軸に沿って、親式を k 軸に沿って分割することにより、バイナリ式を並列化する方法を示しています。

4. 注意力を高める

それでは、GPT-2 のアテンションヘッド、具体的には、HuggingFace 経由の OpenAI の重みを使用して、NanoGPT の 5 層、4 ヘッドの「gpt2」(小) 構成 (レイヤー = 12、ヘッド = 12、埋め込み = 768) を見てみましょう。入力アクティベーションは、256 個のトークンを含む OpenWebText トレーニングサンプルのフォワードパスから取得されます。

この特定のヘッドには特別な点はありません。主に、非常に一般的な注意パターンを計算し、アクティベーションが構造化されて興味深いテクスチャを示すモデルの中央に位置しているために選択されました。

4a 構造

この完全なアテンションヘッドは、入力から始まり、投影された出力で終わる単一の複合式として視覚化できます。 (注: 自己完結性を確保するために、Megatron-LM で説明されているように、出力投影は各ヘッドに対して実行されます。)

この計算には 6 つの行列乗算が含まれます。

 Q = input @ wQ // 1 K_t = wK_t @ input_t // 2 V = input @ wV // 3 attn = sdpa(Q @ K_t) // 4 head_out = attn @ V // 5 out = head_out @ wO // 6

ここで何が起こっているのか簡単に説明してください:

風車の羽根は行列乗算 1、2、3、6 です。最初のセットは Q、K、V への入力の内部投影であり、後者は attn@V から埋め込み次元への外部投影です。
中央には 2 つの行列乗算があります。最初の乗算では注目スコア (後ろの凸型立方体) を計算し、それを使用して値ベクトル (前にある凹型立方体) に基づいて出力トークンを取得します。因果関係とは、注目度スコアが下側の三角形を形成することを意味します。

ただし、ツールの構造と計算を通じて流れる実際の値の両方についてより詳しく理解するには、スクリーンショットや以下のビデオを見るだけでなく、自分でツールを詳しく調べることをお勧めします。

4b 計算して評価する

以下は、アテンションヘッドの計算プロセスのアニメーションです。具体的には、

 sdpa (input @ wQ @ K_t) @ V @ wO

(つまり、上記の行列乗算 1、4、5、および 6。K_t と V は事前に計算されています) は、ベクトル行列積の融合チェーンとして計算されます。シーケンス内の各項目は、入力からアテンションを経て出力まで 1 つのステップで処理されます。このアニメーションオプションについては、後ほど並列化のセクションで詳しく説明しますが、まずは計算された値から何がわかるかを見てみましょう。

興味深いことが沢山あります:

注意計算について説明する前に、低ランクの Q と K_t がいかに顕著であるかを見るのは興味深いことです。 Q@K_t ベクトル行列積アニメーションを拡大すると、さらに興味深いことがわかります。Q と K の多数のチャネル (埋め込み位置) がシーケンス全体でほぼ一定であるように見えます。これは、有用な注意信号が埋め込みの小さなサブセットのみによって駆動される可能性が高いことを意味します。この現象を理解し、活用することは、SysML ATOM 変圧器効率プロジェクトの一部です。
おそらく最もよく知られているのは、注意行列に現れる強力だが不完全な対角線です。これは、このモデル (および多くのトランスフォーマー) の多くのアテンションヘッドに表示される一般的なパターンです。これにより、ローカルアテンションが可能になります。つまり、出力トークンの位置の直前の小さな近傍にある値トークンによって、出力トークンのコンテンツパターンがほぼ決まります。
ただし、この近傍のサイズと、その中の個々のトークンの影響は大きく異なります。これは、注意グリッドの非対角フロストや、注意行列がシーケンスに沿って下降するときの attn[i]@V のベクトル行列積平面の変動パターンに見られます。
ただし、ローカル近傍だけが注目に値するものではないことに注意してください。注目グリッドの左端の列 (シーケンスの最初のトークンに対応) は、ゼロ以外の (ただし変動する) 値で完全に埋められています。つまり、各出力トークンは最初の値トークンによってある程度影響を受けます。
さらに、現在のトークンの近傍と最初のトークンの間の注目スコアの優位性には、不正確ではあるが識別可能な振動があります。この振動の周期は変化しますが、一般的には最初は短く、シーケンスを下に行くにつれて長くなります (同様に、因果関係を考えると、行あたりの候補注目トークンの数に関連します)。
(attn @ V) がどのように形成されるかを理解するためには、attention だけに焦点を当てるのではなく、V も同様に重要です。各出力項は、V ベクトル全体の加重平均です。注意が完全に対角である極端なケースでは、attn@V は V の正確なコピーになります。ここでは、よりテクスチャ化されたものが見られます。特定のトークンが注目行の連続するサブシーケンスで高いスコアを獲得する目に見える帯状の構造が、明らかに V に類似しているが、より太い対角線による垂直方向の遮蔽が多少あるマトリックスに重ねられています。 (補足: mm リファレンスガイドによると、長押しまたは Control キーを押しながらクリックすると、視覚化要素の実際の数値が表示されます。)
覚えておいてください、私たちは中間層 (レイヤー 5) にいるため、このアテンションヘッドへの入力は、生のトークン化されたテキストではなく、中間表現です。したがって、入力に表示されるパターンはそれ自体が明らかになります。特に、強い垂直線は、シーケンスの長い範囲にわたって値が均一に高い特定の埋め込み位置であり、時にはほぼ完全です。
しかし興味深いのは、入力シーケンスの最初のベクトルがユニークであり、これらの高振幅列のパターンを破るだけでなく、ほぼすべての位置で非典型的な値を持っていることです (補足: ここでは視覚化されていませんが、このパターンは複数のサンプル入力にわたって繰り返し表示されます)。
注: 最後の 2 つの箇条書きに関して、ここで視覚化されているのは単一のサンプル入力の計算であることを繰り返し強調しておく価値があります。実際には、各ヘッドには、かなり大規模な例セットにわたって一貫して（同一ではないものの）表現される特徴的なパターンがあることがわかりますが、アクティベーションを含む視覚化を見るときは、入力の完全な分布が、それが刺激するアイデアや直感に微妙な影響を与える可能性があることを念頭に置くことが重要です。

最後に、もう一度アニメーションを直接探索することをお勧めします。

4c アテンションヘッドには多くの興味深い違いがある

先に進む前に、モデルを研究するだけでその仕組みを詳細に理解できるという有用性を示すもう 1 つの例を示します。

これは GPT-2 のもう一つの注目点です。その動作はレベル 5 の 4 番目のヘッドとはまったく異なりますが、これはモデルのまったく異なる部分にあるため、予想どおりです。このヘッダーは最初のレイヤー（レイヤー 0、ヘッダー 2）にあります。

注目すべき点:

このヘッドの注意力は均等に分散されます。これには、比較的重み付けされていない V の平均 (または V の適切な因果接頭辞) を attn@V の各行に連結する効果があります。アニメーションに示すように、注目スコアの三角形を下に移動すると、 attn[i]@V ベクトル行列積は、単純に縮小され、徐々に明らかになる V のコピーではなく、小さな変動を示します。
attn@V は垂直方向の均一性が顕著で、埋め込まれた大きな列内では、シリーズ全体にわたって同じ値のパターンが維持されます。これらは、すべてのトークンが共有するプロパティと考えることができます。
補足: 一方で、注意力が非常に均等に分散される効果を考えると、attn@V にはある程度の一貫性があると予想されるかもしれません。しかし、各行はシーケンス全体ではなく、V の因果サブシーケンスで構成されています。シーケンスを下に移動すると徐々に変形するなど、より多くのバリエーションが生まれないのはなぜでしょうか。目視検査により、V はその長さに沿って均一ではないことが分かります。したがって、答えはその値の分布のより微妙な特性にあるに違いありません。
最後に、外部投影後、このヘッドの出力は垂直方向でさらに均一になります。
このアテンションヘッドによって伝えられる情報のほとんどは、シーケンス内のすべてのトークンによって共有されるプロパティで構成されているという強い印象を受けます。出力投影の重みの構成は、この直感を強化します。

全体的に、この注意ヘッドによって生成された極めて規則的で高度に構造化された情報は、もう少し…贅沢ではない計算手段によって取得できたのではないかと思わずにはいられません。これは確かに未踏の領域ではありませんが、計算信号を視覚化することの明確さと豊かさは、新しいアイデアを生み出すのにも、既存のアイデアについて推論するのにも非常に役立ちます。

4d 再導入: 自由な不変性

振り返ってみると、繰り返しになりますが、アテンションヘッドのような重要な複合操作を視覚化して直感的に理解できるのは、重要な代数的特性 (パラメーターの形状がどのように制限されるか、どの平行軸がどの操作と交差するかなど) に余分な考慮が必要ないためです。それらは、覚えておくべき追加のルールではなく、視覚化されたオブジェクトの幾何学的特性から直接導かれます。

たとえば、これらの注意ヘッドの視覚化では、次のことが明らかです。

Q と attn @ V は同じ長さを持ち、K と V は同じ長さを持ち、これらのペアの長さは互いに独立しています。
Q と K の幅は同じで、V と attn@V の幅も同じです。これらのペアの幅は互いに独立しています。

これらの構造は実際に構築されており、複合構造内の構造コンポーネントの配置場所と向きによって決まります。

この「自由な特性」の利点は、典型的なアーキテクチャのバリエーションを探索するときに特に役立ちます。わかりやすい例としては、一度に 1 つの自己回帰トークンをデコードする際の 1 行の高さの注意行列が挙げられます。

5 並行注意

上記の 5 番目のレイヤーの 4 番目のヘッドのアニメーションは、注目ヘッド内の 6 つの行列乗算のうち 4 つを視覚化します。

これらはベクトルと行列の積の融合チェーンとして視覚化され、幾何学的な直感を裏付けています。入力から出力までの左結合チェーン全体が共有 i 軸に沿って階層化されており、並列化できます。

5a 例: iに沿った分割

実際に計算を並列化するには、入力を i 軸に沿ってブロックに分割します。特定の軸を特定の数のチャンクに分割するように指定することで、ツールでこのパーティション分割を視覚化できます。これらの例では 8 を使用しますが、この数に特別な意味はありません。

とりわけ、この視覚化は、各並列計算が、これらの行列の分割されていない次元に沿って分割された行列に隣接しているため、完全な wQ (内部投影用)、K_t と V (注意用)、および wO (外部投影用) を必要とすることを明確に示しています。

5b 例: デュアルパーティショニング

複数の軸に沿った分割の例もここに示します。この目的のために、私たちはこの分野における最近のイノベーション、つまりFlash Attentionなどの研究結果に基づいたBlock Parallel Transformer (BPT)を視覚化することにしました。論文をご覧ください: https://arxiv.org/pdf/2305.19370.pdf

まず、BPT は上記のように i に沿って分割します。実際、このシーケンスの水平分割は、アテンションレイヤーのもう半分 (FFN) まで拡張されます。 (これの視覚化は後で示されます。)

このコンテキストの長さの問題を完全に解決するために、MHA に 2 番目のパーティション (アテンション計算自体のパーティション (つまり、Q@K_t の j 軸に沿ったパーティション)) が追加されます。これら 2 つのパーティションを組み合わせることで、注目がブロックのグリッドに分割されます。

この視覚化から、次のことがはっきりとわかります。

この二重分割により、アテンション計算における各シーケンス発生の長さを視覚的に分割できるようになるため、コンテキストの長さの問題が効果的に解決されます。
2 番目のパーティションの「範囲」: ジオメトリから、K と V の内部投影計算は、コアの二重行列乗算と一緒にパーティション分割できることは明らかです。
微妙な詳細に注意してください。ここでの視覚的なヒントは、後続の行列乗算 attn @ V も k に沿って並列化し、部分的な結果を分割 k スタイルで合計して、二重行列乗算全体を並列化できることです。しかし、sdpa() の行単位のソフトマックスでは、attn@V の対応する行を計算する前に各行のすべてのセグメントを正規化する必要があり、これにより、アテンション計算と最終的な行列乗算の間に行単位の余分なステップが追加されます。

6 注意層のサイズ

注意層の前半 (MHA) は、その二次複雑度のために計算要件が高いことはよく知られていますが、後半 (FFN) も、その隠し次元の幅 (通常はモデル埋め込み次元の幅の 4 倍) のために独自の要件があります。完全な注意層のバイオマスを視覚化すると、層の 2 つの半分が互いにどのように比較されるかについての直感を構築するのに役立ちます。

6a 完全な注意層の視覚化

以下は、最初の半分 (MHA) が後ろに、2 番目の半分 (FFN) が前に配置されている完全な注意レイヤーです。ここでも、矢印は計算の方向を指しています。

注記：

この視覚化では、単一のアテンションヘッドを描写する代わりに、スライスされていない Q/K/V 重みと、中央の二重行列乗算の周りの投影が表示されます。もちろん、これは MHA 操作全体を忠実に視覚化するものではありませんが、ここでの目標は、各半分で実行される計算の相対的な量ではなく、レイヤーの 2 つの半分における相対的なマトリックスサイズをより明確に把握することです。（また、ここでの重みは実際の重みではなくランダムな値を使用します。）

ここで使用される次元は、ブラウザで（比較的）スケーラブルに保つために縮小されていますが、比率は同じままです（NanoGPT の小さな構成から）：モデル埋め込み次元 = 192（元々は 768）、FFN 埋め込み次元 = 768（元々は 3072）、シーケンス長 = 256（元々は 1024）。ただし、シーケンス長はモデルに基本的な影響を与えません。 (視覚的には、シーケンスの長さの変化は入力ローブの幅の変化として現れ、その結果、注意中心のサイズと下流の垂直面の高さが変化します。)

6b BPTパーティションレイヤーの視覚化

Blockwise Parallel Transformer を簡単に確認すると、アテンションレイヤー全体のコンテキストにおける BPT の並列化の視覚化がここにあります (上記のようにヘッドは省略されています)。特に、i（シーケンスブロック）に沿ったパーティションが MHA と FFN のハーフエッジの両方にどのように拡張されるかに注意してください。

6c パーティション分割 FFN

この視覚化は、上で説明したものと直交する追加の分割を示唆しています。つまり、注意層の FFN 半分で二重行列乗算 (attn_out @ FFN_1) @ FFN_2 を分割し、最初に j に沿って attn_out @ FFN_1 を実行し、次に k に沿って FFN_2 で後続の行列乗算を実行します。このパーティショニングにより、2 つの FFN 重みレイヤーがスライスされ、計算に参加する各コンポーネントの容量要件が削減されますが、部分的な結果の最終的な合計が犠牲になります。

このパーティショニングアプローチを、パーティショニングされていないアテンションレイヤーに適用すると、次のようになります。

BPT 方式で分割されたレイヤーには以下が適用されます。

6dはトークンを1つずつデコードするプロセスを視覚化します

自己回帰トークン単位のデコードプロセスでは、クエリベクトルは単一のトークンで構成されます。このコンテキストで注意層がどのように見えるかを頭の中で想像してみると役に立ちます。つまり、重みの巨大なタイル状平面を通過する埋め込みの 1 行です。

このビューは、アクティベーションと比較した重みの巨大さを強調するだけでなく、K_t と V が 6 層 MLP の動的に生成された層と同様に機能するという概念も呼び起こします。ただし、MHA 自体の mux/demux 計算により、この対応は不正確になります。

7 ロラ

最近の LoRA 論文「LoRA: 大規模言語モデルの低ランク適応」では、微調整中に導入される重み δ が低ランクであるという考えに基づいた効率的な微調整手法について説明しています。論文によれば、これにより「事前トレーニング済みの重みを固定したまま、適応中に密層の変化のランク分解マトリックスを最適化することで、ニューラルネットワーク内のいくつかの密層を間接的にトレーニングできるようになります。」

7a 基本的な考え方

つまり、重要なステップは、行列自体ではなく重み行列の因子をトレーニングすることです。つまり、I x J 重みテンソルを、I x K テンソルと K x J テンソルの行列乗算で置き換えます。ここで、K は小さいままです。

K が十分に小さければ、サイズの面で大きなメリットが得られますが、トレードオフがあります。K を小さくすると、製品が表現できるランクも下がります。サイズの節約と結果への構造化効果を説明するために、128 x 4 の左側引数と 4 x 128 の右側引数のランダム行列乗算、つまり 128 x 128 行列のランク 4 分解の例を示します。 L@R の垂直パターンと水平パターンに注意してください。

7b LoRAを注意頭部に適用する

LoRA がこの分解アプローチを微調整プロセスに適用する方法は次のとおりです。