Resumo Técnico: SuperThoughts: Tokens de Raciocínio em Superposição

Declaração do Problema

Grandes Modelos de Linguagem (LLMs) resolvem problemas complexos gerando sequências explícitas de Cadeia de Pensamento (Chain-of-Thought - CoT), onde cada token representa uma unidade de computação (uma passagem direta/forward pass). Embora modelos recentes como o OpenAI o1 e o DeepSeek-R1 demonstrem que o aumento do computação em tempo de teste melhora significamente o desempenho, essa geração sequencial de tokens é computacionalmente dispendiosa.

Abordagens existentes que tentam contornar a geração discreta de tokens ao raciocinar em espaços latentes contínuos (ex: COCONUT, raciocínio latente) enfrentam um desafio crítico: a falta de supervisão intermediária. O treinamento padrão de CoT depende da perda de entropia cruzada ao nível do token em cada etapa, fornecendo sinais de gradiente densos. Quando o raciocínio ocorre em um espaço latente não restringido, essa supervisão desaparece, levando à instabilidade de treinamento, deriva representacional e baixo desempenho em tarefas de longo horizonte. Por outro lado, métodos que simplesmente comprimem o contexto ou reduzem a contagem de tokens discretos via fine-tuning não reduzem o número de passagens diretas necessárias, limitando seu impacto no throughput de inferência.

A questão central abordada por este trabalho é: Pode um modelo ser treinado para raciocinar em um espaço latente mais rico e superposto, mantendo os sinais de supervisão intermediária da geração de tokens discretos?

Metodologia

Os autores propõem o SuperThoughts, um framework que comprime pares de tokens consecutivos de CoT em representações latentes únicas, permitindo que o modelo decodifique dois tokens por etapa. Esta abordagem reduz pela metade o número de passagens diretas necessárias para uma cadeia de raciocínio de comprimento $L$ (reduzindo-a para $L/2$ etapas) enquanto preserva a supervisão de tokens discretos.

Arquitetura

O modelo consiste em três componentes primários:

1. Compressor: Codifica um par de embeddings de tokens $(c_{2i-1}, c_{2i})$ em um único vetor latente $x_i$. Os autores exploram duas implementações:
   • Projeção Linear: Concatena os embeddings e os projeta via uma matriz aprendível.
   • Bloco Transformer: Processa o par através de uma pequena camada Transformer, extraindo o estado oculto do segundo token.
2. Módulo Principal (Main Module): O backbone de raciocínio primário (inicializado a partir de um LLM pré-treinado). Ele processa o vetor comprimido $x_i$ para produzir um estado oculto $h_i$ e prediz o próximo token de índice ímpar ($c_{2i+1}$).
3. Módulo de Predição Multi-Token (MTP Module): Um módulo leve que prediz o subsequente token de índice par ($c_{2i+2}$). Ele recebe como entrada o token par anterior, o novo token ímpar predito e o estado oculto do módulo Principal, projetando-os em um vetor $x'_i$ para gerar a próxima predição. Ambos os módulos compartilham o mesmo cabeçalho de modelagem de linguagem.

Estratégia de Treinamento

Para lidar com a instabilidade do raciocínio latente, um protocolo de treinamento de dois estágios é empregado:

1. Destilação Latente (Estágio 1): O Compressor é treinado para alinhar o espaço latente comprimido com o manifold de linguagem pré-treinado. Um modelo "professor" processa a sequência de tokens discretos, enquanto um modelo "aluno" processa os pares comprimidos. O Compressor é otimizado para minimizar a distância Smooth-L1 entre os estados ocultos do aluno (após processar $x_i$) e os estados ocultos do professor (após processar o par de tokens correspondente). Isso garante que as representações comprimidas preservem o estado de raciocínio da sequência discreta original.
2. Treinamento Conjunto de Ponta a Ponta (Estágio 2): O módulo Principal, o módulo MTP e o Compressor são treinados conjuntamente usando perda de entropia cruzada padrão para todos os tokens preditos (tanto para as saídas do Principal quanto do MTP).

