私たちが完全に理解している最も複雑なAIモデルとは何か？

驚異的な発展を続けるAI技術に対して、「なんだかよくわからない」という印象を持つ人は決して少なくはありませんが、AIをよく知る人でさえ同じ考えに至るようです。AIに関する教育的デジタルコンテンツを制作するWelch Labsが「現代のAIを理解している者は誰もいない」と冒頭で述べる衝撃的な動画を発信していました。

The most complex model we actually understand - YouTube

AIが生成するトークンは数百億もの計算によって得られるものであり、計算に用いられるパラメータはAIが一度に一つのトークンを予測するよう訓練し続けることで学習してきたものです。

しかし、何兆もの例を通じて次々と小さなテキスト片を予測することを繰り返す学習から何故本物の知性であるかのようなものが生まれるのか、という疑問に対する決定的な答えは誰も持っていません。AIモデルを構成する個々の人工ニューロンの仕組みは簡単ですが、ニューロン同士が接続して複雑化するにつれ理解が追い付かなくなります。どこで人間の理解を超越したのかを突き詰めると「単層Transformerを用いた剰余演算の学習」こそが、私たちが完全に理解している最も複雑なAIモデルであるといえます。

2021年にOpenAIの研究チームは小規模モデルに剰余演算を学習させていました。「x＋y」のような数学的演算をデータセットに変換できるというものです。

・x列・y行のマス目を用意する
・列には様々なx値、行には様々なy値を配置する
・行と列の交わるセルにはx値とy値の和をセットする
・セルの値には最大数が存在し、セットしようとした値が最大値を超えた場合は最大値で割った剰余をセットする

AIモデルはこの剰余演算に対して学習すべき興味を持ち、データの一部のみをテスト用に確保し残りを学習に費やすようになります。AIモデルの視点で考えると、0から4までの数字を表すには5つのトークンが必要であり、加えて等号を表すトークンが1つ必要になります。なおプラスのトークンが不要なのはこの学習では加算しか行わないからです。例えば「1＋2」という算術問題をモデルに入力する場合、問題の最初のトークン「1」をモデルに渡すには、1の位置のみをオンにし他のすべての位置をオフにします。2番目のトークン「2」は2番目の入力のみをオンにします。最後の等号は6番目の入力のみをオンにします。したがって、モデルから見た「1＋2」という算術問題は「1番目のみオン→2番目のみオン→6番目のみオン」という順序で入力されます。さらに、Transformerモデルは入力と同じ次元の出力を返すよう設定されるため、モデルの最終出力も6×3になります。

しかし、このモデルの場合に必要な出力は最終列だけです。「1＋2＝3」の例では最終列の3番目の出力がオンになります。

つまり、モデルが学習しているのは18個の値からなるパターンを6個の値からなる新たなパターンへマッピングする作業であるといえます。

同様にして様々な目標入力と出力パターン、例えば「1＋3＝4」や「2＋3＝0」のパターンなどが与えられます。これらの例を十分に観察することにより問題の根本的な構造を理解するというのが大規模言語モデルの正確な動作原理です。

OpenAIチームが剰余演算でモデルを訓練したところ、初期の結果はかなり期待外れなものでした。モデルは訓練データのパターンを素早く学習して全ての訓練例で正しい出力を生成できたものの、テストセットでは極めて低調な性能を示しました。モデルは剰余演算を実際には学習しておらず、単に訓練データを丸暗記しただけのように見えたのです。

しかしながら、とある研究者がモデル訓練させたまま休暇に出かけるミスをしたところ、休暇の間に膨大な学習ステップを経たモデルは汎化を達成し、訓練データとテストデータの両方で完璧な性能を示すようになりました。いったい何が起こったのでしょうか？

OpenAIチームは様々な算術演算やモデル構成において突発的汎化現象を再現することに成功し、2022年1月にこの現象を「Grokking(グロッキング)」と題した論文を発表しました。SF作家のR.A.ハインラインの小説「異星の客」に登場する人物が使用する火星語「grok」には「表面的な理解ではない内面的な理解」という意味がありますが、この言葉を転用して突発的汎化現象のことをグロッキングと命名したわけです。

Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets
https://arxiv.org/abs/2201.02177

論文ではモデルのパラメータと活性化を深く掘り下げ、非常に納得のいく秀逸な説明を導き出しています。論文では先程の簡略化モデルよりは規模が大きくなりますがNANDで113桁を使用しており、したがってモデルの入力ベクトルは0から112までの数字を示す113ポイントと等号を示す1ポイントを加え、長さは114となります。モデルに「1＋2」を計算させるにはこの114×3の行列を渡しています。行列には最初の列の1番目の位置に1、2番目の列の2番目の位置に1、最後の列の等号の位置に1が配置され、それ以外はすべて0で構成されています。そして入力の113×3行列を学習済み重み行列(埋め込み行列)で乗算することにより128×3の新しいベクトルを生成します。埋め込みベクトルはアテンションブロックを経て多層パーセプトロン演算ブロックへ渡され、多層パーセプトロンからの出力は長さ128となります。この出力を埋め込み行列と乗算し長さ114の最終ベクトルを計算します。モデルの回答はこの最終ベクトル内の最大値によって決定されます。したがってモデルが正常に動作している場合は最大出力値は正解である「1＋2＝3」に対応する3番目の位置で発生するはずです。

