s1: Simple test-time scaling
Referência: https://arxiv.org/html/2501.19393v3
O artigo intitulado “s1: Simple Test-Time Scaling” propõe uma abordagem direta e eficiente para ampliar a capacidade de raciocínio de modelos de linguagem por meio de computação adicional no momento da inferência (test-time compute). Os autores, Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, entre outros pesquisadores de Stanford, Universidade de Washington e do Instituto Allen para IA, apresentam uma solução que combina curadoria de dados cuidadosamente selecionados com uma técnica de controle de geração chamada “budget forcing”.
1. Contextualização e Motivação
Historicamente, os avanços em modelos de linguagem foram alcançados com o aumento de dados e de capacidade computacional durante o treinamento. No entanto, o paradigma emergente de escalonamento em tempo de inferência sugere que é possível melhorar o desempenho do modelo mesmo após o treinamento, aumentando o cálculo dedicado à geração de respostas. O modelo o1 da OpenAI é citado como um exemplo bem-sucedido, mas sua metodologia não foi divulgada, o que incentivou várias tentativas de reprodução.
2. Objetivo
Os autores visam demonstrar que é possível obter ganhos significativos em tarefas de raciocínio utilizando apenas 1.000 exemplos supervisionados, aliados a uma técnica simples de controle computacional durante a inferência. Essa abordagem busca maximizar a eficiência amostral e facilitar a reprodutibilidade.
3. Criação do Dataset s1K
A primeira etapa consistiu na coleta de um conjunto inicial de 59.029 questões provenientes de 16 fontes, categorizadas segundo três princípios: qualidade, dificuldade e diversidade. Fontes como NuminaMATH, AIME, OlympicArena, OmniMath e AGIEval forneceram problemas matemáticos e interdisciplinares. Adicionalmente, os autores criaram dois conjuntos originais: s1-prob, com questões de provas de doutorado em Estatística, e s1-teasers, com enigmas de entrevistas de finanças quantitativas.
Para cada questão, foi gerado um traço de raciocínio e uma solução utilizando a API Gemini 2.0 Flash Thinking Experimental da Google. A partir disso, foram obtidos trios de dados: pergunta, trace de raciocínio e resposta final. Esse processo é conhecido como destilação, em que o conhecimento de um modelo mais poderoso (o Gemini) é transferido para outro modelo menor via supervisionamento.
Exemplo de entrada do dataset s1K:
- Pergunta: “Qual é a área de um círculo de raio 3?”
- Trace de raciocínio gerado pelo Gemini: “Sabemos que a área de um círculo é dada por A = πr^2. Neste caso, o raio r = 3. Então, A = π * 3^2 = π * 9. Portanto, a área é 9π.”
- Resposta final: “9π”
Esses trânsitos de raciocínio são utilizados como labels no treinamento supervisionado. Importa destacar que nem todos os exemplos são 100% corretos do ponto de vista matemático; os autores priorizaram capturar o processo de raciocínio, mesmo que algumas soluções finais estejam incorretas.
As amostras foram descontaminadas com checagem por 8-gramas para evitar sobreposição com benchmarks de teste e, em seguida, deduplicadas. O conjunto foi filtrado em três estágios:
- Qualidade: remoção de exemplos mal formatados ou com erros de API.
- Dificuldade: exclusão de questões resolvidas corretamente pelos modelos Qwen2.5-7B ou 32B.
- Diversidade: distribuição balanceada entre 50 domínios segundo a classificação MSC.
O resultado foi o s1K: um conjunto de 1.000 exemplos com ampla cobertura, dificuldade elevada e traços de raciocínio realistas extraídos do Gemini.
3.1 Enfatizando - Dataset s1k
Inicialmente, foram coletadas 59.029 questões de diversas fontes (como NuminaMATH, AIME, OmniMath, etc.), conforme descrito na seção 2.1 do artigo e resumido neste markdown.
Para cada uma dessas questões, foi gerado automaticamente um reasoning trace (um passo a passo da resolução) e uma resposta final usando o modelo Gemini 2.0 Flash Thinking Experimental da Google. Ou seja, as 59K foram passadas para o Gemini, que gerou os rascunhos de raciocínio e respostas — formando os trios: pergunta, raciocínio, resposta.
Depois disso, houve uma filtragem em três estágios (qualidade, dificuldade e diversidade) para selecionar apenas 1.000 exemplos, que formam o dataset final chamado s1K.
Esse s1K foi usado para fazer o supervised fine-tuning do modelo Qwen2.5-32B-Instruct, resultando no s1-32B.
Ou seja, sim: todos os 59 mil exemplos passaram pelo Gemini para gerar os dados, mas apenas 1.000 foram selecionados manualmente para o treinamento supervisionado.
4. Treinamento Supervisionado (SFT)
O modelo utilizado foi o Qwen2.5-32B-Instruct, um modelo de linguagem de 32 bilhões de parâmetros. O fine-tuning supervisionado (SFT) foi feito em apenas 26 minutos utilizando 16 GPUs NVIDIA H100 com FSDP (Fully Sharded Data Parallel) do PyTorch.
O processo consistiu em treinar o modelo para prever o próximo token com base nos dados de entrada compostos por pergunta + trace de raciocínio. O objetivo era alinhar o comportamento do modelo ao estilo de raciocínio do Gemini, mesmo com uma quantidade reduzida de exemplos.
5. Budget Forcing: Controle de Raciocínio em Tempo de Execução
Para controlar a quantidade de “raciocínio” que o modelo faz antes de responder, os autores propuseram uma intervenção no processo de decodificação denominada budget forcing. Essa técnica manipula diretamente o texto gerado:
- Delimitador de fim de raciocínio: é um token especial que indica para o modelo que é hora de parar de pensar e começar a gerar a resposta. Isso pode ser, por exemplo, a string “Final Answer:”.
- Estender o raciocínio: se o modelo tentar finalizar muito cedo, os autores interceptam esse ponto e inserem a string “Wait” no lugar, forçando o modelo a continuar sua geração.
Exemplo:
Pergunta: “Quantos ‘r’ existem em ‘raspberry’?”
Resposta inicial do modelo: “…A letra ‘r’ aparece 2 vezes.”
Com budget forcing, ao detectar essa tentativa de finalizar, adiciona-se: “Wait”
O modelo continua: “Reanalisando… r - a - s - p - b - e - r - r - y… Três letras ‘r’. Final Answer: 3”
Essa extensão simples pode levar a correções valiosas.
6. Comparativos e Benchmarking
Três benchmarks de raciocínio foram utilizados:
- AIME24: 30 questões da prova de matemática AIME de 2024.
- MATH500: 500 questões de matemática competitiva de dificuldades variadas.
- GPQA Diamond: 198 questões de nível PhD em ciências.
O modelo s1-32B alcançou:
- 56.7% em AIME24
- 93.0% em MATH500
- 59.6% em GPQA
Comparativamente, ele superou o modelo base (Qwen2.5-32B) e foi competitivo com modelos como Gemini 2.0 e o1-preview, mesmo com muito menos dados.
7. Estudos de Ablation
Diversos experimentos foram realizados para entender os efeitos da seleção dos dados:
- 1K-random: seleção aleatória resultou em desempenho significativamente pior.
- 1K-diverse: diversidade sem dificuldade também foi insuficiente.
- 1K-longest: selecionar pelos traces mais longos ajudou em GPQA, mas não superou o s1K completo.
- 59K-full: treinar com todos os exemplos gerou ganhos marginais ao custo de 50x mais GPU horas.
8. Limitações e Perspectivas Futuras
Embora o budget forcing seja eficaz, ele apresenta dois desafios:
- Repetição: o modelo pode entrar em ciclos se for forçado a continuar repetidamente.
- Janela de contexto: existe um limite de tokens que o modelo pode observar.
Os autores sugerem explorações com combinação de métodos paralelos (e.g. REBASE, votação majoritária), uso de penalidades de frequência ou técnicas de reforço para ampliar os ganhos.
9. Conclusão
O artigo oferece uma solução minimalista e surpreendentemente poderosa para razão computacional com modelos de linguagem. Com apenas 1.000 exemplos supervisionados, extraídos do Gemini, e um mecanismo simples de extensão de raciocínio via strings como “Wait”, os autores demonstram que é possível competir com os melhores modelos comerciais. A abordagem s1-32B representa um marco em eficiência amostral e reforça a importância de métodos abertos e reprodutíveis.
Código e dados estão disponíveis em: https://github.com/simplescaling/s1