Loop Transformer の技術的考察と RDMA 設計の経験談

【随笔】Loop Transformer について

空港で飛行機を待っている間に、ふと思い付いたことを書いてみます。

微信公众号（WeChat Official Account）の利点の一つは、過去の投稿がそれぞれの時点での考え方を忠実に記録していることです。数年経ってから自分の技術的判断を検証する際に非常に役立ちます。例えば、2023 年や 2024 年の記事を読み返すと、私はすでに MoE（Mixture of Experts）について議論していました。当時はまだ主流の話題ではありませんでした。しかし、RDMA 転送プロトコルを設計して AlltoAll の問題を解決する必要に迫られていた頃も、CIPU 2.0 の流片前にすべて完了させることができました。最近でも社内では、「ある企業のトップクラスのネットワークカードで EP を有効化すると機能しなくなる」という話や、IB（InfiniBand）に関する噂話が飛び交っています。笑ってしまうほどです。自分が三年前の判断がいかに正しかったかを実感しました。

しかし、今日はその話ではありません。今日の本題は Loop Transformer です。実は GPT-5 の段階で「Universal Transformer を採用している」という噂がありましたが、最近では GPT-6 についても同様の噂が流れています。さらに、智譜（Zhipu AI）の決算報告書でも言及されています。

そこで、ここ一年ほどの間に執筆したいくつかの記事を引用しつつ、なぜ私が Loop Transformer に注目し、好意的に考えているのかを解説します。

昨年、『Transformer の進化について：UT, MoD, MoR...』という記事を書き、Universal Transformer、MoD（Mixture of Delays）、MoR（Mixture of Routers）といった論文の分析を行いました。Loop Transformer の考え方に納得したのには、いくつかの理由があります。

まず、コンピュータアーキテクチャの観点からです。Attention ブロックを CPU などの命令実行ユニットと見なし、FFN（Feed-Forward Network）をメモリとして捉える視点です。この立場では、MoE や Sparse Attention は、いわば分岐判断のような能力を持つことになりますが、それでもまだ完全なコンピュータとして必要な「ループ指令」の機能が欠けているように思えます。したがって、この視点に立てば、Loop Transformer のアルゴリズムを受け入れるのは自然なことだと言えます。

もう一つの視点は情報論です。すべてのトークンが同じだけの新しい情報を保持しているわけではありません。より多くの新情報を提供するトークンに対しては、より多くの計算リソースを投入すべきです。これが私が MoD や MoR といったアルゴリズムを支持する理由の一つでもあります。

さらに、TDA（トポロジカルデータアナリシス）の視点からも、Attention メカニズムには何らかの反復処理を通じて知能を高める必要があるという示唆が得られています。

最後に、Linear Attention も一種の Loop Transformer と見なせるのでしょうか？昨年、『未来の Attention アルゴリズムの選択：Full, Sparse, Linear?』という記事で詳しく分析しました。当時の結論は以下の通りです。

Sparse Attention のアプローチにおいて、KDA（Kernelized Dynamic Attention）のような「状態行列」を導入し、UT や MoR などの再帰構造を追加する。あるいは Google Titan の『Titans: Learning to Memorize at Test Time』のように、テスト時の学習能力を備えたメモリコンテキストを導入することで、Memory as Context (MAC) という概念を実現し、KDA と同様の「状態行列」の効果を達成できる。

このように、Loop Transformer は単なるアーキテクチャの変更ではなく、計算資源の効率的な配分や、情報処理の本質的な理解に基づいた重要な進化であると考えられます。

今年、RSI やテスト時トレーニング（Test-time training）に関する論文を改めて見直すと、ますます面白さを感じています。また、昨年上半年に書いた「Claude Mythos と Loop Transformer について」のアイデアも、より本質的な問い——中国国内の計算リソースで 10T パラメータ規模のモデルを推論するにはどうすればよいか——へと発展しました。

モデルをより大きなパラメータ規模へスケールさせる際、最も単純なアプローチは二つあります。一つは幅を広げる方法（例えば隠れ層次元数を増やす）、もう一つは深さを増す方法（例えばレイヤー数を増やす）です。しかし、両者にはそれぞれ欠点があります。どちらの手法も全体の活性化パラメータ数が増加し、特にモデルを深くする場合は TBT（Time-Bound Training）への影響が顕著になります。この状況では、よりスパースなモデルが必要となります。すなわち、膨大な事前学習知識を保存するための十分な容量と、後続のトレーニングで優れた汎化能力を発揮するための十分な能力を兼ね備えたモデルです。これはモデル自体の要請ですが、同時に計算リソースという制約も考慮する必要があります。単純に幅を広げたり深さを増したりするのは、KV Cache の使用量が倍増するという致命的な問題を引き起こします。これらは中国国内の計算リソースのニーズには適していません。

Loop Transformer を用いれば、例えばモデルを 40 レイヤーのまま維持しつつ、ループをもう一周追加するだけで済みます。この場合、KV Cache は最後のループ分のみを記録すればよく、結果として 80 レイヤーのモデルと比較して KV Cache の使用量が半分になると考えられます。これらのメリットは明白です。しかし、私がイメージしていた Loop Transformer とは少し異なるかもしれません。むしろ私は、Loop Attention を採用し、MoE/FFN レイヤーを半分だけ通過する形を期待していました。

