LFM-1B-Math: 小型モデルは簡潔に推論できるか？

要点

• 本稿では、1.3BのLiquid Foundation Model（LFM）に基づいた小型チャットモデルに数学的推論能力を追加するための原則に基づく後処理トレーニング手法を示します。
• キュレーションされたChain-of-Thought（CoT）データによる広範な教師ありファインチューニングで高い推論能力を引き出し、さらにターゲットを絞った強化学習により推論チェーンを簡潔化します。
• 結果として得られたLFM-1.3B-Mathモデルは、応答制限下でも高い性能を発揮し、リソース制約のあるエッジデバイスへの展開に理想的なトレードオフを実現します。

推論モデルは、複雑な問題を解くためにChain-of-Thought（思考の連鎖）や強化学習（RL）に基づく新たなトレーニングパラダイムを普及させました。これにより、小型モデルでも従来は大型LLMが必要だったタスクを遂行できるようになっています。しかし、ほとんどの研究文献は限られたベースモデルに焦点を当てており、推論手法の有効性と、前処理データやアーキテクチャといったモデル固有の特性を切り分けることが困難です。

本ブログ記事では、小型言語モデルにおける推論手法の探求に基づく洞察を、LFM-1.3B [1] の例を通じて共有します。数学特化の事前学習を受けていない一般的なチャットモデルが、広範なファインチューニングと短期的強化学習を通じて簡潔な推論能力を獲得できる様子を示します。私たちのモデル、事前学習ミックス、アーキテクチャに関する知見を共有することで、エッジ展開に適した小型モデルへの推論能力の有効な埋め込み方に関する理解を広げることを目指します。

教師ありファインチューニング（SFT）

教師ありファインチューニング（SFT）は、キュレーションされたCoTトレース [2] に基づいて、伝統的なモデルに推論行動を付加するための一般的な手法です。強力な教師モデルからのオフライン蒸留スタイルのSFTは、小型推論モデルに特に有効であり、純粋な強化学習が困難な場合にも有効です。

命令データミックス

近年の研究では、わずか1,000サンプルでも効果的なトレーニングが可能とされていますが [3]、我々の調査では100万サンプル以上の大規模オープンソースデータセットを活用することで特に高い性能が得られました。これは、モデルサイズや数学・推論特化の事前学習データとの相関があると推測されます（[4]も参照）。

この観察は他の強力な小型推論モデルの実践とも一致します。例えば、DeepSeek-R1-Distill-Qwen-1.5B [2] は約100万サンプルで学習され、NvidiaのOpenMath-Nemotron-1.5B [5] はこれにさらに300万サンプルを追加で学習しています。

我々はNvidiaのOpenMathReasoning [5] やLlama Nemotron Post-Training [6] のデータセットを含むいくつかのオープンデータセット（Hugging FaceのOpen-R1 [7] やBespoke LabのOpenThoughts（例：OpenThoughts2-1M）[8] を含む）を評価しました。OpenMathReasoningは数学問題において特に有効であると判断し、これにLlama Nemotron Post-Trainingの科学サブセットとSCP-116K [9] を加えることで、科学問題への応用も可能にしました。これにより、最終的なサンプル数は約450万件となりました。

ハイパーパラメータ

これらの大規模かつ高品質なデータセットの利点を最大限に活かすために、通常のポストトレーニングレシピから大きく逸脱しました。具体的には、比較的高い学習率（3e-4）を使用し、合計3エポック（約1000億トークン）で学習しました。これらの積極的なハイパーパラメータの選択は、圧縮や過学習が困難なこれらの大規模データセットから最大限の知識を引き出すために不可欠でした。図1（左）は、マルチエポック学習の累積的な効果を示しており、初回のエポックで平均性能が36%向上し、残り2エポックでさらに7.6%向上しています。

図1. 左：複数エポックにわたる一般的な推論ベンチマークにおける平均性能の向上。右：応答長を短縮するための切り捨てられた推論トレースでのトレーニング。切り捨てによって応答長は大幅に短縮されるものの、32kトークン予算での性能は低下し、制約付きの4kトークン予算でも顕著な改善は見られない。

応答の冗長性

