샤오미 MiMo-7B-RL: 추론 LLM

샤오미의 LLM 코어 팀이 MiMo-7B-RL을 선보이며, AI에서 최고 수준의 추론 능력이 거대 모델에만 요구된다는 통념에 도전합니다. 수학 및 코딩 작업을 위해 특별히 설계된 이 70억 매개변수 모델은 OpenAI의 o1-mini와 같은 훨씬 더 큰 모델 및 전문 시스템과 경쟁하는 성능을 보여줍니다. 이러한 성과는 전체 모델 수명 주기를 최적화하는 포괄적인 전략의 결과이며, 더 효율적인 아키텍처에서도 강력한 추론 능력을 발휘할 수 있음을 입증합니다.

MiMo-7B란 무엇인가

MiMo-7B의 개발은 모델의 기본적인 추론 능력이 사전 학습 단계에서 확립된다는 믿음에 기반합니다. 후반의 미세 조정 단계도 중요하지만, 초기 기반이 핵심입니다. 샤오미는 많은 소형 모델이 복잡한 추론에 어려움을 겪는 이유가 기본 학습에서 논리적 패턴에 대한 충분한 노출이 부족하기 때문이라고 파악했습니다.

이를 해결하기 위해 MiMo의 사전 학습은 "추론 패턴 밀도"를 극대화하도록 세심하게 설계되었습니다. 여기에는 정교한 데이터 처리가 포함되었습니다: 기술 문서 및 코드에서 복잡한 구조를 포착하기 위한 텍스트 추출 강화, 추론 예시를 집중시키기 위한 다차원 필터 적용, 논리적 단계와 문제 해결을 구현하는 방대한 합성 데이터셋 생성. MiMo-7B-Base 모델 구축을 위해 사전 학습 중 약 25조 개의 토큰을 활용한 3단계 데이터 혼합 전략이 사용되었습니다.

또한 샤오미는 보조 학습 목표로 다중 토큰 예측(Multiple-Token Prediction, MTP)을 통합했습니다. 모델이 여러 토큰을 미리 예측하는 이 기법은 복잡한 종속성에 대한 이해를 향상시키고 추측적 디코딩을 통해 추론 속도를 높일 수 있습니다.

고급 강화 학습

미세 조정된 MiMo-7B-SFT 모델을 기반으로 하는 강화 학습(RL) 단계는 수학 및 코딩 능력을 특별히 목표로 합니다. 단위 테스트 또는 수치 검증과 같은 규칙 기반 검사를 통해 모두 검증 가능한 13만 개의 신중하게 선별된 고품질 수학 및 코딩 문제 데이터셋이 학습의 기반을 형성했습니다.

진정한 능력 향상을 보장하고 "보상 해킹"을 방지하기 위해 객관적이고 규칙 기반의 정확도 보상만 사용되었습니다. 복잡한 코드 생성에 내재된 희소 보상 문제를 해결하기 위해 새로운 "테스트 난이도 기반 코드 보상" 시스템이 도입되었습니다. 전부 아니면 전무(all-or-nothing) 방식의 보상 대신, 이 시스템은 문제 내에서 더 쉬운 테스트 케이스를 통과하는 것에 대해 부분 점수를 부여하여 모델이 학습할 수 있는 더 밀도 높은 기울기 신호를 제공합니다.

효율성 또한 핵심이었습니다. 모델이 개선됨에 따라 데이터 재샘플링 전략은 더 쉬운 문제의 가중치를 낮추어 더 어려운 예시에 학습을 집중시켰습니다. 샤오미는 또한 지속적인 생성, 비동기 보상 계산 및 조기 종료를 통합하여 GPU 유휴 시간을 최소화하는 최적화된 RL 인프라인 "Seamless Rollout Engine"을 개발하여 상당한 학습(2.29배) 및 검증(1.96배) 속도 향상을 달성했습니다.

MiMo-7B-RL 제품군: 간략한 개요

샤오미는 개발 단계를 보여주는 여러 모델을 출시했습니다:

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에서 쉽게 사용할 수 있습니다.

모델 액세스:

Hugging Face의 XiaomiMiMo 조직 페이지에서 MiMo-7B-RL 및 시리즈의 다른 모델을 찾을 수 있습니다. 모델 크기는 약 78.3억 매개변수입니다 (BF16 정밀도, Safetensors 형식).

vLLM으로 추론 실행 (권장)

샤오미는 잠재적으로 더 빠른 생성을 위한 다중 토큰 예측 기능을 지원하므로 vLLM (v0.7.3 기반)의 자체 포크를 사용하여 추론할 것을 권장합니다.

• 샤오미 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 빌드를 사용하는 경우, 샤오미가 제공하는 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 ---

사용 권장 사항

최상의 결과를 얻으려면, 특히 벤치마크 점수를 재현하려는 경우 권장 설정인 샤오미의 vLLM 포크 (v0.7.3 기반)와 빈 시스템 프롬프트를 사용하십시오.

결론: 샤오미가 실현한 효율적인 추론?

샤오미의 MiMo-7B-RL은 수학 및 코딩과 같은 전문 분야에서 예외적인 추론 성능이 거대한 모델 크기에 의존하지 않고도 달성 가능함을 보여줍니다. 추론 패턴에 초점을 맞춘 신중한 사전 학습과 혁신적인 강화 학습 기법을 통해 훨씬 더 큰 모델과 효과적으로 경쟁하는 효율적인 모델을 만들었습니다. MiMo 시리즈의 공개 출시는 귀중한 도구와 통찰력을 제공하며, 강력하고 접근 가능한 AI 추론 능력의 발전을 촉진합니다.