Probabilistic Programs of Thought

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) são amplamente utilizados para geração de código e raciocínio matemático, frequentemente exigindo a produção de saídas estruturadas. A abordagem padrão para garantir a correção desses programas envolve um ciclo de amostragem-execução-verificação (geralmente conhecido como Best-of-n ou Beam Search):

1. O modelo gera um programa.
2. O programa é executado e verificado.
3. Se falhar, o processo é repetido, gerando novos programas.

A Limitação: Este processo é extremamente custoso em termos computacionais. Gerar $n$ programas distintos requer $n$ chamadas completas ao LLM, cada uma consumindo recursos intensivos de GPU. Para modelos grandes, mesmo um número pequeno de amostras torna-se proibitivamente caro, limitando a eficácia das estratégias de busca no tempo de teste (test-time).

2. Metodologia: Programas Probabilísticos de Pensamento (PPoT)

Os autores propõem o PPoT, uma nova técnica de decodificação no tempo de teste que explora a distribuição de probabilidade subjacente do LLM dentro do próprio programa gerado, permitindo obter múltiplas amostras com apenas uma geração inicial do modelo.

O funcionamento baseia-se nos seguintes pilares:

• LLM como Programa Probabilístico: Em vez de tratar a saída do LLM como um programa determinístico fixo, o PPoT interpreta os tokens gerados (especialmente dígitos, operadores aritméticos e de atribuição) como variáveis aleatórias.
• Construção do Programa:
  1. O LLM gera um programa inicial e fornece as probabilidades dos próximos tokens (next-token probabilities).
  2. Um subconjunto de tokens críticos (identificados como fontes de incerteza, como números ou operadores) é convertido em variáveis aleatórias categóricas, parametrizadas pelas logits do LLM.
  3. O restante do código permanece fixo.
  4. O resultado é um programa probabilístico que representa compactamente uma distribuição sobre um número exponencial de programas determinísticos possíveis.
• Amostragem Eficiente: Uma vez construído o programa probabilístico, novas soluções são amostradas diretamente dele usando inferência probabilística (executada na CPU), sem a necessidade de novas chamadas ao GPU do LLM.
• Suposição de Independência: O método assume que os tokens ressampleados são independentes dos tokens subsequentes no contexto da sequência original (uma suposição comum em modelos autoregressivos para tornar o problema tratável). Sob essa suposição, a distribuição empírica das amostras converge para a distribuição original do LLM.
• Amostragem Condicionada: Para evitar a repetição da mesma solução incorreta, o sistema pode amostrar condicionando-se a não reproduzir o programa original gerado pelo LLM.

3. Principais Contribuições

1. Novo Paradigma de Decodificação: Introdução do PPoT, uma técnica que transforma a geração de código estruturado em um problema de programação probabilística, permitindo a geração de múltiplas variantes de um programa a partir de uma única chamada ao modelo.
2. Fundamentação Teórica: Demonstração de que a abordagem é distribucionalmente correta (sob certas suposições de independência) e que a distribuição empírica das amostras converge para a distribuição do LLM.
3. Eficiência Computacional: Prova de que o PPoT oferece ganhos significativos de desempenho sem custo adicional de GPU, movendo a carga de trabalho para a CPU (inferência probabilística leve).
4. Validação Empírica: Avaliação abrangente em três benchmarks distintos (raciocínio matemático, geração de código visual e geração de saída estruturada).

4. Resultados Experimentais

Os autores avaliaram o PPoT utilizando os modelos Qwen2.5 Coder (0.5B, 3B, 7B) e Qwen2.5 VL em três conjuntos de dados:

• GSM8k (Raciocínio Matemático):
  • Com apenas 5 amostras adicionais do PPoT por geração do LLM ($m=5$), a precisão (pass@k) aumentou entre 2% e 7% sem custo de GPU adicional.
  • O PPoT conseguiu atingir o desempenho de 20 amostras do LLM usando apenas 8 amostras do LLM + amostras PPoT.
• Plot2Code (Geração de Código para Gráficos):
  • Melhoria de 0,5% a 6% na pontuação de correspondência de texto, corrigindo erros de tempo de execução (como dimensões de vetores incompatíveis) que o LLM original não conseguiu resolver.
• CRUXEval (Geração de Saída Estruturada/Inversão de Código):
  • Aumento de 3% a 7% na precisão, demonstrando a aplicabilidade do método além de programas executáveis, estendendo-se a objetos estruturados (listas, strings).

Eficiência: O tempo de compilação e amostragem do PPoT é computacionalmente barato (dominado pela inferência do LLM inicial). Gráficos mostram que o tempo de execução para $m=20$ é quase indistinguível de $m=0$ (apenas LLM), enquanto a precisão aumenta drasticamente.

5. Significado e Impacto

O trabalho do PPoT representa uma mudança significativa na forma como interagimos com LLMs para tarefas de geração de código:

• Redução de Custos: Permite escalar a busca de soluções (aumentando $k$ efetivamente) sem escalar linearmente o custo de inferência em GPU, que é o gargalo atual.
• Exploração de Incerteza: Em vez de tratar a saída do LLM como um fato, o PPoT explora a incerteza intrínseca do modelo (probabilidades de tokens) para corrigir erros locais (como um operador errado ou um número incorreto) que podem invalidar todo o programa.
• Abstração de Modelos: Propõe ver a geração de código não apenas como uma sequência de tokens, mas como uma distribuição estatística sobre programas, abrindo caminho para o uso de técnicas de inferência probabilística clássica para melhorar a robustez de LLMs.

Em resumo, o PPoT oferece um método "barato" para realizar uma busca mais ampla no espaço de soluções, transformando uma única tentativa falha do LLM em uma oportunidade de gerar dezenas de variações corretas via inferência probabilística na CPU.