MiMo-7B-RL: LLM de Raciocínio da Xiaomi

A Equipe LLM-Core da Xiaomi apresenta o MiMo-7B-RL, desafiando a ideia de que o raciocínio de alto nível em IA requer modelos massivos. Este modelo de 7 bilhões de parâmetros, projetado especificamente para tarefas matemáticas e de codificação, demonstra um desempenho que rivaliza com modelos muito maiores e sistemas especializados como o o1-mini da OpenAI. Essa conquista resulta de uma estratégia abrangente que otimiza todo o ciclo de vida do modelo, provando que um raciocínio potente pode ser desbloqueado em arquiteturas mais eficientes.

O que é o MiMo-7B

O desenvolvimento do MiMo-7B baseia-se na crença de que a capacidade fundamental de raciocínio de um modelo é estabelecida durante o pré-treinamento. Embora as etapas posteriores de ajuste fino (fine-tuning) sejam essenciais, a base inicial é crítica. A Xiaomi identificou que muitos modelos menores têm dificuldade com raciocínio complexo porque seu treinamento base carece de exposição suficiente a padrões lógicos.

Para contrariar isso, o pré-treinamento do MiMo foi meticulosamente projetado para maximizar a "densidade de padrões de raciocínio". Isso envolveu processamento de dados sofisticado: aprimoramento da extração de texto para capturar estruturas complexas em documentos técnicos e código, aplicação de filtros multidimensionais para concentrar exemplos de raciocínio e geração de vastos conjuntos de dados sintéticos que incorporam etapas lógicas e resolução de problemas. Uma estratégia de mistura de dados em três estágios foi empregada durante o pré-treinamento, utilizando aproximadamente 25 trilhões de tokens para construir o modelo MiMo-7B-Base.

Além disso, a Xiaomi incorporou a Previsão de Múltiplos Tokens (MTP) como um objetivo de treinamento auxiliar. Essa técnica, onde o modelo prevê vários tokens à frente, potencialmente aprimora a compreensão de dependências complexas e pode acelerar a inferência através de decodificação especulativa.

Aprendizado por Reforço Avançado

Com base no modelo MiMo-7B-SFT ajustado (fine-tuned), a fase de Aprendizado por Reforço (RL) visa especificamente a proficiência em matemática e código. Um conjunto de dados de alta qualidade com 130.000 problemas de matemática e código cuidadosamente selecionados (curated), todos verificáveis através de verificações baseadas em regras (como testes unitários ou validação numérica), formou a base para o treinamento.

Para garantir uma melhoria genuína da capacidade e evitar "reward hacking", apenas recompensas objetivas e baseadas em regras de precisão foram usadas. Um novo sistema de "recompensa de código impulsionada pela dificuldade do teste" foi introduzido para abordar o problema de recompensa esparsa inerente à geração de código complexo. Em vez de uma recompensa tudo-ou-nada, este sistema concede crédito parcial por passar em casos de teste mais fáceis dentro de um problema, fornecendo um sinal de gradiente mais denso para o modelo aprender.

A eficiência também foi fundamental. À medida que o modelo melhorava, uma estratégia de reamostragem de dados reduzia o peso de problemas mais fáceis, focando o treinamento em exemplos mais desafiadores. A Xiaomi também desenvolveu um "Seamless Rollout Engine", uma infraestrutura de RL otimizada que integra geração contínua, cálculo assíncrono de recompensa e terminação antecipada para minimizar o tempo ocioso da GPU, resultando em acelerações significativas no treinamento (2.29x) e validação (1.96x).

Família MiMo-7B-RL: Um Olhar Rápido

A Xiaomi lançou vários modelos mostrando as etapas de desenvolvimento:

Benchmarks do MiMo-7B-RL

Os resultados da avaliação destacam os pontos fortes do MiMo-7B-RL, particularmente quando comparado a modelos líderes usando uma temperatura de geração de 0.6.

Desempenho Comparativo:

(Benchmarks de matemática/código selecionados mostrados)

O MiMo-7B-RL demonstra desempenho excepcional em matemática e codificação, muitas vezes superando modelos significativamente maiores e modelos de raciocínio especializados como o o1-mini em benchmarks desafiadores como MATH, AIME e versões recentes do LiveCodeBench. Embora suas pontuações gerais de raciocínio sejam fortes para seu tamanho, elas naturalmente ficam atrás dos maiores modelos de fronteira, refletindo seu foco de treinamento especializado.

Desempenho Dentro da Série MiMo:

Esta comparação interna ilustra a eficácia de cada etapa de treinamento. O modelo base mostra um forte raciocínio inicial, que é significativamente impulsionado pelo SFT, e posteriormente refinado para desempenho máximo pela fase final de RL focada em matemática e código. Aplicar RL diretamente à base (RL-Zero) é eficaz, mas a etapa intermediária de SFT parece benéfica para alcançar as pontuações mais altas.

Executando o MiMo-7B-RL

Os modelos estão prontamente disponíveis no Hugging Face Hub.

Acesso ao Modelo:

Encontre o MiMo-7B-RL e outros modelos da série na página da organização XiaomiMiMo no Hugging Face. O tamanho do modelo é de aproximadamente 7.83 bilhões de parâmetros (precisão BF16, formato Safetensors).

Executando Inferência com vLLM (Recomendado)

A Xiaomi recomenda usar seu fork do vLLM (baseado na v0.7.3) para inferência, pois ele suporta o recurso Multi-Token Prediction para geração potencialmente mais rápida.

• Usando o Fork vLLM da Xiaomi (com 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)

• Usando o vLLM Padrão (sem MTP):
  Se não estiver usando o recurso MTP ou estiver usando uma compilação padrão do vLLM, registre a arquitetura MiMo primeiro usando o script register_mimo_in_vllm.py fornecido pela 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 ---

Usando HuggingFace Transformers

A inferência com a biblioteca padrão transformers do HuggingFace também é possível. Lembre-se que trust_remote_code=True é necessário.

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

Recomendações de Uso

Para melhores resultados, especialmente ao tentar replicar pontuações de benchmark, use a configuração recomendada: o fork vLLM da Xiaomi (baseado na v0.7.3) e um prompt de sistema vazio.

Considerações Finais: Raciocínio Eficiente Realizado pela Xiaomi?

O MiMo-7B-RL da Xiaomi demonstra que um desempenho excepcional de raciocínio em domínios especializados como matemática e codificação é alcançável sem recorrer a tamanhos de modelo enormes. Através de um pré-treinamento cuidadoso focado em padrões de raciocínio e técnicas inovadoras de aprendizado por reforço, eles criaram um modelo eficiente que compete efetivamente com contrapartes muito maiores. O lançamento aberto da série MiMo fornece ferramentas valiosas e insights, impulsionando o desenvolvimento de capacidades de raciocínio de IA poderosas e acessíveis.