Dualformer：ランダム化された推論軌跡学習による制御可能な高速思考と低速思考
https://arxiv.org/pdf/2410.09918v2
人間の認知理論では、思考は2つのシステムによって支配されます。迅速で直感的なシステム1と、より遅く推論的なシステム2です。最近の研究では、システム2の思考プロセスをTransformerモデル（大規模言語モデルLLMを含む）に統合することで、推論能力が著しく向上することが示されています。しかし、システム2思考のみを模倣するモデルは、より高い計算コストと遅い応答速度を必要とします。この課題を解決するため、我々は高速思考と低速思考の両方の推論モードをシームレスに統合できる単一のTransformerモデル、Dualformerを提案します。
Dualformerは、ランダム化された推論軌跡を含むデータで訓練されます。訓練中、推論軌跡の異なる部分がランダムに破棄されます。これらの破棄戦略は、思考プロセスを分析し、パターンを通じて思考のショートカットを形成することに似ており、推論軌跡の構造に基づいて特別に設計されています。推論段階では、我々のモデルは、解決策のみを出力する（高速モード）、推論チェーンと最終回答の両方を出力する（低速モード）、またはどちらのモードを有効にするかを自動的に決定する（自動モード）ように設定できます。
すべての場合において、Dualformerはパフォーマンスと計算効率の両面で対応するベースラインモデルを上回ります。(1) 低速モードでは、Dualformerは未確認の30x30迷路ナビゲーションタスクにおいて97.6%の最適解決率を達成し、完全な推論軌跡データで訓練されたSearchformerベースラインモデルの93.3%の性能を上回り、推論ステップ数を45.5%削減しました。(2) 高速モードでは、Dualformerはこれらのタスクで80%の最適率を達成し、回答データのみで訓練されたSolution-Onlyモデル（最適率30%）を大幅に上回りました。(3) 自動モードでは、Dualformerは96.6%の最適率を達成し、Searchformerと比較して推論ステップ数を59.9%削減しました。
数学問題では、大規模言語モデルのファインチューニングを通じて我々の方法も性能向上を達成しており、この技術が特定タスクモデルを超えて広く適用可能であることを示しています。
1 はじめに
心理学において、二重過程理論（Wason and Evans, 1974）は、思考が2つの異なる過程を通じて生じると説明しています。通常、一方は暗黙的（自動的かつ無意識的）であり、もう一方は明示的（制御され、意識的）です。
有名な書籍『ファスト&スロー』（カーネマン, 2017）は、二重過程理論の概念を深く掘り下げています。この本では2つの異なる思考様式が記述されています。システム1は暗黙の過程に対応し、迅速で直感的かつ感情的です。システム2は明示の過程に対応し、より遅く、より推論的で、より論理的です。カーネマンはこの理論を様々な認知活動や行動活動に広く適用し、これら2つのシステムが私たちの意思決定と推論プロセスにどのように影響するかを説明しています。本に記載されているように、システム1は能動的な制御の欠如により認知バイアスや系統的なエラーを引き起こしやすい一方、システム2は注意深い分析、検討、計画を必要とするタスクに適しています。
Transformer（Vaswani et al., 2017）は、シーケンスモデリングツールであり、大規模言語モデル（LLM）を含む様々な基盤モデルの中核技術となっています（Dosovitskiy, 2020; Baevski et al., 2020; Radford et al., 2021; Touvron et al., 2021; Hsu et al., 2021; Touvron et al., 2023; Dubey et al., 2024）。多くの研究で、Transformerは推論や計画の問題解決に広く使用されています。例えば、Zhou et al. (2022); Kojima et al. (2022); Pallagani et al. (2022); Valmeekam et al. (2023a); Chen et al. (2024); Gundawar et al. (2024); Wang and Zhou (2024)などです。興味深いことに、二重過程理論は人工知能の分野にも適用でき、Transformerの推論モードを「高速」と「低速」の2つに分類することができます。高速モードでは、Transformerは推論ステップを含まず最終的な解決策を直接出力します。一方、低速モードでは、モデルは最短経路探索時の探索軌跡のような思考の中間ステップを生成し、最終的な計画とともに出力します。これら2つの推論モードは、人間自身が持つ2つの思考システムにおける長所と短所に非常に似ています。先行研究（Wei et al., 2022; Valmeekam et al., 2023b; Lehnert et al., 2024; Gandhi et al., 2024; Saha et al., 2024）で議論されているように、高速モードで動作するモデルは計算コストが低く、応答速度が速い一方で、精度や最適性においては低速モードのモデルに劣ります。
そこで、「人間が2つの異なる思考システムを持っているように、Transformerベースの推論エージェントに高速モードと低速モードを統合し、互いに補完させることができるか？」という問いが生まれます。
これまでに様々なアプローチが提案されています。一般的なパラダイムは、まず純粋なシステム2モデルを構築し、次にファインチューニングによってシステム1のようにより効率的に動作させるというものです。例えば、システム2の出力をシステム1の形式に「蒸留」する（Yu et al., 2024; Deng et al., 2024）、あるいは知識蒸留（Wang et al., 2023; Shridhar et al., 2023）や記号システムからのブートストラップ（Searchformer（Lehnert et al., 2024）、Stream of Search（Gandhi et al., 2024）など）によって既存のシステム2の推論効率を向上させる方法があります。しかし、これらの方法では、しばしば追加のファインチューニングが必要となり、計算コストが高く、実行時に最終システムをシステム1とシステム2の間で柔軟に切り替えることが困難です。この問題を解決するため、Saha et al. (2024)は、2つの異なるシステムを切り替えるための明示的なメタコントローラを設計しました。
本研究では、推論タスクを解決する際に、システム1とシステム2を即座に構成するために、シンプルなデータ処理方法で十分であるという驚くべき発見を示します。これにより得られたDualformerモデルは、推論プロセス中に高速モードまたは低速モードで動作するように簡単に構成でき、モードが指定されていない場合は、モデル自体がどちらのモードを使用するかを決定することもできます。具体的には、システム2の推論プロセスを模倣するために、我々のTransformerは推論軌跡と最終的な解決策を含むデータで訓練されます。推論ステップの構造的特徴を利用して、生成された軌跡がシステム1が思考プロセスで取る「ショートカット」に似るように、特定の軌跡破棄戦略を設計しました。極端なケースでは、推論軌跡全体を完全に破棄し、Transformerに中間ステップをすべてスキップして最終的な解決策を直接出力するように促します。訓練プロセス中、これらの構造化された軌跡破棄戦略をランダムに選択します。詳細は第4節を参照してください。
我々はまず、このフレームワークをエンコーダ・デコーダ構造のTransformerモデルの訓練に適用し、A*探索アルゴリズムによって推論軌跡が生成される経路探索問題を解決しました。我々はLehnert et al. (2024)の設定と同様に、迷路ナビゲーションとソコバンゲームの2つのタスク領域を検討し、同じトークン化スキームを採用しました。興味深いことに、これらの問題は、現在の最先端の大規模言語モデル（o1-previewやo1-miniなど）にとって挑戦的であり、出力経路がしばしば「壁を突き抜ける」こと（付録Gの例を参照）がわかりました。
各推論モードにおいて、Dualformerは既存のベースラインモデルを上回り、問題解決率と最適率の両方で優れた性能を達成しました。さらに、Dualformerは生成される計画の多様性を著しく向上させ、目標へのより多くのユニークな経路を見つけることができます。注目すべきは、低速モードであっても、Dualformerは非常に高い効率を示し、ベースラインモデルよりもはるかに短い推論軌跡を生成しました。次に、我々はこのフレームワークを大規模言語モデルのファインチューニングに適用し、数学問題に回答させました。Yu et al. (2023)の方法を参照し、訓練サンプルはMATHデータセット（Hendrycks et al., 2021）から得られ、その回答は詳細な中間推論ステップを含むようにLlama-3.1-70B-Instructモデルによって書き換えられました。同様に、得られたLLMもより高い有効性と効率を示しました。
2 関連研究
計画と推論の学習Transformerベースのモデルにおける長期的なタスクでの計画と推論能力を向上させるため、研究者たちは多大な努力を払ってきました。現在、主に2つのアプローチが開発されています。
1つ目のアプローチは、既存の大規模言語モデル（LLMs）を利用するものです。例えば、研究者たちは、Ahn et al. (2022)、Besta et al. (2024)、Sel et al. (2023)、He-Yueya et al. (2023)、Liu et al. (2023)、Silver et al. (2024)の研究のように、大規模言語モデルを訓練して外部の既存の記号ソルバーを呼び出させます。Pallagani et al. (2022, 2024)は、大規模言語モデルを記号ソルバーを使用するようにファインチューニングすることを研究しました。Schick et al.、Hao et al. (2024)は、大規模言語モデルを外部ツールを使用するようにファインチューニングしました。Hao et al.は、大規模言語モデルをワールドモデル（world model）内部で推論と計画を行うようにファインチューニングすることを提案しました。これらの研究の中には、システム1とシステム2の思考モードを大規模言語モデルの推論プロセスに統合しようと試みているものもあります。WestonとSukhbaatar (2023)は、「システム2アテンション」（System 2 Attention, S2A）を提案しました。これは、コンテキストにおける誤った相関へのモデルの依存を減らすことにより、大規模言語モデルの推論能力を向上させることを目的とした、より熟慮的なアテンションメカニズムです。Yu et al. (2024)は、システム2の出力をシステム1モデルに「蒸留」し、中間推論トークンシーケンスなしで、システム2の方法で生成された高品質の結果を大規模言語モデルの生成プロセスに再コンパイルすることを目的としています。
2つ目のアプローチは、Transformerモデルをゼロから訓練し、独自に計画と推論ができるようにすることを目指しています（Lehnert et al., 2024; Saha et al., 2024; Gandhi et al., 2024）。これは通常、特定のタスクの言語表現を使用します。Lehnert et al. (2024)とGandhi et al. (2024)は、言語化された探索プロセスを通じて大規模言語モデルを探索するように訓練することを試みています。Lehnert et al. (2024)はSearchformerを開発しました。このモデルは、経路計画問題におけるA*探索アルゴリズムのプロセスを模倣するように訓練されました。Gandhi et al. (2024)は、彼らのモデルを「カウントダウンゲーム」に適用しました。同時期の研究では、Saha et al. (2024)は2つの独立したモデルを訓練し、外部のメタコントローラによって管理しています。一方のモデルは高速ですが推論軌跡のない応答を生成し、もう一方はより遅いですが完全な推論チェーンを含む応答を生成します。同様に、Lin et al. (2023)は複数のモデルを利用し、エージェントベースのワークフローを通じて低速思考と高速思考を実現しています。我々の研究は上記の研究と密接に関連していますが、重要な違いがあります。Lehnert et al. (2024)やGandhi et al. (2024)が訓練データ中の推論軌跡を変更しないのに対し、我々の方法は訓練プロセス中に推論軌跡をランダム化します。Saha et al. (2024)やLin et al. (2023)とは異なり、我々は明示的なコントローラを使用せず、各モード用に2つのネットワークを別々に訓練することもありません。その代わりに、高速モードと低速モードの機能を単一のモデルに統合しています。
大規模言語モデルを用いた合成データ生成大規模言語モデル（LLM）は、様々な分野で合成データ生成に広く利用されています。例えば、Wei et al. (2021)、Longpre et al. (2023)、Taori et al. (2023)は、自然言語の指示を含む多様なテンプレートからサンプリングすることで、合成指示データセットを構築しました。この方法は、視覚分野（Liu et al., 2024b,a; Zhu et al., 2023; Brooks et al., 2023; Peng et al., 2023）にも応用されています。大規模言語モデルの数学問題解決能力を向上させるため、Yu et al. (2023)は、Llama-3.1-70B-Instructモデルを呼び出して、元のMATHデータセット（Hendrycks et al., 2021）の回答を指定された形式で詳細な中間推論ステップを含むように書き換え、検証し、拡張する方法を開発しました。同様の方法は、Yuan et al. (2023)、Luo et al. (2023)、Lee et al. (2023)、Yue et al. (2023)、Tong et al. (2025)などの他の研究でもさらに探求されています。
3 予備
我々の研究は、Lehnert et al. (2024)の研究に基づいています。計画を実現するために、我々はいくつかのトークンからなるシーケンスをモデル化するためのTransformerモデルを訓練します。このシーケンスは、計画タスク、A*アルゴリズムの計算プロセス、およびA*探索から得られる最適解を順次表現します。トークン化方法は図3.1に示されています。
簡単な例で説明します。3×3の迷路におけるナビゲーションタスクを考えます。目標は、開始セルから目標セルまで、壁セルを通過せずに最短経路を見つけることです。A*アルゴリズムはすでに最適経路を特定しています。我々は一連のトークンを使用して、このタスクと迷路の構造を同時に表現します。これらのトークンはDualformerの入力プロンプトとしても機能します。解は、経路座標を表す「計画トークンシーケンス」として提示されます。
A*アルゴリズムは、図4.1に示すように、探索プロセス中の動的な振る舞いを記録する探索軌跡シーケンスを生成します。A*アルゴリズムは、加重グラフ上の経路探索アルゴリズムであることを思い出してください。「create」句はノード（後続の座標で示される）を探索フロンティア（frontier）に追加し、「close」句はノードをクローズドセット（closed set）に追加します。createまたはcloseの各句の後には、そのノードの座標、開始点からの実際のコスト（cost-since-start）、およびヒューリスティックな推定値をそれぞれ表すx、y、c0、c1の4つのトークンが続きます。A*アルゴリズムとトークン化方法の詳細については、Russell and Norvig (2021)およびLehnert et al. (2024)の研究を参照することをお勧めします。
4 構造化された軌跡破棄とランダム化訓練
Lehnert et al. (2024)によって提案されたSearchformerは、様々な複雑な意思決定タスクの解決において有効性が証明されています。しかし、依然として2つの重要な限界があります。
第一に、このモデルは低速モードでのみ動作し、冗長な推論チェーンを出力するため、推論時間が著しく増加します。この問題は、ブートストラップ（bootstrapping）方法（Lehnert et al., 2024）によって軽減できますが、これは複数回のロールアウトとファインチューニングのループを含む反復最適化技術であり、著しく増加する追加の計算リソースを必要とします。
第二に、Searchformerは多様な解決策を生成することに困難があり、同じ推論経路を繰り返しサンプリングすることがよくあります。例えば、我々がテストした1000個の30x30迷路問題において、Searchformerの推論チェーンは平均1500以上のトークンを含み、64回の応答で7.6種類の異なる実行可能経路しか見つけることができませんでした（第5.1.2節参照）。
これらの課題に対処するため、我々はランダム化された推論軌跡を利用した訓練フレームワークを提案します。我々のアプローチは、2つの研究から着想を得ています。第一に、Searchformerが完全なA*探索軌跡で訓練されているにもかかわらず、生成される推論チェーンはより短く、探索プロセスを大まかにしか描写していないことに気づきました。第二に、人間は意思決定の際に「ショートカット」やパターンに依存することが多く、この概念はシステム1思考（Kahneman, 2017）として知られています。これらの観察と、訓練中にニューラルネットワークの一部ユニットをランダムに破棄するdropout技術（Hinton, 2012; Srivastava et al., 2014）の成功経験が、フレームワークでランダム化された推論軌跡を使用する可能性を探求するきっかけとなりました。我々は、A*探索軌跡の構造的特徴を利用し、各訓練サンプルに対して軌跡内容の一部を選択的に破棄することを目指します。具体的な方法は以下の通りです。
図4.1に示すように、A*探索軌跡には「create」と「close」の2つの句が含まれており、各句にはノードの座標、出発点からの（推定）コスト、目標までのヒューリスティックコストが含まれます。Dualformerを得るために、我々は探索軌跡の構造を利用し、各訓練サンプルからその情報の一部を破棄します。3種類の自然な破棄タイプがあります。
• D1：close句を破棄する
• D2：句中のコストトークン（cost tokens）を破棄する
• D3：create句を破棄する
これらの基本タイプに基づき、各レベルが前のレベルをさらに強化する4つのレベルの破棄戦略を設計しました。
• レベル1戦略：探索軌跡からすべてのclose句を削除する。
• レベル2戦略：レベル1戦略に基づいて、さらにすべての句中のコストトークンを削除する。
• レベル3戦略：より積極的で、追加で30%のcreate句をランダムに破棄する。
• レベル4戦略：最終戦略であり、推論軌跡全体を直接破棄する。
訓練プロセス中にこれらの構造化された破棄戦略をランダムに適用することで、モデルは不完全または高度に簡略化された推論プロセスから直接正しい回答を導き出すことを学習するように促され、それによってシステム1の迅速な直感的な思考を模倣しつつ、システム2の完全な推論能力を保持します。
図4.1は、先に述べた迷路タスクを例に、我々の破棄戦略を示しています。直感的に、レベル1破棄はDualformerにA*探索の「クローズドセット」（close-set）計算を効果的にスキップするよう指示します。レベル2破棄はDualformerに「クローズドセット」とコスト計算を同時にスキップするよう促します。レベル3およびレベル4破棄は、Dualformerに探索ステップの一部またはすべてを省略するよう促します。第5節で示すように、これらの戦略はDualformerがより簡潔で効率的な探索と推論プロセスを学習するのに効果的です。
訓練データの多様性を向上させるため、我々は破棄操作をデータ前処理ステップとして行いませんでした。その代わりに、訓練プロセス中、各バッチ内の各訓練サンプルに対して、分類分布Cat(p₀, p₁, p₂, p₃, p₄)から破棄戦略をランダムにサンプリングします。ここで、p₁からp₄はレベル1からレベル4の破棄戦略を採用する確率を、p₀は完全な推論軌跡を保持する確率を表します。この訓練フレームワークにより、Dualformerは同じ訓練サンプルであっても、異なるバッチで異なる破棄方法で繰り返し現れるため、様々な簡略化された軌跡から学習することができます。
トークンマスキング技術との比較興味のある読者は、我々の訓練フレームワークがBERTなどの有名な大規模言語モデルで使用されているトークンマスキング（token masking）技術（Devlin, 2018; Liu, 2019; Song, 2019; Gauthier and Levy, 2019; Sinha et al., 2021; Kitouni et al., 2024）に似ていると思うかもしれません。しかし、我々の方法はこれらのマスキング技術とは著しく異なります。まず、標準的なマスキング技術は通常、シーケンス内の入力トークンを一様にランダムにマスクしますが、我々の破棄戦略は探索軌跡の部分にのみ適用されます。次に、マスクされた大規模言語モデルは通常、双方向アテンションメカニズムを採用し、マスクされたトークンを予測することを目標としますが、Dualformerは因果アテンションメカニズム（causal attention）を使用し、その訓練目標は次のトークン予測のみであり、全体目標はモデルの推論と計画能力を向上させることです。計算効率の面でも、我々の訓練プロセスは優位性があります。トークンを破棄することで入力シーケンスが短縮され、計算リソースが節約されます。例えば、8枚のTesla V100 32GB GPUでDualformerを訓練して30x30迷路タスクを完了するには30時間かかりますが、完全な推論軌跡を使用すると36時間かかります。さらなる訓練の詳細は、第5節および付録Aを参照してください。
4.1 制御可能な生成
Dualformerの魅力的な特徴の1つは、推論段階でプロンプトを通じて簡単に制御でき、「高速」または「低速」モードで結果を生成できることです。制御メカニズムは非常にシンプルです。標準プロンプト（環境とタスク記述を含む）に加えて、開始トークンbosと制御トークン（control token）であるplanまたはcreateを追加します。planを使用する場合、Dualformerは高速モードに入り、すべての中間推論ステップをスキップして最終的な経路計画を直接出力します。逆に、bosの後にcreateを挿入すると、Dualformerは低速モードに入り、完全な推論軌跡と最終的な解決策を生成します。具体的な例は付録Bを参照してください。標準プロンプトのみを使用する場合（制御トークンなし）、Dualformerは人間の意思決定の二重過程メカニズムを模倣し、具体的な状況に応じて、システム1とシステム2の推論モードに対応する応答タイプを自律的に決定します。
5 実験
我々の実験は以下の質問に答えることを目的としています。
1. Dualformerは、高速モード、低速モード、自動モードで、それぞれ対応するベースラインモデルよりも優れていますか？また、より多様な計画を生成できますか？
2. 低速モードにおいて、Dualformerはより高速な推論、すなわちより短い推論軌跡を実現できますか？
3. 構造化された軌跡破棄技術は、自然言語データセットで訓練された大規模言語モデル（LLMs）に一般化できますか？
質問1と2については第5.1節で回答し、Transformerモデルを訓練して迷路ナビゲーションタスクおよび関連するソコバンゲームを解決する方法について、Searchformer（Lehnert et al., 2024）と同様の方法で説明します。質問3に答えるため、第5.2節ではLlama-3.1-8BおよびMistral-7Bモデルをファインチューニングし、数学問題を解決できるようにします。
5.1 ナビゲーションタスク：迷路とソコバン
我々はLehnert et al. (2024)の方法を踏襲し、迷路とソコバンタスクを検討し、同じデータセットを使用します。両方のデータセットには、完全なA*探索軌跡を持つ10⁵個の訓練サンプルが含まれています。A*アルゴリズムの実装は非決定的です。経路コストが同じ場合、アルゴリズムはランダムに選択します。子ノードの展開順序もランダムです。迷路のサイズは15x15から30x30まで変化します。すべての迷路タスクについて、30%～50%の壁セルを障害物としてランダムに生成し、開始位置と終了位置をランダムにサンプリングします。ソコバンマップのサイズは7x7で、ランダムに配置された2つのドック、2つの箱、および1つの作業員の位置が含まれます。マップ内にはさらに2つの壁セルをランダムに追加します。マップ生成の具体的なプロセスと例の画像については、付録A.1を参照してください。
まず、第4.1節で、Dualformerが高速モードまたは低速モードで動作するように明示的に制御できることを示しました。高速モードでは最終的な解決策のみを出力し、低速モードでは完全な推論軌跡を生成します。第5.1.1節と第5.1.2節では、Dualformerを各モードで対応するベースラインモデルと比較します。我々は、生成される解決策の正確性、最適性、多様性、および推論軌跡の長さなど、複数の指標を用いてモデルの性能を系統的に評価します。最後に、第5.1.5節では設計選択に関するアブレーション実験を行います。
ハイパーパラメータ設定我々は、Lehnert et al. (2024)と同じエンコーダ・デコーダアーキテクチャでDualformerをインスタンス化します。エンコーダはT5アーキテクチャ（Raffel et al., 2020）に基づいており、回転位置エンコーディング（rotary embeddings）を導入しています。デコーダはGPTスタイルのアーキテクチャを採用しています。迷路環境では1500万パラメータ（15M）のモデルを、ソコバン環境では4600万パラメータ（46M）のモデルを使用します。すべてのモデルは10万個の訓練サンプルで訓練されます。迷路環境は4x10⁵エポック、ソコバン環境は8x10⁵エポック訓練されます。モデルアーキテクチャおよびその他のハイパーパラメータの詳細は付録Aを参照してください。
我々は第4節で説明した構造化された軌跡破棄戦略を用いてDualformerを訓練します。各訓練サンプルに適用される破棄戦略については、3つの確率設定を試しました。(1) {p₀ = 0.45, p₁ = p₂ = p₃ = 1/6, p₄ = 0.05}、(2) {p₀ = 0.8, p₁ = p₂ = p₃ = p₄ = 0.05}、(3) {p₀ = 0.7, p₁ = 0.05, p₂ = p₃ = 0.1, p₄ = 0.05}。そして、検証誤差が最も低いものを選択しました。最終的な選択結果は、迷路タスクでは(1)のグループ、ソコバンタスクでは(3)のグループです。
ベースラインモデル高速モードの場合、我々のベースラインは「Solution-Only model」であり、このモデルはDualformerと同じアーキテクチャを持ちますが、最適な最終解のみを含み推論軌跡を含まないシーケンスデータで訓練されます。低速モードの場合のベースラインは「Complete-Trace model」であり、これは完全なA*探索軌跡を含むデータで訓練されたモデルです。このモデルはLehnert et al. (2024)で「探索拡張モデル」（search augmented model）と呼ばれており、探索動的ブートストラップ（bootstrapping）が行われていない元のSearchformerモデルです。第5.1.2節では、ブートストラップされたモデルとDualformerを比較します。これらのモデルはすべて同じパラメータ数です。迷路問題では1500万、ソコバン問題では4600万です。
評価指標我々は、モデルが正しく最適な解決策を生成したかどうかを評価するために、1-Solved-64と1-Optimal-64という2つの指標を使用します。具体的には、各評価タスク（例えば迷路やソコバンゲーム）に対して、訓練済みのモデルから64個の応答をランダムにサンプリングします。生成された推論軌跡がどうであれ、我々は最終的な解決策の部分のみを解析し評価します。64個の解決策のうち少なくとも1つが正しい（実行可能で目標位置に到達できる）場合、そのタスクは1-Solved-64指標で成功とマークされます。少なくとも1つが最適である場合、1-Optimal-64指標で成功とマークされます。我々は1000個の未確認の評価タスクに対してこのプロセスを繰り返し、平均成功率を報告します。
各手法の堅牢性を検証するため、3-Solved-64および3-Optimal-64指標も報告します。これは、少なくとも3つの解決策が正しいか最適である場合にのみ、タスクが成功とマークされるものです。
さらに、生成された解決策の全体的な品質をそのコストに基づいて集計するために、「コスト加重成功率」（Success Weighted by Cost, SWC）（Wu et al., 2019）を採用します。SWC = (1/64) Σ₆₄ᵢ₌₁ I(解決策iが正しい) · (c∗ / cᵢ)ここで、Iは指示関数、c∗は最適解決策のコスト、cᵢはi番目の解決策のコストです。明らかに、SWCスコアが高いほど、生成される解決策は最適に近く、すべての生成された解決策が最適であればSWCは最大値1に達します。
最後に、生成される解決策の多様性を定量化するため、各タスクで64回の応答中に生成された一意の正しい解決策の数を集計し、1000個の評価タスクにおける平均値を報告します。
5.1.1 高速モード
表5.1は、Dualformerとベースラインモデル「Solution-Only model」の迷路およびソコバンタスクにおける性能を示しています。正しい解決策と最適な解決策の生成において、Dualformerは1-Solved-64と1-Optimal-64の両指標でベースラインモデルを著しく上回っています。同時に、3-Solved-64と3-Optimal-64指標でもベースラインを明らかに凌駕しており、Dualformerが解決策生成においてより高い堅牢性を持っていることを示しています。
特に、タスクの難易度が増すにつれて、性能差も拡大します。最大の30x30迷路タスクでは、Dualformerの1-Optimal-64率は「Solution-Only model」の2.8倍、3-Optimal-64率は2.97倍に達しています。DualformerのSWCスコアもベースラインモデルよりも著しく高く、すべての環境で0.9を超えており、Dualformerが生成する各解決策の品質が非常に高く、そのコストが最適解決策に非常に近いことを示しています。
さらに、検討されたすべての問題において、Dualformerは常に多様な解決策を生成できます。付録Cでは、迷路の例を示し、Dualformerとベースラインモデルによって生成された唯一の正しい経路をプロットしています。興味深い観察として、Dualformerの多様性スコア（つまり、64回の応答で生成された唯一の正しい解決策の平均数）は、迷路のサイズが大きくなるにつれて増加しています。直感的には、迷路が大きくなるほど、同じ目標位置への可能な経路が増えるためです。
これは、Dualformerが迷路の構造的特徴を学習したことを示しており、「Solution-Only model」は最適な経路を記憶しているに過ぎない可能性があります。なぜなら、その多様性スコアはすべての迷路サイズで1に近いためです。
5.1.2 低速モード
表5.2は、Dualformerが低速モードで動作した場合の結果を示しています。対応するベースラインモデルは「Complete-Trace model」であり、これは同じアーキテクチャを使用し、完全なA*探索軌跡を含むデータで訓練されています。これまで報告された各指標に加えて、全1000の評価タスクにおいて、64回の応答で生成された推論軌跡の平均長も集計しました。
結果は、Dualformerが計画能力と推論速度の両方で向上を達成したことを示しています。問題解決率、最適率、SWCスコアを含むすべての正確性および最適性指標において、DualformerはComplete-Trace modelを上回っています。さらに、Dualformerが生成する推論軌跡は、ベースラインモデルよりも明らかに短くなっています。平均して、5つのタスクでDualformerは推論軌跡の長さを49.4%削減しました。以前と同様に、Dualformerが生成する解決策はベースラインモデルよりも多様性に富んでいます。具体的な例は付録Cを参照してください。
探索動的ブートストラップ方法との比較
完全軌跡モデルは、Lehnert et al. (2024)で提案されたSearchformerの基本モデルです。この研究では、Anthony et al. (2017)やZelikman et al. (2022)の方法と同様に、ソコバンタスクにおけるモデルの性能を向上させるための「探索動的ブートストラップ」（search dynamics bootstrapping）方法も提案されています。Searchformerモデルの訓練後、新たに構築されたブートストラップデータセットでモデルがファインチューニングされます。具体的には、元のデータセット内の各ソコバンステージについて、32個の応答が生成され、その中で最短の最適応答が新しい訓練データセットに追加されます。このプロセスは複数回繰り返すことができ、これによりSearchformerは徐々により短い応答を生成することを学習します。
表5.4は、Dualformerと最大3ステップのブートストラップファインチューニングを施したSearchformerモデルを比較したものです。結果は、Dualformerがほとんどの指標でブートストラップされたモデルと同等またはそれ以上の性能を達成し、同時に使用する推論ステップ数を45.1%以上削減したことを示しています。
注目すべきは、各ブートストラッププロセスには合計3.2x10⁶個の応答サンプルの生成と、追加で10⁴エポックのファインチューニングが必要であることです。これは、初期の8x10⁵エポックの事前訓練を含めると、Searchformerが3ステップ目のブートストラップまでに合計8.3x10⁵エポックの訓練と9.6x10⁶回のロールアウトを必要とし、計算コストが非常に高くなることを意味します。対照的に、Dualformerは1つの訓練段階で合計8x10⁵エポックしか必要とせず、追加のロールアウトプロセスは一切不要であるため、計算オーバーヘッドが著しく削減されます。
5.1.3 自動モード：二重過程
Dualformerの推論モードをbosの後に制御トークンを挿入して制御するだけでなく、サンプリングを直接行って、人間が意思決定する際の二重過程メカニズムと同様に、自由に動作モードを決定させることもできます。我々はこのモードをDualformerの自動モード（auto mode）と呼んでいます。表5.3に関連する結果を報告します。検討したすべてのタスクにおいて、自動モードのDualformerは、完全軌跡モデルと解決策のみのモデルの両方を上回っています。
Dualformerが問題の難易度に応じて動作モードを自動的に調整できるかどうかは、検討に値する興味深い問題です。これを調査するため、未確認の迷路1000個（壁密度が0.3から0.5の間で変化）それぞれに対して自動モードで64個の応答を生成し、次にすべての実行可能な解決策のうち低速モードの経路が占める割合を分析しました。結果を図5.1に示します。壁密度が増加するにつれて（迷路がより困難になることを意味します）、低速モードの経路の割合も増加します。これは2つの要因によって引き起こされる可能性があります。(1) より困難な問題を解決すること自体がより遅い思考プロセスを必要とするため、(2) Dualformerがより頻繁に低速モードを選択するためです。同様に、迷路のサイズが大きくなり、問題の難易度が上がるにつれて、Dualformerが継続的により多くの低速思考モードを採用することも観察されました。
5.1.4 汎化性能
これまでのところ、我々のすべての実験は、訓練セットと同じサイズと壁密度の迷路でテストされています。次に、Dualformerの分布外（OOD）シナリオにおける汎化能力をさらに調査します。本節では、20x20サイズの迷路で訓練されたDualformerモデルを検討します。訓練時の壁密度は0.3から0.5の間で一様にランダムにサンプリングされます。テスト段階では、壁密度を0.1から0.6まで変化させます。表5.5は、低速モードでのテスト結果を示しており、各ケースで50個の未確認の迷路インスタンスをテストしています。
壁密度が増加すると、迷路の難易度が上がり、同時に入力プロンプトの長さも増加します。したがって、分布内のテストケース（壁密度が0.3、0.4、0.5）では、Dualformer（低速モード）はほぼ100%の最適率を達成しました。しかし、壁密度が0.6（分布外ケース）に増加すると、性能は低下しました。興味深いことに、壁密度が低い（例えば0.1と0.2）場合、Dualformer（低速モード）はどの迷路も成功裏に解決できませんでした。我々の推測では、壁密度が低すぎると、プロンプト情報が短すぎて、モデルが汎化しにくくなるためです。
壁密度を変更するだけでなく、Dualformerの長方形迷路（例えば高さ20、幅10）での低速モードのパフォーマンスもテストしました。モデルは20x20の正方形迷路で訓練されています。表5.6はテスト結果を示しています。驚くべきことに、Dualformerはこれらの状況で良好な汎化能力を示しました。
5.1.5 アブレーション研究
第4節で述べたように、Dualformerの訓練に使用されるランダム化された軌跡は、複数の軌跡破棄戦略の組み合わせに由来し、これらの戦略の組み合わせ方は多様です。本節では、我々が行った設計選択についてアブレーション分析を行います。
まず、高速と低速の2つの推論モードを実行するための簡単な代替方法は、訓練時に「解決策のみのデータ」と「完全な軌跡データ」を混合することです。つまり、我々のランダム化戦略でp₁ = p₂ = p₃ = 0とします。我々はこの変種をMix-pモデルと呼びます。ここでpは訓練データ中の「解決策のみのデータ」の割合を示します。なお、「解決策のみのモデル」は本質的にMix-1モデルであり、「完全な軌跡モデル」はMix-0モデルです。次に、Dualformerを異なるp値のMix-pモデルと2つの推論モードで比較します。
次に、我々の破棄戦略は階層的かつ構造的に設計されています。例えば、レベル2の破棄はレベル1の戦略をさらに拡張したものです。特定のレベルで破棄戦略を停止した場合にモデルのパフォーマンスがどのように変化するかをさらに探求します。
Mix-pモデルとの比較図D.1は、DualformerとMix-pモデルの1-Optimal-64率を比較したものです。我々は8つのp値をテストしました。0（完全な軌跡モデルと同等）、0.05、0.1、0.2、0.4、0.6、0.8、1.0（解決策のみのモデルと同等）です。2つの推論モードにおいて、Dualformerは我々が検討した5つのタスクすべてでMix-pモデルを上回っています。実際、他のすべての指標においてもDualformerはMix-pモデルを全面的に上回っています。詳細は付録Dを参照してください。特に、低速モードでは、Dualformerは推論軌跡が最も短いモデル、つまり推論効率が最も高いモデルです。
ランダム化戦略の組み合わせ分析Dualformerを、特定のレベルで破棄戦略の適用を停止するその変種と比較しました。迷路環境では、各訓練サンプルで完全な推論軌跡を使用する確率p₀を0.5に固定し、他の破棄戦略の確率を調整しました。ソコバン環境では、レベル1の破棄の確率p₀を0.05に固定しました。表E.1は、すべてのモデルが使用した異なる破棄戦略の確率分布を示し、表5.7は実験結果を示しています。
これらの変種は、レベル4の破棄（つまり、軌跡の完全破棄）データで訓練されていないため、高速モードで実行できないことに注意してください。したがって、低速モードでのパフォーマンスのみを報告しています。破棄戦略のレベルが上がるにつれて、生成される推論軌跡の長さは徐々に短くなります。他の指標の面では、レベル1+レベル2+レベル3の破棄を採用したモデルはDualformerと同等の性能を示します。しかし、Dualformerは推論軌跡が短いだけでなく、高速モードで動作する能力も備えており、より包括的な優位性を示しています。
5.2 大規模言語モデル訓練への応用：数学推論
本節では、大規模言語モデル（LLM）を訓練して数学問題を解決する上での構造化された軌跡破棄技術の有効性を示します。具体的には、様々な数学問題と詳細な推論ステップが付属した回答を含むデータセットを使用し、Llama-3-8BおよびMistral-7Bモデルをファインチューニングします。同様に、数学問題の推論軌跡の特定の構造を利用して軌跡破棄方法を設計します。ファインチューニングされたモデルを、このデータセットで直接ファインチューニングされたベースラインモデルと比較します。
データセットすべてのモデルの評価には、Aug-MATHというデータセットを使用します。このデータセットは、元のMATHデータセット（Hendrycks et al., 2021）に由来し、元のデータセットには数学問題と解答の訓練サンプル7500個とテストサンプル5000個が含まれています。Yu et al. (2023)の方法を踏襲し、Llama-3.1-70B-Instructモデルを呼び出して、詳細な中間推論ステップを含むように元の回答を指定された形式で書き換えます。推論軌跡の多様性を向上させるため、各問題に対して温度0.7、top-p=0.9で4つの異なる大規模言語モデル応答をサンプリングします。最終的なデータセットには、訓練サンプル30,000個とテストサンプル5,000個が含まれています。付録F.1では、回答の書き換えに使用したプロンプトテンプレートと、書き換え前後の具体的な訓練例を示しています。
構造化された軌跡破棄とランダム化Llama-3.1-70B-Instructによって書き換えられた数学問題の回答は、平均して6から10個の中間推論ステップを含み、各ステップには複数の文が含まれる場合があります。我々は、第4節で提案されたフレームワークと同様のランダム化訓練方法を採用します。各訓練バッチにおいて、サンプリングされた各訓練サンプルについて、確率pでその中間推論ステップのそれぞれをランダムに破棄します。以下に、異なるp値での実験結果を報告します。
ハイパーパラメータ設定我々は、Mistral-7BとLlama-3-8Bの2つの基本モデルに対し、バッチサイズ32で2エポックのファインチューニングを行います。オプティマイザはAdamW（Loshchilov and Hutter, 2019）を採用し、Mistralモデルの学習率は8e-6、Llama-3モデルは5e-6です。学習率の選択プロセスは以下の通りです。2e-6、5e-6、8e-6の3つの値から探索し、検証損失が最も低い学習率を選択します。その後、選択された学習率を使用して完全な訓練データセットでモデルを再訓練し、結果を報告します。その他のハイパーパラメータについては、デフォルト値を使用します。具体的には、線形学習率ウォームアップ、重み減衰、多段階勾配累積は使用しません。AdamWには以下のパラメータを使用しました。betas=(0.9, 0.999)、eps=1e-8、γ=0.85（乗算ステップごとの学習率減衰）、そしてバッチ戦略として「packing」を採用しました。
評価方法我々は以下の2つの評価指標を採用します。
Greedy@1（Dubey et al., 2024; Yu et al., 2023）：各問題に対し、温度0で1つの応答を生成し、その回答が正しいかどうかを判断します。
pass@20（Chen et al., 2021）：各問題に対し、温度0.5で20個の応答をランダムにサンプリングし、少なくとも1つが正しい割合を統計します。
参考として、元のMATHデータセットでこれらのモデルを直接ファインチューニングした結果も報告します。
結果結果は表5.8に示されています。我々は4つのp値、すなわち0.1、0.2、0.3、0.4をテストしました。実験結果は、提案された訓練戦略が2つの大規模言語モデルの有効性と効率の両方を向上させることを示しています。
まずMistral-7Bモデルの結果を分析します。低速モード推論において、軌跡破棄とランダム化訓練でファインチューニングされたモデルは、Aug-MATHデータセットで直接ファインチューニングされたベースラインモデルよりも優れています。p = 0.1の場合、Greedy@1指標は絶対値で1.7%（相対性能で10%に相当）向上しました。p = 0.2およびp = 0.3の場合、向上は0.9%でした。p = 0.4の場合、向上は0.1%でした。
Pass@20指標では、p = 0.1、0.2、0.3の場合、我々のモデルもベースラインモデルを上回り、最高で61.9%の絶対正解率を達成しました。2つの評価方法において、p値が増加するにつれて、推論軌跡の平均長は徐々に短くなりました。
同様に、高速モード推論においても、我々のモデルはより高い正解率を達成しました。Llama-3-8Bモデルも同様の性能向上傾向を示しました。
最後に、読者の便宜のため、Mistral-7BおよびLlama-3-8Bモデルを元のMATHデータセットでファインチューニングした結果を列挙しました。これらのモデルの性能は、Aug-MATHでファインチューニングしたモデルよりも明らかに劣っています。
6 結論
我々は、Transformerモデルを訓練して推論および計画タスクを解決するためのシンプルで実装が容易なフレームワークを提案しました。推論軌跡の構造を深く分析し、人間の思考プロセスにおける「ショートカット」の振る舞いを模倣するための破棄戦略を設計しました。訓練サンプルにこれらの破棄戦略をランダムに適用することで、得られたDualformerモデルは、推論時に高速モード、低速モード、または自動モードで動作するように制御できます。自動モードでは、モデルがどのような推論方法を採用するかを自律的に決定します。Dualformerは、迷路ナビゲーションタスクとソコバンゲームの両方で性能向上を達成し、同時に必要な推論ステップ数を削減しました。
注目すべきは、我々の方法がゼロから特定のタスクモデルを訓練することに限定されないことです。我々は同じ考え方の技術を大規模言語モデル（LLM）のファインチューニングに適用し、数学問題の解答においても性能向上を達成しました。さらに、提案されたフレームワークは、推論軌跡を破棄することで入力シーケンスの長さを短縮できるため、計算オーバーヘッドも削減します。
今後の研究では、この方法がモデルのスケーラビリティに役立つかどうかをさらに探求し、カリキュラム学習（curriculum learning）や階層的計画（hierarchical planning）などの方法と組み合わせることで、Dualformerがより複雑なタスクシナリオに適応できるかどうかを試すことができます。
付録A：ネットワークアーキテクチャとハイパーパラメータ
我々はLehnert et al. (2024)と同じエンコーダ・デコーダTransformerアーキテクチャをDualformerに使用しています。それはまず各トークンをワンホットベクトルに変換し、次に埋め込み層を介して一連のベクトルに変換します。埋め込みベクトルはその後、Lehnert et al. (2024, 図4)に示されている後続の層を通過します。我々は位置エンコーディングにRoPE埋め込みを使用しており、我々のアーキテクチャではドロップアウトを使用していません。
モデルサイズ、アーキテクチャパラメータ、バッチサイズについては、Lehnert et al. (2024)と同じ設定を使用しています。具体的には、迷路タスクでは、3つのアテンションヘッドと6層、隠れ層サイズ64の1500万パラメータのモデルを使用しています。我々はAdamWオプティマイザ（Loshchilov and Hutter, 2019）を使用してモデルを最適化しており、学習率は2.5e-4、バッチサイズは16、β₀とβ₁はそれぞれ0.9と0.99に設定しています。最初の2000回の勾配更新では、線形ウォームアップ戦略を採用しました。その後は、コサイン学習率スケジューラ（Loshchilov and Hutter, 2016）を使用しています。
ソコバンタスクについては、4つのアテンションヘッドと8層、隠れ層サイズ96の4600万パラメータのモデルを使用しています。バッチサイズは64、学習率は7.5e-5を使用し、その他のハイパーパラメータは上記と同じです。
A.1 迷路とソコバン
我々はLehnert et al. (2024)と同じデータセットを使用しており、このデータセットはhttps://github.com/facebookresearch/searchformerで入手可能です。迷路データセットには1000万の例が含まれ、ソコバンデータセットにも1000万の例が含まれています。我々はmongodbの「id」フィールドに従ってこれらをソートし、最初の10万の例を使用しています（Lehnert et al. (2024)と同じ方法に従っています）。読者の参考のために、データセットは次のように生成されています。迷路タスクでは、まず30〜50%のセルがランダムに壁として指定されます。次に、開始位置と目標位置がランダムに選択されます。その後、A*アルゴリズムを適用して最適な計画を生成します。ソコバンタスクでは、7x7のグリッドマップを使用し、ランダムに2つの壁セルを障害物として挿入します。さらに、2つのドック、2つの箱、1つの作業員の位置もランダムに配置されます。ゲームが生成されたら、それがA*アルゴリズムで解決可能な場合にのみデータセットに追加されます。
A.2 数学推論
我々はhttps://github.com/meta-llama/llama-recipesで提供されている実装を使用してモデルをファインチューニングしています。我々はすべてのモデルに対して2エポックの訓練を、バッチサイズ32で行います。Llamaモデルでは、AdamWオプティマイザを使用し、学習率は5x10⁻⁶です。Mistralモデルでは、学習率は8x10⁻⁶です。学習率の選択プロセスは以下の通りです。2x10⁻⁶、5x10⁻⁶、8x10⁻⁶の3つの値から探索し、検証損失が最も低い学習率を選択します。その後、選択された学習率を使用して完全な訓練データセットでモデルを再訓練し、結果を報告します。その他のハイパーパラメータについては、デフォルト値を使用します。具体的には、線形学習率ウォームアップ、ウェイトディケイ、多段階勾配累積は使用していません。AdamWには以下のパラメータを使用しました。betas=(0.9, 0.999)、eps=1e⁻⁸、γ=0.85（乗算ステップワイズ学習率減衰）、そしてバッチ戦略として「packing」を採用しました。
C 生成計画の多様性
Dualformerは、ユニークで実行可能な解決策を見つける点でベースラインモデルを上回ります。これを視覚的に示すため、迷路タスクの例を選び、高速モードのDualformerを使用して64個の応答を生成しました。図C.1は、高速モードのDualformerが見つけたすべてのユニークで実行可能な経路と、ソリューションのみのベースライン（64個の応答）が見つけた経路をプロットしています。高速モードのDualformerは42個のユニークで実行可能な経路を特定しましたが、ソリューションのみのモデルはわずか3個しか見つけられませんでした。同様に、図C.2は低速モードのDualformerと完全軌跡（Searchformer）ベースラインを比較しています。低速モードのDualformerは39個のユニークで実行可能な経路を発見しましたが、完全軌跡モデルは17個しか見つけられませんでした。