昨日、ある企業の関係者と話した際、彼らは「将来の 10T モデルでは活性化パラメータの 90% が FFN に占められるため、LPU や FFN 専用プロセッサが必要になる」と述べていました。しかし、私はこの見解には強く反対です。私の直感的な答えは、Attention と MoE/FFN の間に何らかのバランスが存在すべきだというものです。

仮に 10T モデルを構築する場合、全体の活性化パラメータ数を 100B〜150B に抑えたいと考えます。その上で、メモリバウンド型の層である MoE や FFN をできるだけ少なく通過させることを目指します。直感的には、これらの MoE は非常にスパースであるべきです。例えば、1024 または 2048 の専門家（Expert）から 8 つだけを選択するような構造が考えられます。少ない topk 値は、推論時の EP（Expert-Parameter）通信オーバーヘッドを低減するのに有利ですが、トレーニング時にはこれらのパラメータが学習した知識をすべて記録し、デッドエキスパート（Dead Expert）の発生を最小限に抑える必要があります。

自然な選択として、topk ルーティングに入る前に、より多く正確な情報を提供してルーティングを支援する仕組みが必要です。そのためには、Hyper Connection のようなメカニズムを用いて topk ルーターへの入力幅を広げることが考えられます。また、Attention については、より高い「解像度」を実現するために計算リソースを投入する必要があります。

もちろん、「Expert にパラメータを積み重ねる以外にも、Engram のような拡張手法があるのではないか」と指摘する人もいます。最近の Engram に関する考察では、100B〜300B 程度の Flash モデルのパラメータ拡張には有効である可能性がありますが、2T〜10T パラメータ規模のモデルには必ずしも適さないと考えられます。特に KV 間でオーバーラップを持つ CSA（Contextual Self-Attention）や、Linear Attention に組み込まれた Convolution が既に n-gram の表現能力を持っている場合などです。ただし、O(1) でパラメータを参照し計算コストをかけないという特性は非常に有望であり、Engram などのアルゴリズムは、潜在空間（Latent Space）の表現能力についてより慎重に検討されるべきかもしれません。

これらの制約条件下では、モデルアーキテクチャの設計をアテンションブロックへとシフトせざるを得なくなります。以前数学専攻で計算手法を学んだ経験からか、私は反復アルゴリズムを好む傾向があります。そのため視線は「Loop Attention」のようなアプローチに向きました。

ループの本質は、アテンションヘッドを実質的に増加させることにあります。これにより、多次にわたるアテンション計算の出力を、トップ K の潜在空間（latent space）内のトークン表現に対してより高次元で十分に展開することが可能になります。その結果、よりスパースなエキスパートルーティングにも柔軟に適応できるようになるのです。

最近、Google の論文《recirculation》[1] を読んでいます。私が期待する働き方と合致しているかもしれません。ループ（looping）とリサーキュレーション（recirculation）の違いは、射の合成方向にあります。各層を状態空間上の自己写像と見なす場合、ループは同一対象上での写像の反復合成であり、依然として深さのカテゴリー内に留まります。一方、リサーキュレーションは時間並進関手を導入し、深さのカテゴリーを時空の二重カテゴリーへと埋め込みます。

ただし、この方式では KV Cache の使用量や処理がより複雑になる可能性があります。KV プレフィックスマッチングも同様に難易度が上がるでしょう。今後、改めてこの論文を詳しく解説する機会があれば、その点について詳しく述べたいと思います。

そもそも Recirculation や Loop が必要とされる根本的な理由は、アテンション計算において潜在空間の Q（クエリ）を書き換え、同時に K（キー）や V（バリュー）も最適化されるプロセスにあるからです。特にスパースな処理においては、各ループごとのトップ K の選択が、単純な切り捨てではなく、より多くのブロックを動的に選出する役割を果たします。

最近、非常に面白いアイデアが浮かびました。DSA や CSA に似た構造が存在すると仮定し、Recirculation/Loop を通じて出力された OV（出力・値）を、MoE のゲート入力となる部分空間として利用できないでしょうか？複数の出力から構成される高次元空間の中に、よりスパースな表現が存在する可能性があります。これにより、トップ K のエキスパート選択がさらに細粒度に行えるようになるかもしれません。

あるいは、「Linear Attention 自体にも再帰やループの属性があるのではないか？」という疑問が浮かぶかもしれません。私の考えとしては、Prefill と Decode 時の並列性を最大限に高めると同時に、KV Cache コンテキストの交換可能性と組み合わせ可能性を向上させることで、推論時のアテンション演算子の並列性を高めることに注力すべきだと考えています。

以上が、最近 Loop Transformer を振り返る中で私が抱いた、まだ未熟な考えの一部です。

もしかすると、これらの工夫は実用的でないかもしれません。モデルアーキテクチャの変更よりも、高品質なデータを直接用意する方がよほど効果的なのではないでしょうか？

参考文献

[1] recirculation: https://arxiv.org/abs/2608.17981

（微信アプリで開くにはこちら）