MACRO、凍結したTransformer層を並べ替え、6モデルの平均精度を5.0ポイント向上

新たな研究で提案されたMACROは、Transformerの固定された逐次的な層ごとの順伝播を、探索可能なプログラムへと書き換える。モデルは特定の層をスキップしたり、以前の層に戻って再実行したり、過去の隠れ状態をresidual streamに加え直したりした後、元のネットワーク後段へ再接続できる。プロセス全体を通じてモデルの重みは凍結され、追加のニューラルルーターは訓練せず、テスト入力ごとの再探索も行わない。生成されるパスはタスク単位で固定され、デプロイ時にはすべてのサンプルが同じパスを共有する。

探索器には、現在の層、残りの計算予算、移動方向、演算タイプを状態とするMarkov policyが用いられる。各ラウンドで30本の候補パスをサンプリングし、分割した訓練データ上で評価した後、高スコアのサンプルを使って遷移確率を更新する。この処理を最大10ラウンド実行する。完了後、システムはtop-k Viterbiを用いて、確率が最も高い有効なパス5本を厳密に求め、最後に独立した検証セットで1本を選択する。これにより、MCTSでルーティングラベルを生成する必要があったDr.LLMの探索時間は、著者らの測定で14.8時間から1.6時間に短縮された。

Qwen3 1.7B、4B、8B、14B、Llama 3.2 3B、Mixtral 8x7B、DeepSeek-R1-Distill-Llama-8Bなどの構成を用い、13の推論および知識ベンチマークで評価したところ、MACROは通常の層順序に比べて平均5.0ポイント精度を向上させ、Dr.LLMを7.2ポイント上回った。特に顕著なのがQwen3-1.7BのGSM8Kでの結果だ。第0層から第7層まで実行した後に第3層へ戻り、5つのblockを再実行してから末端まで進むことで、精度は43.4%から69.5%へ上昇した。このパスではblockが計33回呼び出されるのに対し、元のモデルでは28回だった。

したがって、MACROは計算コストなしで高速化できる手法ではない。一部のパスでは層の呼び出し回数、KV cacheのスロット数、レイテンシが増加し、平均的な改善効果も主に性能の低い小規模モデルに集中している。また、タスクごとに最大1,000件の訓練サンプルと独立した検証セットが必要であり、オープンエンドなワークロードや継続的に変化するワークロードには実用的でない可能性がある。ライブラリではすでに6種類のモデル構成、3つのrandom seedを用いた再現コマンド、MIT Licenseが提供されている。次に確認すべきなのは、固定パスを類似タスク間で転移できるか、そして実際の生成処理、long context、バッチサービスでも追加計算を相殺できるかどうかだ。