推論は、モデルが複数の解法戦略を順次試し、間違いを修正し最終的な解答を自己検証することを可能にします。これにより解答精度は大きく向上しますが、応答は自然と長くなり、目標に完全には沿わない冗長性が生じることがあります。これは、メモリや生成速度が制約となる多くのエッジ用途とは相容れません。

我々は、SFTのみで応答を短くするいくつかの方法を試しました：

• 上位20%の最長サンプルを除去（全トークンの約50%をフィルタリング）
• システムプロンプトによる後処理条件付け（例："thinking: {low; medium; high}"）
• 応答全体が3kトークン以内に収まるよう推論トレースを切り捨て

後者については、4kトークンの応答予算下で評価した際、切り捨てられた推論トレースの学習でも性能向上につながるか検証しました。図1（右）が示すように、応答長の短縮には成功しましたが、性能向上は見られず、32kトークンの非制限評価では性能が低下しました。したがって、元のモデルの生成長を直接制約する方が効果的であると判断しました。

強化学習

SFT（教師ありファインチューニング）だけでは、性能を大きく損なわずに推論を短縮するには不十分であることがわかりました。そこで私たちは、RL（強化学習）を活用してモデルを簡潔な推論行動へと洗練させました。本手法は、Group Relative Policy Optimization（GRPO）[10] に基づいています。これは、モンテカルロ法によるアドバンテージ推定と検証可能な成果報酬の組み合わせにより、単純な最適化目標を提供するためです。さらに、学習の安定性を高めるために、アドバンテージの正規化を除去し、過長応答のフィルタリングを適用しました [11,12]。これらの変更のもとで、GRPO、RLOO [13]、REINFORCE [14] の類似性が見られたため、一貫性の観点からGRPOを採用しています。

RLによる冗長性制御

モデルの冗長性を削減する最初のステップとしてよく使われるのが、あらかじめ定義された長さのヒューリスティックに基づいた専用の報酬項を導入することです。しかし、これは問題の複雑さが多様である場合、推論の予算を柔軟に調整する能力を妨げる可能性があります。特に小型モデルはゼロショット能力が限られており、難しい問題を解くためには複数の試行錯誤が必要となる場合があります。それでも、応答を簡潔に正則化することは依然として重要です。そこで私たちは、GRPO-LEAD [15] に基づく難易度認識型の報酬定式化とアドバンテージの再重み付けを採用し、Chain-of-Thought（CoT）の長さと問題の難易度の整合性を高めました。

以下では、状態をs、行動をa、報酬をrとし、各プロンプトに対してG個の応答をサンプリングし、個々の生成aiをトークンレベルの行動ai,ₜから構成されるものと定義します。成果報酬は、長さを考慮した指数的スケーリングによって定義されます。

ここで、sg() はストップグラディエント演算子を表し、アドバンテージAiは以下に基づいて重み付けされます：

ここで、μGとσGは、正解応答の長さのグループ平均と標準偏差を表します。

目的関数は、[15] のアドバンテージ再重み付け定式に従い、KL正則化を伴わない単純なポリシー勾配設定内で構成されます。

ここで、sg() はストップグラディエント演算子を表し、アドバンテージAiは以下に基づいて重み付けされます：

ここで、ρGは、サンプルiに関連付けられたグループGの正解応答率を表します。

過長応答を大規模にフィルタリングするのを避けるため、Elastic Reasoning [16] で採用されているように、最終応答の生成前に明示的な推論予算を設定しました。切り捨てられていない推論トレースに対して成功裏にファインチューニングを行った後、RL段階では切り捨てられた推論トレースが最終応答の品質向上に十分な信号を含んでいることが多いとわかりました。生成されるサンプルは、プロンプト、推論トレース、最終回答から構成され、

ここで、推論トレースは部分的であり、最終回答は終了していない可能性があります。

RL用データミックス

本研究で使用したデータセットは、いくつかの一般的なオープンソースデータセットのサンプルを組み合わせたものです。全体的に、データの正確な混合比は、問題の難易度バランスを除いて、性能への強い影響を示しませんでした。GRPOの学習信号を十分に確保するために、SFTチェックポイントに基づいて経験的な解答率を評価し、最終的なミックスには、4096トークンの生成予算下で解答率が20～80%のサンプルのみを含めました。さらに、大規模データセットからランダムにサブサンプリングして、合計11,373件のサンプルミックスを構成しました（以下の分布に従う）：

