Latent Reasoning with Supervised Thinking States

Resumo Técnico: Raciocínio Latente com Estados de Pensamento Supervisionados

1. O Problema

O raciocínio via Chain-of-Thought (CoT) — onde o modelo gera uma sequência de passos lógicos em linguagem natural antes da resposta final — aumentou significativamente a precisão de LLMs em tarefas complexas. No entanto, o CoT apresenta dois grandes problemas:

• Custo de Inferência: A geração de longos raciocínios aumenta o número de tokens, elevando a latência e o custo computacional.
• Extensão do Contexto: O raciocínio explícito consome espaço na janela de contexto do modelo.

Abordagens anteriores de "raciocínio latente" tentaram substituir o texto por embeddings contínuos para economizar tokens, mas enfrentaram dificuldades na supervisão (exigindo o custoso Backpropagation Through Time - BPTT) ou na capacidade de expressividade (ao tentar comprimir o pensamento apenas na profundidade das camadas do modelo).

2. Metodologia: Thinking States

Os autores propõem uma arquitetura que permite o raciocínio recorrente durante o processamento da entrada, sem aumentar o comprimento do contexto. O método baseia-se em três componentes principais:

• Processamento por Blocos (Chunking): A sequência de entrada é dividida em blocos não sobrepostos ($X_i$).
• Bloco de Pensamento (Thinking Block - T): Um decodificador Transformer leve que, a partir das representações de um bloco, gera uma sequência de tokens de pensamento em linguagem natural ($Z_i$).
• Bloco de Compressão (Compression Block - C): Um codificador que transforma essa sequência de texto variável em um estado latente de tamanho fixo ($S_i$).

O Fluxo de Trabalho:

1. O estado de pensamento do bloco anterior ($S_i$) é injetado nas representações do bloco atual ($X_i$) em uma camada rasa (inicial) do modelo principal.
2. O modelo processa o bloco e gera representações profundas.
3. O Bloco $T$ decodifica essas representações em pensamentos legíveis.
4. O Bloco $C$ comprime esses pensamentos no próximo estado $S_{i+1}$, que será usado no próximo bloco.

Treinamento com Teacher-Forcing:
Diferente de métodos latentes puramente contínuos, o Thinking States utiliza anotações em linguagem natural. Durante o treinamento, os autores utilizam o "estado de ouro" (baseado no texto real do CoT) para injetar no modelo. Isso permite o treinamento totalmente paralelo, eliminando a necessidade de BPTT e tornando o treinamento muito mais rápido.

3. Principais Contribuições

• Raciocínio Recorrente sem Expansão de Contexto: O modelo "pensa" enquanto lê, mas os pensamentos são comprimidos em estados que não ocupam espaço na janela de contexto.
• Supervisão via Linguagem Natural: Ao manter os pensamentos como tokens de texto durante o treinamento, o modelo beneficia-se de dados de CoT existentes e mantém a interpretabilidade.
• Eficiência de Treinamento: A técnica de teacher-forcing permite que o custo de treinamento seja quase constante, independentemente da profundidade do raciocínio, ao contrário de métodos que exigem BPTT.
• Algoritmo de Prefill Rápido: Introduziram um método de preenchimento especulativo que acelera a inferência ao assumir que muitos blocos de entrada não exigem raciocínio complexo (estados triviais).

4. Resultados Experimentais

Os testes foram realizados em tarefas de rastreamento de estado (State Tracking), QA de múltiplos saltos (Multi-Hop QA) e problemas matemáticos (GSM8K).

• Generalização de Comprimento: Em tarefas de rastreamento de estado (como Paridade e Atribuição de Variáveis), o Thinking States superou o CoT em generalização para sequências muito mais longas do que as vistas no treinamento. Isso sugere que o modelo aprendeu um mecanismo de raciocínio recorrente robusto, em vez de apenas memorizar padrões.
• Desempenho vs. Eficiência:
  • No GSM8K, o método aproximou-se da precisão do CoT, mas com um ganho de velocidade de aproximadamente 2.66x.
  • Em Multi-Hop QA, o método igualou ou superou o desempenho do CoT com latência reduzida.
• Comparação com Baselines: Superou consistentemente métodos de raciocínio latente existentes (como Coconut e iCoT) em precisão e velocidade.

5. Significância e Conclusão

O trabalho demonstra que é possível obter os benefícios cognitivos do Chain-of-Thought sem os custos proibitivos de tempo e memória. A capacidade de comprimir o raciocínio em estados latentes supervisionados abre caminho para:

1. LLMs mais eficientes: Redução do consumo de energia e latência em aplicações de produção.
2. Melhoria via RL: O modelo pode servir como um ponto de partida (warm-start) para Aprendizado por Reforço, onde o modelo aprende a otimizar seus próprios estados de pensamento internos além do que os humanos conseguem escrever.
3. Interpretabilidade: Ao contrário de outros modelos latentes "caixa-preta", os estados de pensamento aqui podem ser lidos e auditados, pois derivam de tokens de linguagem natural.