モデルを剰余演算で訓練するとOpenAIチームが観察したのと同じ成長パターンが見られます。モデルは約140ステップ後に訓練データを記憶し、1万ステップの訓練を経て汎化します。

モデルの中間出力(活性化)について詳しく見るため、多層パーセプトロンブロックの第2層にある合計512個のニューロン群の出力を詳細に観察します。

「0＋0」という問題をネットワークに入力すると、この層の最初のニューロンは出力値1.17、2番目のニューロンは0.6、といった具合に出力を返します。

では入力計算式を変更した際の値の変化を可視化しましょう。xの値を0に固定し、y値の範囲を探索します。「0＋0」「0＋1」「0＋2」……と進め、yの全113通りの値を網羅すると、一部のニューロン出力に正弦波のような形状が現れます。

さらに調査を深めるためにこれらのニューロンの全ペア間の相関を探るべく、各ニューロンペアの散布図を7×7グリッドで作成します。例えば2番目の散布図(1行目)では最初のニューロンの出力をy値、2番目のニューロンの出力をx値としてプロットしていきます。するとプロットは美しいループ形状を描きます。

各ニューロン出力ペアで同様のプロットを作成していくうち、モデルは明らかに何らかの構造を学習している構造を示すことが見て取れます。

ただ、この構造が学習に関連しているのかどうかは気になるところですので訓練を遡って確認してみます。

学習データを単純に記憶したモデル段階ではこれらの構造が完全に消失しているのがわかります。つまり初期モデルでは学習後に示された波形やループといった構造が全く見られませんので、学習の段階に関連している可能性があります。

モデルがグロッキング動作中に内部で生成する波形とループは、入力xとyのsin値や正負を計算し活用している可能性を示唆しています。活性化パターンに対して離散フーリエ変換を適用するとモデルが学習した波の周波数を計算できます。最初の波では最大周波数成分が8π/113を示し、3番目の波では最大周波数成分が6π/113を示しています。

これらの波をモデルの出力上にプロットすると見事に一致することがわかります。

モデルが示す波形が学習によるものかを識別するために「スパース線形プローブ」と呼ばれる手法を用います。埋め込みベクトルのさらに数箇所の値をサンプリングすると同様の半周期的な曲線が見られます。

8つの曲線を重み付き和で結ぶと周波数8π/113のcos曲線に極めて近い曲線が得られることが判明しました。この重み付き和が重要なのは、重み付き和の計算こそがアテンションブロックや多層パーセプトロンブロックの主要な機能だからです。

この段階で、x×8π/113のsin値に対しても同様のスパース線形プローブを計算できます。

最初の埋め込みベクトルは最初の入力xのみに依存し、2番目の埋め込みベクトルは2番目の入力yのみに依存しています。これらの入力はアテンションブロックで結合されますが、3つの入力を独立に処理するのに同じ埋め込み行列が使用されるため、2番目の埋め込みベクトルに対しても同じスパース線形プローブを適用できるのです。適用すると同じ美しいcos曲線とsin曲線が現れますが今度はyの関数としてです。つまりモデルは非常に早い段階で入力のsin値とcos値を計算することを学習しているのです。

今度はxとyの両方を変化させた場合にこれらの出力がどのように変化するかを探り、ネットワークがこれらの変数をどのように結びつけているのかを解明できるかどうか見てみます。単一ニューロンの出力をグラフ化し、yを0に固定してxの全可能な値を走査するとおなじみとなったcos曲線が得られます。

次にグラフに軸を追加し、yを変化させながらニューロンの出力をプロットします。

xとyの全値の組み合わせを探り、ニューロンの出力をサーフェスグラフの高さとしてグラフ化します。サーフェスグラフの色はニューロンの出力値に対応させています。これまで見てきた多くの出力と同様にこのグラフもほぼ波状になります。

ネットワークが学習したこの波構造を最もよく捉えるsin曲線とcos曲線の組み合わせを調べるため今回も離散フーリエ変換を適用しますが、今回はxとyの両方に対して行います。

上位周波数を抽出することでサーフェスグラフをいくつかの主要成分に分解できます。下図で分解したグラフのうち一番下の青で示したグラフの成分はcos(x)、下から2番目の黄色で示したグラフの成分はcos(y)です。そして下から3番目のオレンジ色で示したグラフの成分が最も強く表れておりかつ最も興味深いものです。なぜならcos(x)とcos(y)の積に等しいからです。

次に多層パーセプトロンをさらに1層深く掘り下げ、この層のニューロンの出力をxとyの関数としてプロットしてみましょう。やはり波のような形状が見られますが、波は先程のものより不規則で斜め向きとなっていることがわかります。

ニューロンの出力が最大化されている2つの波頭について、波の頂点に該当する入力値の組み合わせを見ていきます。

