QED-Nano: Teaching a Tiny Model to Prove Hard Theorems

Resumo Técnico: QED-Nano

1. O Problema

Sistemas de IA proprietários (como Gemini 3 Pro) têm demonstrado capacidades impressionantes na resolução de problemas matemáticos de nível olímpico (IMO), atingindo desempenho de nível de medalha de ouro. No entanto, essas soluções dependem de modelos massivos (centenas de bilhões de parâmetros), pipelines de treinamento não divulgados e custos de inferência proibitivos. Isso cria uma lacuna entre o que é teoricamente possível e o que a comunidade de pesquisa de código aberto pode reproduzir, estudar ou melhorar.

A questão central abordada pelo artigo é: Modelos pequenos e abertos podem ser pós-treinados para competir com sistemas proprietários muito maiores na geração de provas matemáticas rigorosas de nível olímpico?

O desafio específico é que a escrita de provas matemáticas difere da resolução de problemas com respostas finais verificáveis. Exige-se a geração de argumentos longos, globalmente coerentes e matematicamente rigorosos, em vez de apenas encontrar um valor numérico final.

2. Metodologia

Os autores desenvolveram o QED-Nano, um modelo de linguagem de 4 bilhões de parâmetros (baseado no Qwen3), totalmente treinado para provar teoremas em linguagem natural, sem depender de ferramentas externas como Lean. O pipeline de treinamento consiste em três estágios principais:

A. Estágio 1: Ajuste Fino Supervisionado (SFT) com Distilação

• Objetivo: Instilar um estilo básico de escrita de provas e cobrir estratégias de prova.
• Dados: O modelo foi ajustado em um conjunto de dados de ~7.500 respostas geradas pelo DeepSeek-Math-V2 (um modelo de 685B parâmetros especializado em matemática).
• Processo: O conjunto de dados foi filtrado para manter apenas soluções estruturalmente válidas. O objetivo foi "ensinar" ao modelo de 4B como estruturar um argumento lógico, servindo como uma inicialização forte para o RL.
• Desafio Identificado: O SFT induziu um fenômeno de "explosão de comprimento", onde o modelo tendia a repetir-se ou divagar até esgotar a janela de contexto, imitando a forma superficial de provas longas sem a estrutura lógica necessária.

B. Estágio 2: Aprendizado por Reforço (RL) com Recompensas Baseadas em Rubricas

• Algoritmo: Utilização do GRPO (Group Relative Policy Optimization) com atualizações assíncronas.
• Recompensa: Em vez de recompensas binárias (sucesso/fracasso), o modelo recebe feedback denso baseado em rubricas de correção. Um avaliador (LLM judge, especificamente o GPT-OSS-20B) pontua as provas de 0 a 7 com base em checkpoints intermediários, falhas comuns e deduções específicas.
• Curriculum: O conjunto de prompts foi filtrado para remover problemas muito fáceis (onde o modelo já acertava) e muito difíceis (onde nenhuma tentativa tinha sucesso), focando em uma distribuição unimodal de dificuldade para maximizar a eficiência do treinamento.

C. Estágio 3: Expansão com Cache de Raciocínio (Reasoning Cache - RC)

• Inovação: Para lidar com provas extremamente longas (milhares de tokens) sem a variância alta típica do RL em longos horizontes, os autores integraram o framework Reasoning Cache.
• Mecanismo: O raciocínio é decomposto em ciclos iterativos de "resumir e refinar". Em cada iteração, o modelo gera um segmento de raciocínio, resume o progresso em uma representação textual compacta ("estado") e condiciona a próxima geração ao problema original e a esse resumo.
• Treinamento: O modelo é treinado com RL para melhorar a geração condicionada a esses resumos. Isso permite que o modelo explore horizontes de raciocínio equivalentes a centenas de milhares de tokens durante a inferência, mantendo os rollouts de treinamento curtos e estáveis.

3. Principais Contribuições

1. QED-Nano (4B): Um modelo de código aberto de apenas 4B parâmetros que supera modelos abertos muito maiores (como Nomos-1 de 30B e GPT-OSS-120B) em benchmarks de prova matemática.
2. Pipeline de Treinamento Aberto: A primeira divulgação completa de um pipeline de pós-treinamento (SFT + RL com rubricas + RC) capaz de gerar provas de nível olímpico, incluindo datasets (FineProofs-SFT, FineProofs-RL) e código.
3. Eficácia do Test-Time Compute: Demonstração de que pequenos modelos podem escalar seu desempenho significativamente ao serem acoplados a scaffolds de inferência (como o RSA ou Reasoning Cache), permitindo que "pensem" por mais tempo (até 2 milhões de tokens por problema) sem aumentar o tamanho do modelo.
4. Superação de Limitações do SFT: O uso de RL corrigiu as patologias introduzidas pelo SFT (como a repetição excessiva e a falta de estrutura), ensinando o modelo a ser mais conciso e rigoroso.

4. Resultados

O desempenho foi avaliado em três benchmarks principais: IMO-ProofBench, ProofBench e IMO-AnswerBench.

• Desempenho Base (sem scaffold): O QED-Nano alcançou 40% no IMO-ProofBench e 45% no ProofBench, superando significativamente outros modelos de 4B e competindo com modelos de 30B-120B.
• Desempenho com Scaffold (Agente): Ao acoplar o modelo com o scaffold RSA (Recursive Self-Assembly) durante a inferência (permitindo múltiplos rollouts e síntese):
  • IMO-ProofBench: 56.9% (vs. 58.7% do Gemini 3 Pro proprietário).
  • ProofBench: 62.6%.
  • IMO-AnswerBench: 76.5%.
• Comparação: O QED-Nano (Agente) supera o Nomos-1 (30B) e o GPT-OSS-120B, aproximando-se do desempenho do DeepSeek-Math-V2 (685B) e do Gemini 3 Pro, com uma fração do custo de inferência.
• Qualidade: Análises qualitativas mostraram que as provas do QED-Nano são mais estruturadas e explícitas do que as de outros modelos pequenos, embora tendam a usar abordagens computacionais intensivas em vez de insights estruturais elegantes.

5. Significado e Impacto

• Democratização da IA Matemática: O trabalho prova que não é necessário um modelo de 100B+ parâmetros para atingir o estado da arte em provas matemáticas. Modelos pequenos e especializados, quando adequadamente treinados, podem igualar ou superar sistemas gerais massivos.
• Custo-Efetividade: A abordagem de "pequeno modelo + inferência computacionalmente intensiva" (test-time scaling) oferece um trade-off custo-desempenho muito superior, tornando a pesquisa em matemática avançada acessível a instituições com recursos limitados de GPU.
• Direção Futura: O artigo sugere que o foco deve mudar do aumento puramente paramétrico para o desenvolvimento de pipelines de pós-treinamento robustos (SFT + RL + Scaffolding) e a melhoria da sinergia entre eles, especialmente para mitigar a "explosão de comprimento" e encorajar insights criativos ("aha moments") em vez de apenas refinamento computacional.

Em resumo, o QED-Nano estabelece um novo marco para modelos abertos em raciocínio matemático, demonstrando que a especialização via treinamento cuidadoso pode superar a escala bruta.