GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

Título: GTS: Escalonamento no Tempo de Inferência de Raciocínio Latente com um Amostrador de Pensamento Gaussiano Aprendível

1. O Problema

O Escalonamento no Tempo de Inferência (ITS - Inference-Time Scaling) é uma técnica crucial para melhorar o raciocínio de Grandes Modelos de Linguagem (LLMs) alocando mais poder computacional durante a inferência para gerar e selecionar múltiplas trajetórias de raciocínio.

• Modelos Discretos (Baseados em Tokens): O ITS funciona bem aqui porque existe uma distribuição de probabilidade condicional explícita sobre os próximos tokens, permitindo amostragem controlada (ex: temperature scaling, nucleus sampling).
• Modelos de Raciocínio Latente (Contínuo): Modelos que realizam raciocínio em múltiplos passos diretamente no espaço de estados ocultos (sem gerar tokens intermediários) carecem dessas distribuições explícitas.
• Limitação Atual: Para introduzir diversidade nessas trajetórias latentes, os métodos existentes dependem de perturbações heurísticas, como dropout ou ruído gaussiano fixo.
  • O Problema Central: Aumentar a força da perturbação (estocasticidade) não garante melhor qualidade de amostragem. Perturbações fortes podem causar grandes deslocamentos de distribuição que afastam o modelo das regiões que levam à resposta correta, resultando em "super-exploração" (ruído excessivo) ou "sub-exploração" (muito pouco desvio). Além disso, essas perturbações não definem uma distribuição de amostragem condicional explícita, tornando o processo difícil de controlar ou otimizar.

2. Metodologia: Gaussian Thought Sampler (GTS)

Os autores propõem o GTS, um módulo leve que reformula a exploração latente como uma amostragem condicional de uma distribuição aprendida sobre estados de raciocínio contínuos.

• Conceito Principal: Em vez de injetar ruído fixo, o GTS aprende uma política de amostragem probabilística explícita.
• Mecanismo de Funcionamento:
  1. Modelo Base Congelado: O modelo de raciocínio latente (o "backbone") permanece congelado.
  2. Política Gaussiana Condicional: O GTS é um pequeno módulo (MLP) que, dado o estado oculto determinístico $h_t$ em cada passo de raciocínio, prevê os parâmetros de uma distribuição gaussiana condicional: média ($\mu$) e desvio padrão ($\sigma$).
  3. Perturbação: A perturbação $z_t$ é amostrada dessa distribuição e adicionada ao estado oculto: $\tilde{h}_t = h_t + z_t$.
  4. Treinamento (Reinforcement Learning): O GTS é treinado usando uma abordagem estilo GRPO (Group Relative Policy Optimization), similar ao PPO, mas adaptada para trajetórias contínuas.
     • Recompensa: Baseada na correção da resposta final e em um termo de "confiança" (shaping reward) que penaliza respostas corretas com baixa confiança e respostas incorretas com alta confiança.
     • Objetivo: Otimizar a política de perturbação para maximizar a probabilidade de encontrar trajetórias que levem a respostas corretas, mantendo o backbone intacto.

3. Contribuições Chave

1. Identificação de uma Limitação Fundamental: O artigo demonstra que o sucesso do ITS em espaços latentes não depende apenas de injetar mais estocasticidade, mas de controlar a qualidade da amostragem. Perturbações heurísticas fixas frequentemente falham em direcionar a exploração para regiões úteis do espaço de raciocínio.
2. Proposta do GTS: Introdução de um amostrador leve que transforma a perturbação heurística em uma política de amostragem probabilística explícita e otimizable. O GTS aprende onde e quanto perturbar o estado latente dependendo do contexto.
3. Validação Empírica: Demonstração de que essa abordagem baseada em princípios supera consistentemente as heurísticas (dropout e ruído gaussiano fixo) em múltiplos benchmarks e arquiteturas de raciocínio latente.

4. Resultados Experimentais

Os experimentos foram conduzidos em dois modelos de raciocínio latente (COCONUT e CODI) e no conjunto de dados GSM8K (problemas matemáticos).

• Desempenho de Escalonamento (Pass@N):
  • O GTS superou significativamente as linhas de base heurísticas (Dropout e Ruído Gaussiano Padrão) à medida que o orçamento de amostragem ($N$) aumentava.
  • Enquanto métodos heurísticos mostravam ganhos estagnados ou degradados com $N$ alto (devido à super-exploração), o GTS manteve um crescimento robusto no desempenho.
• Qualidade da Amostragem (Análise Diagnóstica):
  • Ganho de Amostragem (SG): O GTS produziu ganhos positivos consistentes na probabilidade da resposta correta.
  • Deslocamento de Distribuição (JS Divergence): O GTS conseguiu um equilíbrio ideal: gerou diversidade suficiente para explorar novas trajetórias, mas sem causar deslocamentos de distribuição excessivos que arruinariam a coerência do raciocínio (ao contrário do dropout forte, que causou grande deslocamento com baixo ganho).
• Generalização: O GTS treinado no GSM8K demonstrou boa generalização em benchmarks fora da distribuição (MultiArith, SVAMP, GSM8K-Hard), mantendo sua superioridade sobre métodos heurísticos, que mostraram comportamento inconsistente entre diferentes tarefas.
• Análise de SNR (Relação Sinal-Ruído): O GTS aprendeu dinâmicas adaptativas. Em alguns modelos, alternou entre passos mais determinísticos e mais estocásticos, alinhando-se à estrutura interna de raciocínio do modelo base, em vez de aplicar um ruído uniforme.

5. Significado e Conclusão

O trabalho estabelece que o Escalonamento no Tempo de Inferência (ITS) em modelos de raciocínio contínuo exige uma mudança de paradigma:

• Do Heurístico para o Aprendível: A simples injeção de ruído não é suficiente. A exploração deve ser modelada como um problema de otimização de política.
• Controle vs. Diversidade: O sucesso não vem da maximização da diversidade de respostas, mas da maximização da utilidade da diversidade (amostrar trajetórias que realmente melhoram a decisão final).
• Eficiência: O GTS oferece uma solução leve (apenas alguns milhões de parâmetros adicionais) que não requer o re-treinamento do modelo base, tornando-o prático para aplicação em modelos existentes.

Em resumo, o GTS demonstra que um amostrador de pensamento gaussiano aprendível e condicional é superior a perturbações estáticas, permitindo um controle fino sobre a exploração no espaço latente e resultando em raciocínios mais robustos e escaláveis.