1番目の波の最初の頂点は「x＝0, y＝65」です。波峰に沿って左下に移動していくと、中間値が「x＝20, y＝45」「x＝40, y＝25」「x＝60, y＝5」となっており、「x＝65, y＝0」で終わります。入力ペアはいずれも合計値が65となります。つまりこのニューロンは「x＋y」が65の時に最大限に発火するのです。したがってニューロンは独自の方法で加算を学習した、より正確には合計が65となるすべての入力ペアに対して発火するということになります。

2番目の波の頂点は「x＝66, y＝112」で始まり、「x＝91, y＝87」といった値を経て「x＝112, y＝66」で終わり、入力ペアの値を合計するといずれも178となります。本モデルは除数113による剰余演算で訓練されていることを考慮すると、2番目の波の頂点についても演算結果が65となる入力ペアで構成されているといえます。

ではたった1層のニューロンで、cos(x)×cos(y)からどうやってxとyそのものを加算する処理へ移行しているのでしょうか？

先程、第1層ニューロンの出力で最も強い周波数成分がcos(x)×cos(y)であり、第2層ニューロンで最も強い周波数成分がsin(x)×sin(y)であることを確認しました。ここで以下の仮定を行います。

・cos(x)×cos(y)ニューロンに割り当てられた重みは1である
・sin(x)×sin(y)ニューロンに割り当てられた重みは-1である

この負の重みは2番目のグラフを上下反転させます。ここで重み付けされたグラフを加算すると、sinとcosが驚くほど適切に干渉し合って次の層のニューロンに見られる対角線対称性を生み出し、結果として入力の合計が65となる組み合わせでニューロンが発火するのです。

cos(x)×cos(y)－sin(x)×sin(y)は実は三角関数の公式「加法定理」であり、cos(x＋y)に等しくなります。この公式によりsinとcosの積の和をxとyの和に変換できるのですが、まさにこれがネットワークが計算をする上で必要となるものだったのです。

驚くべきことに、ネットワークはこの三角関数の公式を効果的に活用して剰余演算を解決することを学習したようです。

モデルの最後のunbed部ではさらに重み付き和を計算します。今回は多層パーセプトロンにおける最終層ニューロンの出力を見てみます。いくつかのニューロンの出力をグラフ化すると様々なシフトやスケールを持つ対角線対称性が確認できます。

以下は出力が7を返す結果のサーフェスグラフです。

65になるすべての数の組み合わせを検出した多層パーセプトロンニューロンと同様に、このグラフは7になるすべてのxとyの組み合わせで最大値に達します。この剰余演算問題を解決するためにネットワークは入力の数値の符号と近接性を推定し、それらの関数の積を計算します。その後、三角関数の公式を巧妙に利用して剰余演算を解くために必要な対角線対称性を創出し、これらの結果パターンを複数組み合わせることで最終的な答えを導き出していたのです。

では、モデルが剰余演算を解く仕組みを詳細に観察することでなぜ成長するのかを理解すべく、モデルが学習した様々な構造の進化を可視化しながらトレーニングプロセスをもう一度見てみましょう。学習開始から数百ステップでモデルはトレーニングデータに完全に適合しますが、まだ汎化の兆候や親和性の痕跡は見られません。

学習が進んでもモデルの性能は横ばい状態が続き、何も起きていないように見えます。モデルの学習過程ではトレーニング性能とテスト性能を可視化するのが一般的であり、両指標がこれほど長く横ばい状態であれば通常は学習が完了し安定した解に収束したと推測するものです。しかし内部を詳細に観察するとモデルが剰余演算問題を解くために必要な構造を組み立て始めていることが明らかになります。

Neel Nanda氏らは論文で「除外損失(excluded loss)」という巧妙な新指標を提案しています。これまで学習過程におけるモデルの精度をプロットしてきた代わりにモデルの交差エントロピー損失をプロットしますので、値が低い方が良好な成績となります。

モデルがいくつかの主要周波数で動作していることがわかったので、性能測定前に主要周波数に関する情報をモデルの最終出力から除去してみます。つまり、先程発見した8π/113周波数を除去し、学習過程におけるこの除外損失をプロットします。新しい指標は学習損失と共に急激に低下しますが、モデルがsin曲線とcos曲線による表現を構築するにつれて緩やかに上昇します。どうして除外損失が増加するのかというとモデルが主要周波数を利用する能力を奪ったからです。

着目すべきなのは、長い期間にわたる学習・テスト性能の横ばい状態の中で除外損失が徐々に上昇している点で、モデルが主要周波数のパターンを次第に活用していることを示しています。興味深いことに、Nanda氏らは成長がsin・cos構造の完成時ではなく、モデルが初期学習で依存した記憶例を除去する「クリーンアップ段階」の直後に発生することを示唆しています。このダイナミクスは極めて興味深く、本モデルが剰余演算で優れた性能を発揮するに至った理由を明快に説明しています。除外損失という指標はモデルの記憶依存から学習への緩やかな進展を明確に示し、グロッキング現象を実に巧みに説明しています。これだけの明瞭さは現代AIにおいては美しくも稀な例外であり、ブラックボックスだらけの世界における稀有な「透明な箱」であると言えます。