ハイパーパラメータ

128プロンプトのバッチをサンプリングし、それぞれに対して温度1.0で13個の応答を生成します。文脈長は4096トークン、最終応答の予算は1024トークンとしています。オプティマイザーにはAdamWを用い、学習率は一定の5e-6、KL正則化は使用していません。GRPO-LEADのパラメータは以下の通りです：α = 0.05、k1 = 0.4、k2 = 1.1、k3 = 10、ρ0 = 0.75。

評価

我々の推論チェックポイントの性能を評価するために、GPQA Diamond、MATH500、AIME24、AIME25、AMC23といった標準的な推論ベンチマークを用いました。GPQA Diamondは生物学、物理学、化学に関する多様な質問を含みますが、その他のベンチマークは数学に特化しています。我々は、32kトークンのフル出力予算と、エッジ制約のある4kトークン設定という2つの異なる応答予算を考慮しました。

表3では、32kトークン出力予算下におけるLFM-1.3Bのチェックポイントと他の小型推論モデルの性能を示しています。DeepSeek-R1-Distill-Qwen-1.5Bは、数学データで事前学習されたQwen2.5-Math [20] に基づく長いCoTを使ってSFTされたモデルであることに注意が必要です。LFM-1.3B-Mathは、数学特化の事前学習がないにもかかわらず、これらの小型推論モデルと競合できる性能を示しています。特に、長いCoTでのSFTにより、LFM-1.3B-Distillの平均性能は14.08%から59.87%に大幅に向上し、困難なAIMEサブセットを含むすべてのベンチマークで競争力のある性能を発揮しました。

小型推論モデルは、表示や計算コストの制約（例：スマートフォンでのスクロール）によって長すぎる応答がユーザー体験を損なうエッジアプリケーションに特に適しています。理想的なモデルは、問題の難易度に応じて応答長を適切に調整しつつ、低い平均応答長でフルトークン予算下でも高い性能を発揮するものです。

そこで我々は、表4のような制約付き設定（応答長が4kトークンに制限される）を検討しました。ここでは、多くのモデルで性能が大きく低下し、長いCoT推論への依存が顕著であることが示されました。この挙動を緩和するために、応答長を制限した強化学習を活用してLFM-1.3B-Mathを作成しました。結果として得られたモデルは、制約付き設定下での性能を30.9%から46.98%へと大幅に向上させ（表4）、非制約下での平均性能（表3）の低下はごくわずかでした。さらに、図2が示すように、平均応答長の低下によって冗長性の大幅な削減も確認されました。LFM-1.3B-Mathは、複雑な科学および数学の問題を解く能力を備えつつ、問題の複雑さに応じて応答長を適切に調整することができる強力なエッジ向けモデルとして位置付けられます。

図2. 正解および不正解の応答における平均応答長（各モデル間での比較）。左側はCoT対応のベースモデル、右側はファインチューニング済みモデルであり、ファインチューニング済みモデルの平均性能は上部に（フル予算｜制約付き予算）の形式で表示されています。

結論

LFM-1.3B-Mathは、小型言語モデルであっても数学特化の事前学習に依存せずに、競争力のある推論性能を達成できることを示しています。大規模な教師ありファインチューニングとターゲットを絞った強化学習を慎重に適用することで、一般的な1.3Bパラメータのモデルを、専門モデルと競合可能な優れた推論モデルへと変換することに成功しました。ここでは、キュレーションされたChain-of-Thoughtデータによる広範なSFTが高い推論能力を引き出し、短期的なRLが推論チェーンの簡潔さを維持する役割を果たしています。

その結果として得られた、積極的なトレーニングハイパーパラメータ、難易度に応じた報酬設計、明示的な推論予算の組み合わせは、限られた応答予算でも高い性能を可能にする上で不可欠であり、リソース制約のあるエッジデバイスへの展開に理想的なトレードオフとなります。

私たちは、大型モデルで観察される推論能力は、単にスケールやドメイン特化の学習データに依存するものではなく、工夫されたトレーニング手法によって効果的に引き出せると考えています。