MiMo-7B-RL : le LLM de raisonnement de Xiaomi

L'équipe LLM-Core de Xiaomi présente MiMo-7B-RL, remettant en question l'idée que le raisonnement de premier ordre en IA nécessite des modèles massifs. Ce modèle de 7 milliards de paramètres, spécialement conçu pour les tâches mathématiques et de codage, démontre des performances qui rivalisent avec des modèles beaucoup plus grands et des systèmes spécialisés comme o1-mini d'OpenAI. Cette réussite résulte d'une stratégie globale optimisant l'ensemble du cycle de vie du modèle, prouvant qu'un raisonnement puissant peut être débloqué dans des architectures plus efficaces.

Qu'est-ce que MiMo-7B ?

Le développement de MiMo-7B repose sur la conviction que la capacité de raisonnement fondamentale d'un modèle est établie lors du pré-entraînement. Bien que les étapes ultérieures de fine-tuning soient essentielles, la base initiale est cruciale. Xiaomi a constaté que de nombreux modèles plus petits ont du mal avec le raisonnement complexe parce que leur entraînement de base manque d'exposition suffisante aux schémas logiques.

Pour contrer cela, le pré-entraînement de MiMo a été méticuleusement conçu pour maximiser la « densité des schémas de raisonnement ». Cela impliquait un traitement sophistiqué des données : améliorer l'extraction de texte pour capturer des structures complexes dans les documents techniques et le code, appliquer des filtres multidimensionnels pour concentrer les exemples de raisonnement et générer de vastes ensembles de données synthétiques incarnant des étapes logiques et la résolution de problèmes. Une stratégie de mélange de données en trois étapes a été employée pendant le pré-entraînement, utilisant environ 25 000 milliards de jetons pour construire le modèle MiMo-7B-Base.

De plus, Xiaomi a intégré la prédiction à jetons multiples (MTP) comme objectif d'entraînement auxiliaire. Cette technique, où le modèle prédit plusieurs jetons à l'avance, améliore potentiellement la compréhension des dépendances complexes et peut accélérer l'inférence grâce au décodage spéculatif.

Apprentissage par renforcement avancé

S'appuyant sur le modèle MiMo-7B-SFT affiné, la phase d'apprentissage par renforcement (RL) cible spécifiquement les compétences en mathématiques et en code. Un ensemble de données de haute qualité de 130 000 problèmes de mathématiques et de code soigneusement sélectionnés, tous vérifiables par des contrôles basés sur des règles (comme des tests unitaires ou une validation numérique), a constitué la base de l'entraînement.

Pour garantir une véritable amélioration des capacités et éviter le « piratage des récompenses », seules des récompenses de précision objectives et basées sur des règles ont été utilisées. Un nouveau système de « récompense de code basée sur la difficulté des tests » a été introduit pour s'attaquer au problème de la récompense éparse inhérent à la génération de code complexe. Au lieu d'une récompense tout ou rien, ce système accorde un crédit partiel pour la réussite de tests plus faciles dans un problème, fournissant un signal de gradient plus dense pour que le modèle puisse en tirer des leçons.

L'efficacité était également essentielle. Au fur et à mesure que le modèle s'améliorait, une stratégie de rééchantillonnage des données a réduit le poids des problèmes plus faciles, en se concentrant sur l'entraînement sur des exemples plus difficiles. Xiaomi a également développé un « moteur de déploiement transparent », une infrastructure RL optimisée qui intègre la génération continue, le calcul asynchrone des récompenses et l'arrêt anticipé pour minimiser le temps d'inactivité du GPU, ce qui a permis d'accélérer considérablement l'entraînement (2,29x) et la validation (1,96x).

Famille MiMo-7B-RL : un aperçu rapide

Xiaomi a publié plusieurs modèles présentant les étapes de développement :

Benchmarks MiMo-7B-RL

Les résultats de l'évaluation mettent en évidence les points forts de MiMo-7B-RL, en particulier par rapport aux principaux modèles utilisant une température de génération de 0,6.

Performances comparatives :

(Benchmarks mathématiques/code sélectionnés affichés)

MiMo-7B-RL démontre des performances exceptionnelles en mathématiques et en codage, dépassant souvent les modèles significativement plus grands et les modèles de raisonnement spécialisés comme o1-mini sur des benchmarks difficiles comme MATH, AIME et les versions récentes de LiveCodeBench. Bien que ses scores de raisonnement général soient solides pour sa taille, ils traînent naturellement les plus grands modèles de pointe, reflétant son objectif d'entraînement spécialisé.

Performances au sein de la série MiMo :

Cette comparaison interne illustre l'efficacité de chaque étape d'entraînement. Le modèle de base montre un fort raisonnement initial, qui est considérablement renforcé par SFT, et encore affiné pour des performances maximales par la phase RL finale ciblant les mathématiques et le code. L'application de RL directement à la base (RL-Zero) est efficace, mais l'étape intermédiaire SFT semble bénéfique pour obtenir les scores les plus élevés.

Exécution de MiMo-7B-RL

Les modèles sont facilement disponibles sur le Hugging Face Hub.