Inferência Adaptativa

Reconhecendo que nem todas as etapas de raciocínio são igualmente difíceis, os autores introduzem um mecanismo adaptativo baseado em confiança. Em cada etapa, se a confiança da predição do módulo MTP (probabilidade máxima) cair abaixo de um limiar $\tau$, o modelo rejeita a predição do MTP. Ele então retorna para a decodificação padrão para aquela etapa, alimentando o módulo Principal diretamente com o token único. Isso permite que o modelo aloque mais computação para tokens "difíceis" enquanto mantém a eficiência nos tokens "fáceis".

Principais Contribuições

1. Arquitetura SuperThoughts: Um framework inovador que funde pares de tokens em embeddings superpostos e decodifica dois tokens por etapa usando um módulo Principal e um módulo MTP, efetivamente reduzindo pela metade as passagens diretas necessárias.
2. Protocolo de Treinamento de Dois Estágios: Um método para estabilizar o raciocínio latente ao primeiro alinhar o espaço comprimido via destilação antes do treinamento conjunto de ponta a ponta com supervisão discreta.
3. Mecanismo de Inferência Adaptativa: Um sistema de fallback dinâmico que troca throughput por precisão em etapas de raciocínio incertas, evitando a degradação de desempenho em tarefas complexas.
4. Avaliação Empírica: Avaliação abrangente em benchmarks de raciocínio matemático e científico (MATH500, AMC23, OlympiadBench, GPQA-Diamond) demonstrando redução significativa no comprimento do CoT com perda mínima de precisão.

Resultados

Os autores avaliaram o SuperThoughts em modelos Qwen2.5-Math (variantes de 1.5B, 7B e 14B) em múltiplos benchmarks:

• Redução de Comprimento: Com o mecanismo adaptativo, o SuperThoughts alcança uma redução de 20–35% no comprimento do CoT.
• Precisão: O método mantém a precisão dentro de 1–2 pontos do baseline de CoT discreto padrão na maioria das tarefas. Por exemplo, no modelo de 1.5B com $\tau=0.999$, a precisão no MATH500 atingiu 73.0% (vs. 72.4% do baseline) com uma redução de 36% no comprimento do CoT.
• Efeitos de Escala: Modelos maiores (7B) toleram melhor a compressão agressiva do que modelos menores (1.5B). Sem a inferência adaptativa, o modelo de 1.5B sofreu quedas significativas de precisão (14–21 pontos), enquanto o modelo de 7B caiu apenas 5–12 pontos.
• Comparação com HAMburger: O SuperThoughts superou o baseline HAMburger (Liu & Zhang, 2025) em todos os níveis de compressão, alcançando maior precisão para o mesmo comprimento de CoT ou comprimentos de CoT menores para a mesma precisão.
• Aceleração de Inferência: Embora o speedup teórico seja proporcional à redução de comprimento, o speedup real de tempo de relógio (wall-clock time) é ligeiramente menor devido ao overhead dos módulos Compressor e MTP. No entanto, esse overhead diminui conforme o tamanho do modelo aumenta; para o modelo de 14B, uma redução de ~33% no comprimento traduziu-se em uma redução de ~28% no tempo de geração.

Significância

O artigo afirma que o SuperThoughts demonstra com sucesso que modelos podem raciocinar em um espaço latente superposto sem sacrificar a estabilidade proporcionada pela supervisão de tokens discretos. Ao comprimir pares de tokens e utilizar inferência adaptativa, o framework oferece um caminho prático para reduzir o custo computacional de tarefas de raciocínio de longo horizonte. O trabalho sugere que modelos maiores possuem capacidade por etapa suficiente para lidar com tokens superpostos de forma confiável, sugerindo que o escalonamento futuro poderia permitir uma compressão ainda maior (ex: superpor três ou mais tokens). Além disso, os autores observam que, à medida que os módulos MTP nativos se tornarem padrão em futuros LLMs (ex: Qwen3-Next, MiMo), a aplicação do SuperThoughts se tornará ainda mais direta e eficaz.