ユニバーサル・トランスフォーマーにはメモリが必要：適応的再帰推論における深さ状態のトレードオフ

概要: 本研究では、数独エクストリーム（組合せ推論ベンチマーク）上で、適応計算時間（ACT）を備えた単一ブロックのユニバーサル・トランスフォーマー（UT）における計算用のスクラッチパッドとして学習済みメモリトークンを調べます。メモリトークンは経験的に必要であることが分かりました。テストしたすべての構成において――3つのシード、複数のトークン数、2つの初期化方式、ACTおよび固定深さでの処理――メモリトークンなしの構成では、非自明な性能を達成できませんでした。最適なトークン数は鋭い下限しきい値を示します（T=0は常に失敗、T=4は境界、T=8は81セルのパズルで確実に成功）その後、安定したプラトー（T=8-32、完全一致 57.4% +/- 0.7%）が続き、T=64では注意の希釈（attention dilution）によって崩壊します。
実験の過程で、70%以上の学習実行が失敗する原因となるルータ初期化の罠を特定しました。デフォルトのゼロ・バイアス初期化（p ~ 0.5）と、Gravesが推奨する正のバイアス（p ~ 0.73）のどちらも、初期化直後に約2ステップでトークンが停止し、浅い平衡状態（halt ~ 5-7）に落ち着くため、モデルはそこから抜け出せません。バイアスを反転して-3とする（「deep start」、p ~ 0.05）ことで、この失敗モードは解消されます。この罠が、アーキテクチャ選択のアーティファクトではなく、ACT初期化に固有のものであることを、アブレーションにより確認しました。
信頼できる学習が確立できたことで、(1) ACTは固定深さ処理よりも一貫した結果を与えること（3つのシードにわたって56.9% +/- 0.7% vs 53.4% +/- 9.3%）；(2) lambdaウォームアップを伴うACTは、34%少ないponderステップで同等の精度（57.0% +/- 1.1%）を達成すること；(3) 注意ヘッドが再帰的深さにわたって、メモリリーダー、制約伝播器、統合器へと専門化すること、を示します。コードは https://github.com/che-shr-cat/utm-jax で利用可能です。