Accès au modèle :

Trouvez MiMo-7B-RL et d'autres modèles de la série sur la page de l'organisation XiaomiMiMo sur Hugging Face. La taille du modèle est d'environ 7,83 milliards de paramètres (précision BF16, format Safetensors).

Exécution de l'inférence avec vLLM (recommandé)

Xiaomi recommande d'utiliser sa fork de vLLM (basée sur v0.7.3) pour l'inférence, car elle prend en charge la fonctionnalité de prédiction à jetons multiples pour une génération potentiellement plus rapide.

• Utilisation de la fork vLLM de Xiaomi (avec MTP) :

# Assurez-vous que la fork vLLM de Xiaomi est installée
from vllm import LLM, SamplingParams

# --- EXTRAIT DE CODE FACTUEL DÉBUT ---
# Source : https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Remplacez par votre chemin de téléchargement

llm = LLM(
    model=model_path,
    trust_remote_code=True,  # Requis pour le code personnalisé de MiMo
    num_speculative_tokens=1, # Active le décodage spéculatif MTP
    disable_log_stats=False
)
# Température d'échantillonnage recommandée pour la réplication des benchmarks
sampling_params = SamplingParams(temperature=0.6)

# Structure de conversation d'exemple (invite système vide recommandée)
conversation = [
    {
        "role": "system",
        "content": "" # Utilisez une invite système vide
    },
    {
        "role": "user",
        "content": "Écrivez une fonction python pour calculer le nième nombre de Fibonacci.",
    },
]

# Générer la réponse
outputs = llm.chat(conversation,
                   sampling_params=sampling_params,
                   use_tqdm=False)

# Traiter et imprimer la sortie
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}")
# --- EXTRAIT DE CODE FACTUEL FIN ---

print("=" * 80)

• Utilisation de vLLM standard (sans MTP) :
  Si vous n'utilisez pas la fonctionnalité MTP ou si vous utilisez une version vLLM standard, enregistrez d'abord l'architecture MiMo à l'aide du script register_mimo_in_vllm.py fourni par Xiaomi.

# --- EXTRAIT DE CODE FACTUEL DÉBUT ---
# Source : https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
# Assurez-vous que register_mimo_in_vllm.py est accessible
import register_mimo_in_vllm

from vllm import LLM, SamplingParams

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Remplacez par votre chemin de téléchargement
llm = LLM(
    model=model_path,
    trust_remote_code=True,
    # Ne définissez pas num_speculative_tokens si vous n'utilisez pas MTP
    disable_log_stats=False
)
sampling_params = SamplingParams(temperature=0.6)

# La configuration de la conversation et l'appel de génération sont les mêmes que dans l'exemple MTP...
conversation = [
    {"role": "system", "content": ""},
    {"role": "user", "content": "Écrivez une fonction python pour calculer le nième nombre de Fibonacci."},
]
outputs = llm.chat(conversation, sampling_params=sampling_params, use_tqdm=False)
# Le traitement de la sortie est le même...
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}")
# --- EXTRAIT DE CODE FACTUEL FIN ---

Utilisation de HuggingFace Transformers

L'inférence de la bibliothèque transformers standard de HuggingFace est également possible. N'oubliez pas que trust_remote_code=True est nécessaire.

# --- EXTRAIT DE CODE FACTUEL DÉBUT ---
# Source : https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Remplacez par votre chemin de téléchargement

# Charger le modèle et le tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True, # Essentiel pour charger MiMo
    device_map="auto"       # Utiliser le GPU si disponible
)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Préparer l'invite d'entrée
prompt = "Écrivez une fonction python pour calculer le nième nombre de Fibonacci."
# Tokeniser l'entrée
inputs = tokenizer([prompt], return_tensors='pt').to(model.device)

# Générer la séquence de sortie
output_sequences = model.generate(
    **inputs,
    max_new_tokens=256,      # Contrôler la longueur de la sortie
    temperature=0.6,         # Température recommandée
    do_sample=True           # Utiliser l'échantillonnage pour les températures != 1.0
)

# Décoder la sortie
generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
print(generated_text)
# --- EXTRAIT DE CODE FACTUEL FIN ---

Recommandations d'utilisation

Pour de meilleurs résultats, en particulier lorsque vous essayez de reproduire les scores de référence, utilisez la configuration recommandée : la fork vLLM de Xiaomi (basée sur v0.7.3) et une invite système vide.

Réflexions finales : un raisonnement efficace réalisé par Xiaomi ?

MiMo-7B-RL de Xiaomi démontre que des performances de raisonnement exceptionnelles dans des domaines spécialisés comme les mathématiques et le codage sont réalisables sans recourir à des tailles de modèles énormes. Grâce à un pré-entraînement minutieux axé sur les schémas de raisonnement et à des techniques d'apprentissage par renforcement innovantes, ils ont créé un modèle efficace qui rivalise efficacement avec des homologues beaucoup plus grands. La publication ouverte de la série MiMo fournit des outils et des informations précieux, faisant progresser le développement de capacités de raisonnement IA puissantes et accessibles.