継続的思考マシン
(pub.sakana.ai)- Continuous Thought Machine(CTM) は、ニューラルネットワーク計算に 時間的な神経ダイナミクス を明示的に組み込み、静的な feed-forward 処理だけでは扱いにくい「思考の展開」をモデル内部で実現しようとする新しいアーキテクチャである
- 構造の中心には、データ順序と切り離された internal ticks、各ニューロンごとの MLP が最近の pre-activation 履歴を処理する neuron-level models、ニューロン対の時間的同期を表現として使う仕組みがある
- ImageNet、2D 迷路、parity、Q&A MNIST、CIFAR-10/100、実数ソート、強化学習の実験で、中核構造を維持したまま入力・出力モジュールだけを変えて適用できる可能性を示した
- 迷路実験では、位置埋め込みなしで L/R/U/D/W の経路を直接予測し、39×39 迷路と長さ 100 までで学習したモデルが、99×99 迷路と約 6 倍長い経路にも一般化する例を示した
- CTM は同期表現を通じて、記憶の想起、適応的計算、解釈可能な注意の移動、内部ワールドモデルの形成といった振る舞いを示したが、実際のニューロンを文字通り模倣するモデルではない
CTM が狙う問題
- 既存のニューラルネットワークは、生物学的な脳で見られる 時間的な神経ダイナミクス を意図的に単純化し、大規模ディープラーニングに適した静的な活性値中心の処理を行ってきた
- 脳には spike-timing-dependent plasticity(STDP)、神経振動、spike timing や synchrony に基づく時間コーディングがあるが、現代のニューラルネットワークは主に単純性と計算効率を優先している
- 人間の認知の柔軟性と汎用性に比べて、現在の AI には不足している部分があり、その一部は 時間処理 に関係している可能性がある
- CTM の貢献は 3 つの要素に整理できる
- 分離された内部次元: 人工神経システムにおいて思考が展開できる時間軸
- neuron-level models(NLMs): 各ニューロンが固定 ReLU のような静的関数ではなく、入力信号の履歴を処理して活性化される
- neural synchronization: 観測と予測に直接使われる潜在表現
推論モデルと recurrence
- 既存モデルを大型化する方法は大きな進展をもたらしたが、計算コストとデータ要求が増大しており、長期的な持続可能性には疑問が残る
- 系列データでは RNN 系が長く使われてきたが、Transformer ベースの手法に多く置き換えられ、近年は recurrence がモデル複雑性を拡張する経路として再び注目されている
- テキスト生成の reasoning models は、中間生成によってテスト時の計算量を増やす recurrence の形を用いている
- CTM は recurrence 自体よりも、recurrence が可能にする 正確なタイミングと神経活動の相互作用 を中核とみなす
- 既存アプローチとの違いは 3 点ある
- データモダリティと分離された内部次元で逐次的な思考が可能
- 各ニューロンの private model が精密な神経タイミングを扱う
- 神経同期を課題解決の表現として直接利用する
CTM の 3 つの中核メカニズム
- CTM は、データを処理している間に内部で神経活動を複数の tick にわたって展開する構造である
- 1 つの internal tick では、最近の pre-activation 履歴が収集され、NLMs がそれを処理して post-activation を生成する
- 時間に沿った post-activation 履歴はニューロン間の同期計算に使われ、その結果が Synchronization Representation になる
- 技術的な詳細は Technical Report にあり、GitHub repository も公開されている
-
Internal ticks: 思考が展開する内部次元
- CTM は
t ∈ {1, …, T}という連続的な内部次元を導入する - RNN や Transformer のように単語やフレームなどデータ自体の順序に従って処理するのではなく、自ら生成した internal ticks に従って動作する
- この内部展開により、画像や迷路のような静的または非系列データでも、表現を繰り返し作り直して洗練できる
- ページ上部のインタラクティブな迷路デモは 75 ticks を使用している
- CTM は
-
Recurrent weights と neuron-level models
- CTM の synapse model は U-NET 形状の recurrent MLP であり、各 tick で pre-activation を生成する
- 直近
M個の pre-activation が各ニューロンの入力履歴として集められる - 各ニューロン
dは固有パラメータθdを持つ private MLP で自分の pre-activation 履歴を処理し、post-activation を生成する - すべてのニューロンの post-activation は attention output と結合され、次の tick の recurrent 計算に入る
-
Synchronization representation
- CTM は、ある時点のニューロン状態のスナップショットではなく、時間に沿った ニューロン活動ダイナミクス に従って入力を取り込み出力を生成するよう設計されている
- post-activation 履歴
Ztの内積St = Zt · (Zt)^Tによってニューロン間の synchronization matrix を計算する - この行列は
O(D²)に増大するため、実運用では(i, j)のニューロン対を部分サンプリングしてSoutとSaction表現を作る Soutは出力空間に射影されて logits などの予測を作り、Sactionは attention query のように観測のための行動に使われる- モデル幅
Dが大きくなるほど、可能な同期表現はD × (D+1) / 2規模まで増える
-
データ入力方式
- データは各 internal tick で、現在の synchronization に基づく attention によって観測される
- ほとんどの実験では標準的な cross attention を使用した
- FeatureExtractor がデータから key と value に使う局所特徴を作り、synchronization から射影された query がそれを参照する
- attention output は次の recurrence サイクルで post-activation とともに使われる
学習損失: internal ticks 全体を最適化
- CTM は各 internal tick ごとに出力を生成する
- 各 tick の予測に対して cross-entropy などの標準損失を計算し、certainty は
1 - normalized entropyで計算する - 最終損失は 2 つの tick を動的に選んで平均する
- 損失が最小の tick
t1 = argmin(L) - 確信度が最大の tick
t2 = argmax(C)
- 損失が最小の tick
- この方式は最後の step 1 つだけに依存しないため、複数の internal ticks で意味のある計算が起きるようにする
- この損失設計は curriculum 効果を自然に生み、問題の難易度に応じて計算量を調整できるようにする
ImageNet 実験
- ImageNet 実験の目的は新たな state-of-the-art を達成することではなく、CTM が データと相互作用する方法 を示すことにある
- CTM は画像を観測しながら予測を積み上げ、その過程で neural synchronization を直接表現として利用する
- internal steps は途中で打ち切ることができ、adaptive compute が可能である
- ある時点以降の精度向上は小さいが、追加の利得は存在する
- デモでは 16 個の attention heads の attention weights、クラス予測、時間に沿った certainty が表示される
- 神経活動は UMAP 射影で 2D 可視化され、各ニューロンは点で表され、サイズは絶対値、色は値の符号と大きさを示す
2D 迷路実験
- 迷路解きは、モデルが出発点から到着点までの経路を段階的に構成しなければならない課題として設定された
- CTM は経路を画像ではなく L/R/U/D/W のステップ列として直接予測するよう学習された
- ページ上部には小さなバージョンのインタラクティブデモがあり、より大きなモデルのデモも提示されている
- デモでは 75 internal ticks のあいだに経路が作られていく過程を示し、壁を通過する予測は有効経路表示から除外される
- 16 個の attention heads の重みと平均 attention があわせて表示され、モデルがどこに集中しているかを確認できる
-
一般化とワールドモデル
- 長さ 100 までの経路を 39×39 迷路で解くよう学習した CTM を、より大きい 99×99 迷路に適用した
- 例の迷路の全経路は学習条件よりおよそ 6 倍長い
- CTM は位置埋め込みを使わず、経路をクラス文字列として直接予測しなければならなかった
- 位置埋め込みがないため、CTM はデータに問い合わせて迷路を探索するための internal world model を作る必要がある
- 研究チームは、ゲームや動画のようなより複雑な環境で、明示的な位置エンコーディングなしに CTM がどう経路探索するのかを見たいとしている
Parity 実験
- parity task は、2 値系列の累積 parity を 64 個の位置それぞれで予測する課題として構成された
- 長さ 64 の binary vector 全体が一度に与えられるため、単純な逐次入力よりも難しい設定である
- CTM は internal ticks 数を変えて学習され、パラメータ数を揃えた LSTM と比較された
- 75 internal ticks を超える CTM はこの課題を安定して解くことができ、一部の実行では 100% 精度に到達した
- LSTM は 10 internal ticks を超えると学習が難しくなり、内部思考次元を展開するのに適していない結果を示した
- デモでは、一部の attention heads が 後ろから前へ データをなぞり、最初の attention head が negative parity の位置にだけ attend するなど、解釈可能な振る舞いが現れた
- 2 つの CTM の例は異なる戦略を学習した
- 1 つはデータを逆順に attend した後、累積 parity を一度に予測する
- もう 1 つは順方向に attend しながら parity を段階的に予測する
- どちらも完全な精度に到達した
Q&A MNIST 実験
- Q&A MNIST は、CTM の記憶と想起能力を評価するための課題である
- モデルはまず MNIST digit の系列を見て、その後にどの digit を想起し、どの modular operation を適用するかを指定する index と operator embedding を受け取る
- すべての digit と index/operator embedding が提示された後、zero-tensor flag が最終回答の生成を指示する
- 実験では、CTM の memory length は MNIST digit が neuron-level models の activation history window の外に来るよう設定された
- そのため CTM は、後で digit を想起するには活性化を組織して情報を保持しなければならない
-
結果と一般化
- LSTM は input あたり internal tick が 1 つのとき CTM より高性能だが、internal ticks が増えるとより不安定になる
- CTM は internal ticks が増えるほど性能が向上し、最も難しい in-distribution task で 95% 以上の精度 を達成した
- CTM はかなり前の timestep で見た digit 値を想起でき、これはニューロンの組織化と同期から生じた結果と解釈される
- 一般化実験では、学習時より多くの digits または index-operator embeddings が入力されたときの精度を測定した
- CTM と LSTM baseline はどちらも operation 数の増加には一般化できた
- 実証結果では、新しい index embedding が提示されるたびにモデルが指定された operation 結果を計算して保存するため、final answer flag を待たずに処理を継続できた
- CTM は internal ticks が多いほど性能が向上し、LSTM は逆の傾向を示した
追加実験
-
CIFAR-10: human、feed-forward、LSTM との比較
- CIFAR-10 実験は、CTM を human performance、feed-forward baseline、LSTM baseline と比較するための設定である
- 差異を明確にするため、制限された backbone を使用した
- human label データセットとして CIFAR-10D と CIFAR-10H を使用した
- CIFAR-10D は難易度レベルの補正に関連する
- CIFAR-10H は人間の不確実性を定量化するために用いられたデータセットである
- CIFAR-10D は こちら、CIFAR-10H は こちら にある
- calibration の計算には CIFAR-10H の確率が使われ、CTM は人間と比べても最良の calibration を示した
- CTM の神経活動は豊かで多様かつ複雑なダイナミクスを示し、periodic driving function がなくても periodic behavior が現れた
- CTM と LSTM の神経活動の違いは、neuron-level models と synchronization representation が計算メカニズムとして神経ダイナミクスを可能にしている証拠となる
-
CIFAR-100 ablation
- CIFAR-100 実験では、ニューロン数、すなわちモデル幅を変えつつ、他の条件と学習時間を固定した
- より広いネットワークでは、より長い学習時間や異なる hyper-parameters が必要な可能性があり、一部で精度低下が見られた
- neuron-level models がどれほど固有になるかを見るため、ニューロン間ダイナミクスの cosine similarity を測定した
- モデル幅が広がるほどニューロン間の多様性が減るのではなく、むしろ増える現象が観察された
- internal ticks 数と予測の関係も分析された
- 25、50、100 internal ticks 設定で、CTM が最も高い確信を持つ step の分布を確認した
- 各設定で 2 つの集中領域が現れ、CTM がデータに応じて別個の内部プロセスをたどっていると解釈された
-
実数ソート
- CTM は
N(0, I30)から生成された 30 個の実数をソートするよう学習された - 目的は、制御された環境で CTM がいつより多く、あるいは少なく計算するかを確認し、CTC loss によって系列出力を学習できるかを検証することである
- この CTM は長さ 30 の実数リストを約 80% の確率でソートできた
- CTM は
-
強化学習
- CTM は continuous thought dimension によって非系列データを処理するだけでなく、外部環境と相互作用する課題にも適用された
- proximal policy optimization により、navigation task と partially observable CartPole、および Acrobot の変種を学習した
- この設定では CTM は観測を受け取り、固定数の internal thought steps で処理した後、次の action を出力する
- activation history は環境 step 間で連続的に維持されるため、過去の environment step の activation が現在の意思決定に影響しうる
- その結果、CTM は LSTM baseline と同等の性能を示し、continuous environment でも学習可能であることを示した
結論と限界
- CTM は pointwise activation function を private neuron-level models に置き換えることで、より豊かなニューロンダイナミクスを生み、activation vector ではなく neural synchronization を新しい表現として利用する
- このアプローチにより、画像分類における時間に沿った表現構築、位置埋め込みなしの迷路 attention と内部地図形成、adaptive computation、activation history の外にある記憶の保存と想起が可能になる
- CTM の中核アーキテクチャは多様な課題で概ね維持され、主に入力・出力モジュールの調整だけが必要だった
- 迷路探索のような複雑なシナリオでも CTM は少ないチューニングで動作した一方、LSTM はかなり調整しても苦戦した
- CTM は生物学的ニューロンを厳密に文字通り模倣しようとするモデルではない
- 実際のニューロンは CTM のように activation history にアクセスしない可能性がある
- それでも traveling waves のような emergent phenomena が現れた
- CTM は生物学から概念を借りつつ、実用性と生物学的インスピレーションを折衷する手法であり、現在の AI に欠けている能力を開く研究方向になりうる
まだコメントはありません。