LLM関連の構造理解(論文要約) "Looped Transformers as Programmable Computers"
良書であればある程、引用文献や引用論文を更に芋づる式に掘り起こしたくなるのですが、数理系やデータサイエンス系は中々骨が折れるので、漸く今年の1月の研究論文に追いついてきました。AI×ビックデータみたいな話をいまだにされるのですが、大規模言語モデルからすると時代遅れ過ぎていて、改めて前提知識も踏まえた整理をするための要約(意図的な忘却の為の記録)。
"Looped Transformers as Programmable Computers"
論文概要
トランスフォーマーネットワークを一般的なコンピュータとして使用するためのフレームワークが提案されています。特定の重みをプログラムしてループ内に配置することで、入力シーケンスは命令とデータの読み書きのためのメモリとして機能します。
先に結論
トランスフォーマーネットワークが特定の重みとループ内で配置されることで、複雑なアルゴリズムとプログラムをシミュレートできることを示しています。
基本的な計算機、線形代数ライブラリ、バックプロパゲーションを用いたインコンテキスト学習アルゴリズムをエミュレートすることができる13層のトランスフォーマーを示しています。
セクション１：読み書き操作と条件分岐
トランスフォーマーがスクラッチパッド（一時的なメモリ領域）から特定の位置にデータを読み取り、書き込むことができることが証明されています。また、条件が真であればプログラムカウンターを特定の位置に設定する条件分岐命令も導入されています。
SUBLEQと一命令コンピュータ（OISC）
SUBLEQ命令は、一命令コンピュータ（OISC）を構築するための単一の命令です。この研究では、SUBLEQがチューリング完全であることが証明されています。
FLEQ: より柔軟な注意ベースのコンピュータ
FLEQは、SUBLEQを一般化した新しい概念で、より一般的な関数を生成できます。これにより、システムの柔軟性が高まり、非線形計算や線形代数計算、反復最適化アルゴリズムなども実装可能になります。
関数の統一テンプレート形式
非線形関数や基本的な線形代数演算をトランスフォーマーで実装する方法を示しています。これらのテクニックは、次のセクションで反復アルゴリズムを構築する際に重要です。
非線形関数のエンコーディング
注意機構内でさまざまな関数をエンコードするための主要な要素が紹介されています。具体的には、ソフトマックスをシグモイド関数として動作させ、クエリと値の重み行列に複数の係数を格納する方法が説明されています。
行列の転置と乗算
ソフトマックスを線形化することで、行列の転置と乗算が可能であることが示されています。この手法は、行列のベクトル化形式を作成してから、固定された列の順列でその転置のベクトル化形式を作成する方法で行列の転置を実装します。
注意力と全結合ネットワークの利点
全結合ネットワークも普遍的な関数近似器であるため、以前のセクションで提示された関数と全体的な辞書機能を実装することが可能です。しかし、簡単な関数（例えば、(x^2)）を計算するためには、ReLUネットワークに対数的な深さが必要であることが指摘されています。これに対して、注意力ベースのネットワークでは、本質的に2層で(x^2)を実装できることが示されています。
主なポイント
非線形関数や基本的な線形代数演算をトランスフォーマーで実装する方法が示されています。
注意機構内で非線形関数をエンコードする新しい手法が紹介されています。
行列の転置と乗算についても詳しく説明されています。
注意力ベースのネットワークが全結合ネットワークに比べて計算効率が高いことが強調されています。
セクション２：基本的な計算機能の実装
FLEQフレームワーク内で基本的な計算機能を実装する方法が示されています。具体的には、加算、減算、乗算、除算、平方根、パーセンテージなどの6つの基本的な関数をトランスフォーマーで実装する方法が説明されています。
関数ブロックの定義
加算や減算を行うためのトランスフォーマーに基づいた関数ブロックが存在することが示されています。これらの関数ブロックは、3層、1つのヘッド、次元数(O(1))で構成されています。
乗算と除算の実装
乗算と除算もトランスフォーマーに基づいた関数ブロックで実装できることが示されています。乗算は2層、1つのヘッド、次元数(O(d))で、除算は乗算と逆数を取る操作で実装できます。
平方根と逆数の近似
平方根と逆数は、シグモイド関数の和で近似できることが示されています。これにより、標準的なトランスフォーマーに基づいた関数ブロックでこれらの関数を実装できます。
その他の関数とアルゴリズム
基本的な計算機能以外にも、代数関数や三角関数など、より多くの操作を含むように拡張することが可能です。これらの関数をトランスフォーマーで実装する際には、目標とする関数(f(x))をシグモイドで近似する方法と、反復的な数値アルゴリズムを使用する方法があります。
主なポイント
基本的な計算機能をトランスフォーマーで実装する方法が詳しく説明されています。
加算、減算、乗算、除算などの基本的な関数は、トランスフォーマーに基づいた関数ブロックで実装できます。
平方根や逆数などの複雑な関数も、シグモイド関数の和で近似して実装できます。
さまざまな代数関数や三角関数もトランスフォーマーで実装可能であり、その方法が示されています。
セクション３：数値線形代数アルゴリズムの実装
数値線形代数のアルゴリズムをトランスフォーマーで実装する方法が示されています。具体的には、QR分解、Gauss-Seidel法、Arnoldi反復法、Lanczosアルゴリズムなどが実装可能であるとされています。
学習アルゴリズムのエミュレーション
この部分では、確率的勾配降下法（SGD）をエミュレートする方法が紹介されています。線形モデルから始め、2層のニューラルネットワークに対するバックプロパゲーションアルゴリズムの実装に至るまでが説明されています。
ループトランスフォーマーの効果
ループトランスフォーマーは、多くのモデルでのインコンテキスト学習を効果的に行い、高い精度を達成できるとされています。これは、以前の研究で一回の推論呼び出しに限定されていたインコンテキスト学習を拡張するものです。
複雑な反復プログラムの実装
確率的勾配降下法（SGD）などの複雑な反復プログラムを実装するには、ループ構造を持つトランスフォーマーか、プログラムの深さに応じてサイズが増加するトランスフォーマーが必要であるとされています。
一般的な損失関数とモデルに対するSGD
この研究は、線形回帰を超えた一般的な損失関数とモデルに対してトランスフォーマーがSGDを実装できる最初の作業であるとされています。
主なポイント
数値線形代数の基本的なアルゴリズムをトランスフォーマーで実装する方法が示されています。
確率的勾配降下法（SGD）をエミュレートする方法が詳細に説明されています。
ループトランスフォーマーは、多くのモデルで高い精度を達成できるとされています。
複雑な反復プログラムもトランスフォーマーで実装可能であるとされています。
一般的な損失関数とモデルに対してもSGDを実装できるとされています。
セクション４：SGD（確率的勾配降下法）のシミュレーション
確率的勾配降下法（SGD）をシミュレートするプログラムが紹介されています。このプログラムは、トランスフォーマーの一連の命令を使用して、SGDの各ステップをエミュレートします。具体的には、重みとバイアスの更新、ポインタの操作、ループカウンタのインクリメントなどが行われます。
ネットワークの深さと計算コスト
この研究では、一つの隠れ層を持つニューラルネットワークに対するバックプロパゲーションをエミュレートするアルゴリズムが設計されています。しかし、この構造は任意の深さのネットワークに一般化できるとされています。ただし、ネットワークが深くなるほど、トレーニングの計算コストも増加すると指摘されています。
結論と未解決の問題
トランスフォーマーがプログラム可能な計算ユニットとして使用できることが示されています。特に、注意機構の多様性と、単一のループで複雑な反復アルゴリズムや一般的なプログラムを模倣できるモデルを作成できることが強調されています。さらに、自然言語でこれらのモデルを制御・プログラムする可能性が開かれています。
主なポイント
SGDをエミュレートするためのトランスフォーマーに基づいたプログラムが紹介されています。
ネットワークの深さが増えると、計算コストも増加する可能性があるとされています。
トランスフォーマーがプログラム可能な計算ユニットとしての潜在能力があり、注意機構の多様性が強調されています。
自然言語でトランスフォーマーを制御・プログラムする新しい可能性が開かれています。
セクション５-６証明と補足
主に数学的な証明と補足(後述)。具体的には、ポインタの加算、非線形関数のシグモイド和に関する補足があります。
ポインタの加算について
補足では、1つの隠れ層を持つフィードフォワードReLUネットワークが、2つの非負整数を二進数ベクトルで表現し、その和を二進数ベクトルで出力できることが示されています。
非線形関数とシグモイド
非線形関数をシグモイドの和で近似できることが示されています。具体的には、3層、mヘッド、次元数(O(d))を持つトランスフォーマーによって、この近似が可能であるとされています。
主なポイント
このセクションは主に数学的な証明と補足に焦点を当てています。
ポインタの加算や非線形関数のシグモイド和に関する具体的な証明が提供されています。
これらの証明は、トランスフォーマーが計算能力を持つことをさらに裏付けるものとなっています。
行列の転置
行列の転置に関する証明では、トランスフォーマーが4層、1つのヘッド、次元数(r = 2d + 2 \log d = O(d))を持つ関数ブロックで、行列(A)を転置した行列(A^T)を出力できることが示されています。
行列の乗算
行列の乗算に関する証明では、トランスフォーマーが2層、1つのヘッド、次元数(r = O(d))を持つ関数ブロックで、行列(A)と(B)の乗算(A^T B^T)を出力できることが示されています。
エラー分析
入力行列(X)の各要素がある定数(G)で制限されている場合の読み取り/書き込み操作のエラーについて説明されています。
主なポイント
行列の転置や行列の乗算に関するトランスフォーマーの応用が詳細に説明されています。
これらの証明は、トランスフォーマーが高度な計算能力を持つことをさらに裏付けています。
エラー分析部分では、入力行列の要素が制限されている場合のエラーについても考慮されています。
エラー分析と関数近似
エラー分析と関数近似についても詳しく説明されています。特に、行列乗算におけるエラーが計算され、そのエラーが前述の証明で使用された定数に依存することが示されています。
累積エラー
T回の操作後の累積エラーについても考慮されています。具体的には、各反復での入力が理想的な入力に近い場合、その後の反復でもエラーが制限されることが示されています。
Turing CompleteなSUBLEQ命令
この部分では、研究者たちが使用しているSUBLEQ命令がTuring Completeであることが示されています。これは、Minskyマシンと呼ばれるTuring Completeなマシンを用いて証明されています。
単一命令セット
最後に、単一の命令セットが紹介されています。この命令セットは、メモリの特定の位置に格納されたデータを操作するもので、スカラー、ベクトル、または行列としてデータを扱うことができます。
主なポイント
エラー分析と関数近似について詳細な説明があります。
累積エラーについても考慮され、その制限が説明されています。
SUBLEQ命令がTuring Completeであることが証明されています。
単一の命令セットが紹介され、その命令セットがどのようにデータを操作するかが説明されています。
その他：詳細なステップと計算
ループトランスフォーマーがプログラム可能なコンピュータとして動作するための詳細なステップと計算が説明されています。具体的な操作、行列の操作、そしてReLU層の使用について詳しく説明されています。
ステップ1：命令のコピー
プログラムカウンタ`pt`が指す`t-th`命令をスクラッチパッドの`Current Instruction`ブロックにコピーする必要があります。このステップでは、行列`X`の特定の行が考慮されます。
ステップ2：データのコピー
`pa`、`pb`、`pc`それぞれに対して一つのヘッドを使用し、データをスクラッチパッドメモリにコピーします。
ステップ3：データの移動
特定の関数ブロックにデータを移動するためのfeedforward ReLU層が使用されます。
ステップ4：関数の適用
各関数は独自の注意ヘッドを持ち、その結果はそれぞれの行ブロックに書き戻されます。
ステップ5：データの戻し
関数ブロックからスクラッチパッドメモリにデータを戻すためのfeedforward ReLU層が再び使用されます。
ステップ6：エンコーディングとメモリ
このステップでは、エンコーディング行ブロック、メモリストレージ行ブロック、および入力のその他の行に焦点を当てます。
ステップ7：セクション4.3と同一
このステップは、以前のセクション4.3と同一の操作を行います。
主なポイント
ループトランスフォーマーがプログラム可能なコンピュータとして動作するための詳細なステップと計算が説明されています。
行列操作、ReLU層の使用、データのコピーと移動について具体的な説明があります。
各ステップで行われる操作が詳細に説明されています。
行列乗算における誤差
誤差( A )は
$$
(e^{cx_{ij}} )
$$
などの項によって定義されています。この誤差は、定数( C_1 )と( C_2 )によって制御されることが示され、特定の条件が満たされた場合、合計誤差は( \epsilon )より小さくなるとされています。
関数近似における誤差
Lemma5における誤差は、定理3の直接的な結果。この誤差は、使用されるヘッドの数( m )に対して
$$
(\frac{1}{\sqrt{m}})
$$
に比例する。
( T )回の操作後の誤差の蓄積
( t )回目の反復での入力( X_t )は、誤差項
$$
(\epsilon_t M_t)
$$
によって摂動されます。論文では、次の反復( t+1 )での誤差
$$
(\epsilon_{t+1})
$$
が制限されることが示されています。
行列( A )との行列乗算によって、新しい誤差項
$$
(\epsilon_0)
$$
が生じるとされています。
トランスフォーマーをプログラム可能なコンピュータとして使用するためのアルゴリズムに関する詳細な説明。特に、スクラッチパッドメモリ（一種の一時的なデータストレージ）とその「現在の命令」ブロックに命令をコピーするプロセスに焦点が当たっている。
D.1 ステップ1
このステップでは、プログラムカウンタ ( pt ) によって指し示される ( t )-th 命令をスクラッチパッドの「現在の命令」ブロックにコピーする必要があります。
命令 ( ct ) は、複数のパラメータ
$$
(pat, pbt, pct, pflagt, ppt, pmt, dh, dw )
$$
マスクは
$$
(b^{(1)}{\text{mask}}, b^{(2)}{\text{mask}}, b^{(3)}_{\text{mask}})
$$
マトリックス ( X ) と操作
( X ) マトリックスは、この操作中に使用または変更される行の関連するサブセットのみを考慮。実行の開始時には、「現在の命令」行ブロックは空であり、入力 ( X ) は特定の形状を持っています。
注意ヘッドと ( K, Q, V ) マトリックス
特定の ( K, Q, V ) マトリックスを使用すると、( X ) マトリックスが更新され、新しい ( ct ) 値が生成されます。
ReLUによるマスク生成
$$
(b^{(1)}{\text{mask}}, b^{(2)}{\text{mask}}, b^{(3)}_{\text{mask}})
$$
ReLU関数を使用して生成されます。
補足計算
定数
$$
(\epsilon, \delta \in [0, 1])
$$
が与えられたとき、1つの隠れ層としきい値活性化関数を持つニューラルネットワーク ( f ) が存在し、隠れ層に ( d ) 個の活性化関数があると仮定する。このとき、すべての
$$
(x \in [-C, -\delta] \cup [\delta, C] )
$$
に対して、
$$
[|f(x) - \frac{1}{x} | \leq \epsilon]
$$
が成立する。ただし、
$$
(d = \Omega\left(\frac{\log(1/(\epsilon\delta))}{\epsilon\delta} + \log C\right))
$$
証明の概要
区間 ( [\delta, C] ) を特定のサブ区間に分割する。
各サブ区間
$$
([a_i, a_i(1 + \epsilon a_i)))
$$
に対して、
$$
( \frac{1}{x})
$$
を近似する活性化関数を用いる。このようなサブ区間の数を計算し、それが隠れ層で必要な活性化関数の数になる。
主なポイント
( f(x) ) と ( \frac{1}{x} ) の差が ( \epsilon ) 以下になるようなニューラルネットワークが存在する。
この近似の精度は、隠れ層の活性化関数の数 ( d ) に依存する。
( d ) の下限は下記。
$$
(\Omega\left(\frac{\log(1/(\epsilon\delta))}{\epsilon\delta} + \log C\right))
$$
特定の条件下で
$$
(\frac{1}{x})
$$
を近似するニューラルネットワークが存在することを示しています。そして、その近似の精度は隠れ層の活性化関数の数に依存するとされています。
この記事が気に入ったらサポートをしてみませんか？