TEMPO: Scaling Test-time Training for Large Reasoning Models

Resumo Técnico: TEMPO

1. O Problema

Os Grandes Modelos de Raciocínio (LRMs) demonstram capacidades notáveis em tarefas complexas, mas seus parâmetros permanecem estáticos após o treinamento offline. A Treinamento em Tempo de Teste (TTT) visa superar essa limitação, permitindo que o modelo atualize seus parâmetros durante a inferência usando dados de teste não rotulados.

No entanto, os métodos existentes de TTT para LRMs (como TTRL e EMPO) enfrentam duas limitações fundamentais:

1. Estagnação de Desempenho (Plateaus): Eles dependem de sinais de recompensa heurísticos e auto-gerados (ex: entropia, consistência própria). À medida que o modelo evolui, esses sinais tornam-se enviesados, levando a um teto de desempenho onde o modelo não consegue mais melhorar.
2. Colapso de Diversidade: Em busca de maior desempenho médio, esses métodos tendem a convergir para um conjunto estreito de padrões de raciocínio, suprimindo soluções válidas alternativas e degradando a qualidade geral do raciocínio a longo prazo.

A causa raiz identificada é a falta de calibração externa: o sinal de recompensa deriva apenas das próprias saídas do modelo, criando um ciclo de auto-reforço que distorce a estimativa de correção.

2. Metodologia: TEMPO

O authors propõem o TEMPO (Test-time Expectation-Maximization Policy Optimization), um framework de TTT que desacopla a geração de recompensa da otimização da política através de um design alternado de Actor-Critic, formalizado pelo algoritmo Expectation-Maximization (EM).

O processo opera em dois estágios alternados:

• Passo E (Expectation) - Recalibração do Crítico:

  • O modelo crítico (Critic) é periodicamente atualizado usando um conjunto de dados rotulados ($D_L$) com recompensas verificáveis.
  • O objetivo é estimar a distribuição posterior de respostas corretas, servindo como um sinal de recompensa "aterrado" (grounded) e livre de viés do modelo atual.
  • Isso corresponde à estimativa da variável latente (correção da resposta) no framework EM.

• Passo M (Maximization) - Refinamento da Política:

  • O modelo ator (Policy) gera trajetórias de raciocínio em questões de teste não rotuladas ($D_u$).
  • A política é otimizada (via Gradiente de Política/RL) usando as recompensas fornecidas pelo crítico recalibrado.
  • Isso corresponde à atualização dos parâmetros do modelo para maximizar a probabilidade de gerar respostas corretas, dado o sinal do crítico.

Inovação Teórica: O artigo demonstra que métodos anteriores são variantes degeneradas do EM que executam apenas o Passo M (otimização da política) sem o Passo E (recalibração da distribuição posterior). Ao reintroduzir o Passo E periodicamente, o TEMPO aperta o Evidence Lower Bound (ELBO), permitindo melhoria sustentada.

3. Contribuições Principais

1. Framework TEMPO: Um método de TTT que alcança ganhos de desempenho sustentados através da otimização alternada Actor-Critic, evitando o colapso de diversidade e os platôs de desempenho observados em métodos de auto-treinamento anteriores.
2. Análise Unificada via EM: Uma caracterização teórica que identifica a ausência do Passo E (recalibração posterior) como a causa fundamental da falha de escalabilidade em métodos existentes. A proposta de restaurar a calibração periódica do crítico em dados rotulados é apresentada como a solução principista.
3. Validação Empírica Abrangente: Experimentos extensivos em diversas famílias de modelos (OLMO3, Qwen3) e escalas (7B, 8B, 14B), cobrindo benchmarks de matemática (AIME 2024/2025/2026), lógica e STEM.

4. Resultados Experimentais

O TEMPO demonstrou superioridade significativa em precisão e diversidade em comparação com baselines como RLVR (PPO padrão), TTRL e EMPO.

• Ganhos em Matemática (AIME 2024):
  • OLMO3-7B: Aumentou a precisão de 33.0% para 51.1%.
  • Qwen3-14B: Aumentou a precisão de 42.3% para 65.8%.
  • O modelo continuou a melhorar consistentemente por mais de 350 passos de treinamento, enquanto os baselines estagnaram ou colapsaram.
• Preservação de Diversidade:
  • Enquanto métodos como TTRL sofrem de colapso de diversidade (queda no pass@k mesmo com aumento no avg@k), o TEMPO mantém altas pontuações de pass@k, indicando que explora múltiplas trajetórias de raciocínio válidas.
• Generalização:
  • O método generalizou bem para tarefas não matemáticas, incluindo quebra-cabeças lógicos (ZebraLogic) e tarefas de nível de especialista (GPQA-Diamond), com ganhos absolutos de até +24.5 pontos em AGI Eval para o modelo OLMO3-7B.
• Ablação:
  • Experimentos mostraram que remover o Passo E (mantendo o crítico congelado) leva a uma estagnação rápida, confirmando que a recalibração periódica é essencial para evitar o desalinhamento entre o crítico e a política em evolução.

5. Significado e Impacto

O TEMPO representa um avanço crucial na capacidade de modelos de IA de aprender continuamente durante a inferência. Ao formalizar o treinamento em tempo de teste como um processo EM, o trabalho oferece uma solução teórica e prática para o problema de "alucinação de recompensa" em sistemas de auto-reforço.

A principal implicação é que computação adicional em tempo de teste pode ser convertida em ganhos de capacidade mensuráveis e escaláveis, desde que haja um mecanismo de calibração externa para evitar o viés de auto-referência. Isso abre caminho para agentes de IA mais robustos que podem adaptar-se dinamicamente a novos problemas complexos sem necessidade de retreinamento offline massivo.

Limitações Notadas: O método requer maior memória de GPU (para manter ator e crítico) e depende da disponibilidade de um conjunto de dados rotulados para a etapa de recalibração do crítico.