[論文レビュー] Adaptive Loops and Memory in Transformers: Think Harder or Know More?

Markus Frey, Behzad Shomali+4 more|arXiv (Cornell University)|Mar 9, 2026

Topic Modeling被引用数 0

ひとこと要約

要約: 本論文は、適応的な層ごとのループとゲート付き局所/グローバル記憶を組み合わせたトランスフォーマーを提案し、ループ（思考を深めること）と記憶（より多くを知ること）が数学的推論と常識推論に与える影響を研究する。ループは数学性能を向上させ、記憶は常識に寄与し、両者を組み合わせるとアイソFLOPの基準よりも強い結果が得られる。

ABSTRACT

Chain-of-thought (CoT) prompting enables reasoning in language models but requires explicit verbalization of intermediate steps. Looped transformers offer an alternative by iteratively refining representations within hidden states. This parameter efficiency comes at a cost, as looped models lack the storage capacity of deeper models which use unique weights per layer. In this work, we investigate transformer models that feature both adaptive per-layer looping, where each transformer block learns to iterate its hidden state via a learned halting mechanism, and gated memory banks, that provide additional learned storage. We find that looping primarily benefits mathematical reasoning, while memory banks help recover performance on commonsense tasks compared to parameter and FLOP matched models. Combining both mechanisms yields a model that outperforms an iso-FLOP baseline, with three times the number of layers, across math benchmarks. Analysis of model internals reveals layer specialization: early layers learn to loop minimally and access memory sparingly, while later layers do both more heavily.

研究の動機と目的

• 適応的な層ごとのループがトランスフォーマーの推論を改善するかを調査する。
• ゲート付き記憶バンクがループモデルで欠如する容量を回復できるかを検証する。
• ループと記憶が数学と常識タスクの性能にどう相互作用するかを分析する。
• ループと記憶を使用する際の内部レイヤーの専門化を特徴づける。

提案手法

• デコーダーのみのトランスフォーマーを適応的ループと局所/グローバル記憶バンクで拡張する。
• 学習済みのハルティング機構を用いて最大Nmax回のイテレーションにわたる中間ループ状態に重み付けを行う。
• 局所（レイヤーごと）およびグローバル（共有）記憶KVバンクを導入し、ゲート付き記憶統合を行う。
• 約12層・約200Mパラメータのモデルで14Bトークンを用いて訓練し、常識と数学のベンチマークをBPBと精度で評価する。
• ループと深さの増加の効率性を評価するために、アイソパラメータ・アイソFLOPのベースラインと比較する。

実験結果

リサーチクエスチョン

• RQ1 適応的な層ごとのループは標準的なトランスフォーマーと比べて数学的推論を改善するか。
• RQ2 特に常識タスクにおいて、ループ化されたモデルの容量低下を記憶バンクが補えるか。
• RQ3 ループの深さと記憶の成長は性能と内部ダイナミクスにどのように影響し合うか。
• RQ4 初期層と後半層ではループと記憶の利用に差があるか。

主な発見

Model	CS Acc	CS BPB	Math BPB
IsoPar	0.477	0.859	2.163
Loop-3	0.501	0.813	1.687
Loop-5	0.503	0.823	1.737
Loop-7	0.498	0.832	1.659
IsoFLOP	0.523	0.780	1.801
IsoPar-M	0.459	0.823	2.108
Mem (g0=-3)	0.472	0.810	1.619
Mem (g0=0)	0.481	0.810	1.662
Mem (g0=3)	0.511	0.794	1.616
IsoFLOP-M	0.535	0.749	1.761

• 適応的ループ（Nmax=3）は数学のBPBを約0.476ポイント改善（2.163から1.687へ）、常識正解率をわずかに向上（0.477から0.501へ）。
• より多くのループは数学での利得が次第に小さくなり、追加の反復で常識性能がわずかに低下する可能性。
• 局所/グローバル記憶を追加すると、ループだけよりも数学と常識の両方でさらなる改善が見られ、相補的な効果が示唆される。
• 記憶補強付きループは、数学タスクでアイソパラメータのベースラインより優れることが多く、深さを増やすベースラインに対する常識のギャップを縮小する。
• レイヤー分析では、初期レイヤーはループを最小限にとどめ、記憶にも依存が少ない一方、後半レイヤーはより多くループし、記憶を多く利用しており、思考を深めることと知識を増やすことの専門化を示す。
• ループと記憶の組み合わせは、三分の一の層数でアイソFLOPベースラインを上回りつつ、数学ベンチマークで競争力のある性能を達成できる。