Xiaomiの推論LLM、MiMo-7B-RL

XiaomiのLLM-Coreチームは、AIにおける最高レベルの推論には巨大なモデルが必要であるという考え方に挑戦するMiMo-7B-RLを発表します。この70億パラメータのモデルは、数学およびコーディングタスクのために特別に設計されており、はるかに大きなモデルやOpenAIのo1-miniのような専門システムに匹敵するパフォーマンスを発揮します。この成果は、モデルライフサイクル全体を最適化する包括的な戦略の結果であり、より効率的なアーキテクチャでも強力な推論を引き出せることを証明しています。

MiMo-7Bとは

MiMo-7Bの開発は、モデルの基本的な推論能力は事前学習中に確立されるという信念に基づいています。後のファインチューニング段階は不可欠ですが、初期の基盤が重要です。Xiaomiは、多くの小さなモデルが複雑な推論に苦労するのは、そのベーストレーニングに論理的なパターンへの十分な露出が不足しているためであると特定しました。

これに対抗するため、MiMoの事前学習は「推論パターンの密度」を最大化するために綿密に設計されました。これには高度なデータ処理が含まれます：技術文書やコードの複雑な構造を捉えるためのテキスト抽出の強化、推論例を集中させるための多次元フィルターの適用、そして論理的なステップと問題解決を体現する膨大な合成データセットの生成です。MiMo-7B-Baseモデルを構築するために、事前学習中に約25兆トークンを使用した3段階のデータ混合戦略が採用されました。

さらに、Xiaomiは補助的な学習目的としてMultiple-Token Prediction (MTP) を組み込みました。この技術は、モデルが数トークン先を予測することで、複雑な依存関係の理解を深め、投機的デコーディングによる推論の高速化を可能にします。

高度な強化学習

ファインチューニングされたMiMo-7B-SFTモデルを基盤として、強化学習（RL）フェーズは、数学とコーディングの習熟度を特にターゲットとしています。ルールベースのチェック（単体テストや数値検証など）で検証可能な、厳選された13万件の高品質な数学およびコーディング問題のデータセットが、学習の基礎となりました。

真の能力向上を保証し、「報酬ハッキング」を回避するために、客観的でルールベースの精度報酬のみが使用されました。複雑なコード生成に内在する疎な報酬問題を解決するために、斬新な「テスト難易度駆動型コード報酬」システムが導入されました。これは、全か無かの報酬ではなく、問題内の簡単なテストケースに合格した場合に部分的な報酬を与えることで、モデルが学習するためのより密な勾配シグナルを提供します。

効率性も重要でした。モデルが向上するにつれて、データ再サンプリング戦略により簡単な問題の重みを下げ、より難しい例に学習を集中させました。Xiaomiはまた、「シームレスロールアウトエンジン」という最適化されたRLインフラストラクチャを開発しました。これは、連続生成、非同期報酬計算、早期終了を統合し、GPUアイドル時間を最小限に抑えることで、学習（2.29倍）および検証（1.96倍）の大幅な高速化を実現しました。

MiMo-7B-RLファミリー：概要

Xiaomiは、開発段階を示すいくつかのモデルをリリースしています。

MiMo-7B-RLベンチマーク

評価結果はMiMo-7B-RLの強みを浮き彫りにしています。特に生成温度0.6を使用して主要モデルと比較した場合に顕著です。

比較パフォーマンス：

（選択された数学/コードベンチマークを示す）

MiMo-7B-RLは、数学およびコーディングにおいて卓越したパフォーマンスを発揮し、MATH、AIME、および最近のLiveCodeBenchバージョンなどの難しいベンチマークで、しばしばはるかに大きなモデルやo1-miniのような専門的な推論モデルを超えています。その一般的な推論スコアはサイズに対して強力ですが、その専門的な学習の焦点が反映され、最大の最先端モデルには自然に劣ります。

MiMoシリーズ内のパフォーマンス：

この内部比較は、各学習段階の有効性を示しています。ベースモデルは強力な初期推論を示し、SFTによって大幅に強化され、数学とコードをターゲットとする最終的なRLフェーズによってピークパフォーマンスにさらに洗練されます。ベースモデルに直接RLを適用する（RL-Zero）ことも有効ですが、SFTの中間ステップが最高スコアを達成する上で有益であるようです。

