MiMo-7B-RL: El Modelo de Lenguaje de Razonamiento de Xiaomi

El equipo LLM-Core de Xiaomi presenta MiMo-7B-RL, desafiando la idea de que el razonamiento de primer nivel en IA requiere modelos masivos. Este modelo de 7 mil millones de parámetros, diseñado específicamente para tareas matemáticas y de codificación, demuestra un rendimiento que rivaliza con modelos mucho más grandes y sistemas especializados como o1-mini de OpenAI. Este logro es el resultado de una estrategia integral que optimiza todo el ciclo de vida del modelo, demostrando que se puede desbloquear un razonamiento potente en arquitecturas más eficientes.

Qué es MiMo-7B

El desarrollo de MiMo-7B se basa en la creencia de que la capacidad de razonamiento fundamental de un modelo se establece durante el preentrenamiento. Si bien las etapas posteriores de ajuste fino son esenciales, la base inicial es crítica. Xiaomi identificó que muchos modelos más pequeños tienen dificultades con el razonamiento complejo porque su entrenamiento base carece de suficiente exposición a patrones lógicos.

Para contrarrestar esto, el preentrenamiento de MiMo fue diseñado meticulosamente para maximizar la "densidad de patrones de razonamiento". Esto implicó un procesamiento de datos sofisticado: mejorar la extracción de texto para capturar estructuras complejas en documentos técnicos y código, aplicar filtros multidimensionales para concentrar ejemplos de razonamiento y generar vastos conjuntos de datos sintéticos que incorporen pasos lógicos y resolución de problemas. Se empleó una estrategia de mezcla de datos en tres etapas durante el preentrenamiento, utilizando aproximadamente 25 billones de tokens para construir el modelo MiMo-7B-Base.

Además, Xiaomi incorporó la Predicción de Múltiples Tokens (MTP) como un objetivo de entrenamiento auxiliar. Esta técnica, en la que el modelo predice varios tokens por adelantado, mejora potencialmente la comprensión de dependencias complejas y puede acelerar la inferencia mediante decodificación especulativa.

Aprendizaje por Refuerzo Avanzado

Basándose en el modelo MiMo-7B-SFT ajustado, la fase de Aprendizaje por Refuerzo (RL) se enfoca específicamente en la competencia matemática y de codificación. Un conjunto de datos de alta calidad de 130.000 problemas de matemáticas y código cuidadosamente seleccionados, todos verificables mediante comprobaciones basadas en reglas (como pruebas unitarias o validación numérica), constituyó la base para el entrenamiento.

Para garantizar una mejora genuina de la capacidad y evitar el "hackeo de recompensas", solo se utilizaron recompensas de precisión objetivas y basadas en reglas. Se introdujo un novedoso sistema de "recompensa de código impulsada por la dificultad de la prueba" para abordar el problema de la recompensa dispersa inherente a la generación de código complejo. En lugar de una recompensa de todo o nada, este sistema otorga crédito parcial por pasar casos de prueba más fáciles dentro de un problema, proporcionando una señal de gradiente más densa de la que el modelo puede aprender.

La eficiencia también fue clave. A medida que el modelo mejoraba, una estrategia de remuestreo de datos redujo el peso de los problemas más fáciles, enfocando el entrenamiento en ejemplos más desafiantes. Xiaomi también desarrolló un "Motor de Despliegue Continuo" (Seamless Rollout Engine), una infraestructura de RL optimizada que integra generación continua, cálculo de recompensas asíncrono y terminación temprana para minimizar el tiempo de inactividad de la GPU, lo que resultó en aceleraciones significativas en el entrenamiento (2.29x) y la validación (1.96x).

Familia MiMo-7B-RL: Un Vistazo Rápido

Xiaomi ha lanzado varios modelos que muestran las etapas de desarrollo:

Benchmarks de MiMo-7B-RL

Los resultados de la evaluación destacan las fortalezas de MiMo-7B-RL, particularmente cuando se compara con modelos líderes utilizando una temperatura de generación de 0.6.

Rendimiento Comparativo:

(Se muestran benchmarks seleccionados de matemáticas/código)

MiMo-7B-RL demuestra un rendimiento excepcional en matemáticas y codificación, a menudo superando a modelos significativamente más grandes y modelos de razonamiento especializados como o1-mini en benchmarks desafiantes como MATH, AIME y versiones recientes de LiveCodeBench. Si bien sus puntuaciones de razonamiento general son sólidas para su tamaño, naturalmente se quedan por detrás de los modelos de frontera más grandes, lo que refleja su enfoque de entrenamiento especializado.

Rendimiento Dentro de la Serie MiMo:

Esta comparación interna ilustra la efectividad de cada etapa de entrenamiento. El modelo base muestra un fuerte razonamiento inicial, que se potencia significativamente mediante SFT, y se refina aún más hasta alcanzar el máximo rendimiento mediante la fase final de RL dirigida a matemáticas y código. Aplicar RL directamente a la base (RL-Zero) es efectivo, pero el paso intermedio de SFT parece beneficioso para lograr las puntuaciones más altas.

Ejecutando MiMo-7B-RL

Los modelos están fácilmente disponibles en el Hugging Face Hub.

Acceso al Modelo:

Encuentra MiMo-7B-RL y otros modelos de la serie en la página de la organización XiaomiMiMo en Hugging Face. El tamaño del modelo es de aproximadamente 7.83 mil millones de parámetros (precisión BF16, formato Safetensors).

Ejecutando Inferencia con vLLM (Recomendado)

Xiaomi recomienda usar su fork de vLLM (basado en v0.7.3) para la inferencia, ya que soporta la característica de Predicción de Múltiples Tokens para una generación potencialmente más rápida.

• Usando el Fork de vLLM de Xiaomi (con MTP):

# Asegúrate de que el fork de vLLM de Xiaomi esté instalado
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" # Reemplaza con tu ruta de descarga

llm = LLM(
    model=model_path,
    trust_remote_code=True,  # Requerido para el código personalizado de MiMo
    num_speculative_tokens=1, # Habilita la decodificación especulativa MTP
    disable_log_stats=False
)
# Temperatura de muestreo recomendada para replicar benchmarks
sampling_params = SamplingParams(temperature=0.6)

# Estructura de conversación de ejemplo (se recomienda un prompt de sistema vacío)
conversation = [
    {
        "role": "system",
        "content": "" # Usa un prompt de sistema vacío
    },
    {
        "role": "user",
        "content": "Write a python function to compute the nth Fibonacci number.",
    },
]

# Genera la respuesta
outputs = llm.chat(conversation,
                   sampling_params=sampling_params,
                   use_tqdm=False)

# Procesa e imprime la salida
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)

• Usando vLLM Estándar (sin MTP):
  Si no usas la característica MTP o usas una compilación estándar de vLLM, registra primero la arquitectura MiMo usando el script register_mimo_in_vllm.py proporcionado por Xiaomi.

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
# Asegúrate de que register_mimo_in_vllm.py sea accesible
import register_mimo_in_vllm

from vllm import LLM, SamplingParams

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Reemplaza con tu ruta de descarga
llm = LLM(
    model=model_path,
    trust_remote_code=True,
    # No establezcas num_speculative_tokens si no usas MTP
    disable_log_stats=False
)
sampling_params = SamplingParams(temperature=0.6)

# La configuración de la conversación y la llamada de generación es la misma que en el ejemplo MTP...
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)
# El procesamiento de la salida es el mismo...
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 ---

Usando HuggingFace Transformers

También es posible realizar inferencia con la librería estándar transformers de HuggingFace. Recuerda que trust_remote_code=True es necesario.

# --- 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" # Reemplaza con tu ruta de descarga

# Carga el modelo y el tokenizador
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True, # Esencial para cargar MiMo
    device_map="auto"       # Usa GPU si está disponible
)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Prepara el prompt de entrada
prompt = "Write a python function to compute the nth Fibonacci number."
# Tokeniza la entrada
inputs = tokenizer([prompt], return_tensors='pt').to(model.device)

# Genera la secuencia de salida
output_sequences = model.generate(
    **inputs,
    max_new_tokens=256,      # Controla la longitud de la salida
    temperature=0.6,         # Temperatura recomendada
    do_sample=True           # Usa muestreo para temperaturas != 1.0
)

# Decodifica la salida
generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
print(generated_text)
# --- FACTUAL CODE SNIPPET END ---

Recomendaciones de Uso

Para obtener mejores resultados, especialmente al intentar replicar las puntuaciones de los benchmarks, utiliza la configuración recomendada: el fork de vLLM de Xiaomi (basado en v0.7.3) y un prompt de sistema vacío.

Reflexiones Finales: ¿Razonamiento Eficiente Realizado por Xiaomi?

MiMo-7B-RL de Xiaomi demuestra que se puede lograr un rendimiento de razonamiento excepcional en dominios especializados como las matemáticas y la codificación sin recurrir a tamaños de modelo enormes. Mediante un preentrenamiento cuidadoso centrado en patrones de razonamiento y técnicas innovadoras de aprendizaje por refuerzo, han creado un modelo eficiente que compite eficazmente con contrapartes mucho más grandes. La publicación abierta de la serie MiMo proporciona herramientas e información valiosas, impulsando el desarrollo de capacidades de razonamiento de IA potentes y accesibles.