MiMo-7B-RL: LLM Penalaran dari Xiaomi

Tim LLM-Core Xiaomi mempersembahkan MiMo-7B-RL, menantang gagasan bahwa penalaran tingkat atas dalam AI membutuhkan model besar. Model 7 miliar parameter ini, yang secara khusus direkayasa untuk tugas matematika dan pengkodean, menunjukkan kinerja yang menyaingi model yang jauh lebih besar dan sistem khusus seperti o1-mini milik OpenAI. Pencapaian ini merupakan hasil dari strategi komprehensif yang mengoptimalkan seluruh siklus hidup model, membuktikan bahwa penalaran yang kuat dapat dibuka dalam arsitektur yang lebih efisien.

Apa itu MiMo-7B

Pengembangan MiMo-7B bergantung pada keyakinan bahwa kemampuan penalaran fundamental model ditetapkan selama pra-pelatihan. Meskipun tahap fine-tuning selanjutnya penting, fondasi awal sangat krusial. Xiaomi mengidentifikasi bahwa banyak model yang lebih kecil kesulitan dengan penalaran kompleks karena pelatihan dasarnya kurang memiliki paparan yang cukup terhadap pola logis.

Untuk mengatasi hal ini, pra-pelatihan MiMo dirancang dengan cermat untuk memaksimalkan "kepadatan pola penalaran". Ini melibatkan pemrosesan data yang canggih: meningkatkan ekstraksi teks untuk menangkap struktur kompleks dalam dokumen teknis dan kode, menerapkan filter multi-dimensi untuk memusatkan contoh penalaran, dan menghasilkan dataset sintetis yang luas yang mewujudkan langkah-langkah logis dan pemecahan masalah. Strategi campuran data tiga tahap digunakan selama pra-pelatihan, memanfaatkan sekitar 25 triliun token untuk membangun model MiMo-7B-Base.

Selain itu, Xiaomi memasukkan Multiple-Token Prediction (MTP) sebagai tujuan pelatihan tambahan. Teknik ini, di mana model memprediksi beberapa token ke depan, berpotensi meningkatkan pemahaman ketergantungan kompleks dan dapat mempercepat inferensi melalui speculative decoding.

Pembelajaran Penguatan Tingkat Lanjut

Membangun di atas model MiMo-7B-SFT yang telah di-fine-tune, fase Pembelajaran Penguatan (RL) secara khusus menargetkan kemahiran matematika dan kode. Dataset berkualitas tinggi berisi 130.000 soal matematika dan kode yang dikurasi dengan cermat, semuanya dapat diverifikasi melalui pemeriksaan berbasis aturan (seperti unit test atau validasi numerik), menjadi dasar untuk pelatihan.

Untuk memastikan peningkatan kemampuan yang asli dan menghindari "reward hacking", hanya hadiah akurasi berbasis aturan yang objektif yang digunakan. Sistem "test difficulty driven code reward" yang baru diperkenalkan untuk mengatasi masalah sparse reward yang melekat dalam pembuatan kode kompleks. Alih-alih hadiah all-or-nothing, sistem ini memberikan kredit parsial untuk melewati kasus uji yang lebih mudah dalam suatu masalah, memberikan sinyal gradien yang lebih padat untuk dipelajari oleh model.

Efisiensi juga menjadi kunci. Seiring dengan peningkatan model, strategi re-sampling data mengurangi bobot masalah yang lebih mudah, memfokuskan pelatihan pada contoh yang lebih menantang. Xiaomi juga mengembangkan "Seamless Rollout Engine", infrastruktur RL yang dioptimalkan yang mengintegrasikan pembuatan berkelanjutan, perhitungan hadiah asinkron, dan penghentian dini untuk meminimalkan waktu idle GPU, menghasilkan percepatan pelatihan (2.29x) dan validasi (1.96x) yang signifikan.

Keluarga MiMo-7B-RL: Sekilas

Xiaomi telah merilis beberapa model yang menunjukkan tahap pengembangan:

Benchmark MiMo-7B-RL

Hasil evaluasi menyoroti kekuatan MiMo-7B-RL, terutama jika dibandingkan dengan model-model terkemuka menggunakan suhu generasi 0.6.

Kinerja Komparatif:

(Benchmark matematika/kode terpilih ditampilkan)

MiMo-7B-RL menunjukkan kinerja luar biasa dalam matematika dan pengkodean, seringkali melebihi model yang jauh lebih besar dan model penalaran khusus seperti o1-mini pada benchmark yang menantang seperti MATH, AIME, dan versi LiveCodeBench terbaru. Meskipun skor penalaran umumnya kuat untuk ukurannya, secara alami tertinggal dari model frontier terbesar, mencerminkan fokus pelatihan khususnya.

Kinerja dalam Seri MiMo:

Perbandingan internal ini menggambarkan efektivitas setiap tahap pelatihan. Model dasar menunjukkan penalaran awal yang kuat, yang secara signifikan ditingkatkan oleh SFT, dan selanjutnya disempurnakan hingga kinerja puncak oleh fase RL terakhir yang menargetkan matematika dan kode. Menerapkan RL langsung ke dasar (RL-Zero) efektif, tetapi langkah perantara SFT tampaknya bermanfaat untuk mencapai skor tertinggi.

Menjalankan MiMo-7B-RL

Model-model ini tersedia di Hugging Face Hub.

Akses Model:

Temukan MiMo-7B-RL dan model lain dalam seri ini di halaman organisasi XiaomiMiMo di Hugging Face. Ukuran model sekitar 7.83 miliar parameter (presisi BF16, format Safetensors).

Menjalankan Inferensi dengan vLLM (Direkomendasikan)

Xiaomi merekomendasikan penggunaan fork vLLM mereka (berdasarkan v0.7.3) untuk inferensi, karena mendukung fitur Multi-Token Prediction untuk potensi generasi yang lebih cepat.

• Menggunakan Xiaomi vLLM Fork (dengan 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)

• Menggunakan vLLM Standar (tanpa MTP):
  Jika tidak menggunakan fitur MTP atau menggunakan build vLLM standar, daftarkan arsitektur MiMo terlebih dahulu menggunakan skrip register_mimo_in_vllm.py yang disediakan oleh Xiaomi.

# --- 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 ---

Menggunakan HuggingFace Transformers

Inferensi library transformers standar HuggingFace juga dimungkinkan. Ingat trust_remote_code=True diperlukan.

# --- 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 ---

Rekomendasi Penggunaan

Untuk hasil terbaik, terutama saat mencoba mereplikasi skor benchmark, gunakan pengaturan yang direkomendasikan: fork vLLM Xiaomi (berdasarkan v0.7.3) dan prompt sistem kosong.

Pemikiran Akhir: Penalaran Efisien Diwujudkan oleh Xiaomi?

MiMo-7B-RL Xiaomi menunjukkan bahwa kinerja penalaran yang luar biasa dalam domain khusus seperti matematika dan pengkodean dapat dicapai tanpa menggunakan ukuran model yang sangat besar. Melalui pra-pelatihan yang cermat yang berfokus pada pola penalaran dan teknik pembelajaran penguatan yang inovatif, mereka telah menciptakan model yang efisien yang bersaing secara efektif dengan rekan-rekannya yang jauh lebih besar. Rilis terbuka seri MiMo menyediakan alat dan wawasan yang berharga, mendorong pengembangan kemampuan penalaran AI yang kuat dan mudah diakses.