MiMo-7B-RLの実行

モデルはHugging Face Hubで簡単に利用できます。

モデルへのアクセス：

MiMo-7B-RLおよびシリーズの他のモデルは、Hugging FaceのXiaomiMiMo組織ページで入手できます。モデルサイズは約78.3億パラメータです（BF16精度、Safetensors形式）。

vLLMを使用した推論の実行（推奨）

Xiaomiは、Multi-Token Prediction機能をサポートしており、生成速度が向上する可能性があるため、推論にはvLLMのフォーク（v0.7.3ベース）を使用することを推奨しています。

• Xiaomi vLLMフォークの使用（MTPあり）：

# Ensure Xiaomi's vLLM fork is installed
from vllm import LLM, SamplingParams

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Replace with your download path

llm = LLM(
    model=model_path,
    trust_remote_code=True,  # Required for MiMo's custom code
    num_speculative_tokens=1, # Enables MTP speculative decoding
    disable_log_stats=False
)
# Recommended sampling temperature for benchmark replication
sampling_params = SamplingParams(temperature=0.6)

# Example conversation structure (empty system prompt recommended)
conversation = [
    {
        "role": "system",
        "content": "" # Use an empty system prompt
    },
    {
        "role": "user",
        "content": "Write a python function to compute the nth Fibonacci number.",
    },
]

# Generate the response
outputs = llm.chat(conversation,
                   sampling_params=sampling_params,
                   use_tqdm=False)

# Process and print output
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}")
    print("-" * 20)
    print(f"Generated text: {generated_text!r}")
# --- FACTUAL CODE SNIPPET END ---

print("=" * 80)

• 標準vLLMの使用（MTPなし）：
  MTP機能を使用しない場合、または標準のvLLMビルドを使用する場合は、まずXiaomiが提供するregister_mimo_in_vllm.pyスクリプトを使用してMiMoアーキテクチャを登録してください。

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
# Ensure register_mimo_in_vllm.py is accessible
import register_mimo_in_vllm

from vllm import LLM, SamplingParams

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Replace with your download path
llm = LLM(
    model=model_path,
    trust_remote_code=True,
    # Do not set num_speculative_tokens if not using MTP
    disable_log_stats=False
)
sampling_params = SamplingParams(temperature=0.6)

# Conversation setup and generation call is the same as the MTP example...
conversation = [
    {"role": "system", "content": ""},
    {"role": "user", "content": "Write a python function to compute the nth Fibonacci number."},
]
outputs = llm.chat(conversation, sampling_params=sampling_params, use_tqdm=False)
# Processing output is the same...
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}\n{'-'*20}\nGenerated text: {generated_text!r}")
# --- FACTUAL CODE SNIPPET END ---

HuggingFace Transformersの使用

標準のHuggingFace transformersライブラリを使用した推論も可能です。trust_remote_code=Trueが必要であることを忘れないでください。

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Replace with your download path

# Load the model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True, # Essential for loading MiMo
    device_map="auto"       # Use GPU if available
)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Prepare the input prompt
prompt = "Write a python function to compute the nth Fibonacci number."
# Tokenize the input
inputs = tokenizer([prompt], return_tensors='pt').to(model.device)

# Generate the output sequence
output_sequences = model.generate(
    **inputs,
    max_new_tokens=256,      # Control output length
    temperature=0.6,         # Recommended temperature
    do_sample=True           # Use sampling for temperatures != 1.0
)

# Decode the output
generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
print(generated_text)
# --- FACTUAL CODE SNIPPET END ---

使用上の推奨事項

最良の結果を得るには、特にベンチマークスコアを再現しようとする場合、推奨されるセットアップを使用してください：XiaomiのvLLMフォーク（v0.7.3ベース）と空のシステムプロンプト。

最終的な考察：Xiaomiによる効率的な推論の実現か？

XiaomiのMiMo-7B-RLは、数学やコーディングのような専門分野における卓越した推論パフォーマンスが、巨大なモデルサイズに頼ることなく達成可能であることを示しています。推論パターンに焦点を当てた綿密な事前学習と革新的な強化学習技術を通じて、彼らははるかに大きなモデルと効果的に競合する効率的なモデルを作成しました。MiMoシリーズのオープンリリースは、貴重なツールと洞察を提供し、強力でアクセスしやすいAI推論能力の開発を推進します。