Resumo Técnico: Raciocínio Latente Comprimido (CoLaR)

Definição do Problema

Os Grandes Modelos de Linguagem (LLMs) alcançam um desempenho superior em raciocínio matemático através do prompting de Cadeia de Pensamento (Chain-of-Thought - CoT), que gera cadeias de raciocínio explícitas, ao nível de tokens. No entanto, essas sequências extensas são computacionalmente caras e ineficientes, criando cargas significativas nos servidores e dificultando a escalabilidade, particularmente sob condições de alta concorrência. Embora os métodos de eficiência existentes tentem pular tokens redundantes ou usar notações compactas, eles permanecem limitados por representações esparsas baseadas em tokens. Além disso, abordagens recentes de raciocínio baseadas em latentes (ex: Coconut, CODI) frequentemente dependem de cadeias de raciocínio de comprimento fixo e geração latente determinística, o que limita sua capacidade de adaptação e exploração.

Metodologia: Raciocínio Latente Comprimido (CoLaR)

O CoLaR é um novo framework projetado para comprimir dinamicamente processos de raciocínio em um espaço latente denso, preservando a capacidade de explorar diversos caminhos de raciocínio. O método opera através de uma abordagem de treinamento em duas etapas:

1. Ajuste Fino Supervisionado (SFT) com Objetivos Auxiliares

O CoLaR estende além da predição padrão de próximo token ao introduzir um objetivo auxiliar de predição de próximo embedding comprimido.

• Compressão Dinâmica: Durante o treinamento, um fator de compressão $c$ é amostrado aleatoriamente de um intervalo predefinido $[1, c_{max}]$. Os embeddings de $c$ tokens de raciocínio consecutivos são fundidos em um único embedding comprimido usando um módulo de Compressão de Embedding. Este módulo escala a soma dos embeddings por $1/\sqrt{c}$ para preservar as propriedades da distribuição original, evitando a distorção causada pelo simples pooling de média.
• Arquitetura de Cabeça Dupla: O modelo utiliza duas cabeças:
  • Cabeça Latente: Uma cabeça probabilística (um MLP de duas cabeças) que prediz a média ($\mu$) e o desvio padrão ($\sigma$) da distribuição do próximo embedding comprimido. Isso permite a predição auto-regressiva de variáveis latentes densas em vez de valores determinísticos.
  • Cabeça de Linguagem: Prediz a resposta final e, crucialmente, fornece supervisão densa ao predizer os tokens de raciocínio originais correspondentes aos embeddings comprimidos. Isso aproxima uma tarefa de classificação multirrótulo para garantir que o modelo aprenda a semântica dos grupos comprimidos.
• Função de Perda: A perda de treinamento combina a perda de modelagem de linguagem (para predição de tokens) e uma perda latente. A perda latente utiliza uma formulação de soft-MSE que combina o Erro Quadrático Médio com um termo de regularização de entropia. Isso incentiva o modelo a manter a capacidade de exploração (latentes diversas) enquanto se ajusta aos dados, abordando os problemas de subajuste observados com a perda NLL (Negative Log-Likelihood) padrão em conjuntos de dados mais simples.

2. Aprendizado por Reforço (RL) para Exploração e Explotação

Para ir além da imitação do professor de CoT e desbloquear o potencial do raciocínio latente, o CoLaR emprega Otimização de Política Relativa de Grupo (GRPO - Group Relative Policy Optimization).

• Amostragem Probabilística: Como a Cabeça Latente produz uma distribuição de probabilidade, o modelo pode amostrar diversas cadeias de raciocínio latente para a mesma entrada.
• Mecanismo de Recompensa: Para cada questão, um grupo de saídas (cadeias latentes + respostas) é gerado. As recompensas são atribuídas com base na correção da resposta (1 para correto, 0 para incorreto).
• Estratégia de Equilíbrio: O objetivo do GRPO calcula recompensas normalizadas pelo grupo. Crucialmente, a recompensa é calculada pela média sobre o número de tokens/latentes na cadeia. Este design incentiva inerentemente o modelo a explotar caminhos de raciocínio mais curtos e corretos (pois a recompensa por token é maior), enquanto ainda explora diversos caminhos para encontrar soluções corretas.

Principais Contribuições

1. Raciocínio Latente de Velocidade Dinâmica: O CoLaR introduz um framework que prediz auto-regressivamente variáveis latentes que encapsulam a semântica de múltiplos tokens de palavras. Isso permite o ajuste dinâmico da velocidade de raciocínio durante a inferência simplesmente através do prompting de um fator de compressão desejado.
2. Cabeça Latente Probabilística e Integração com RL: Ao contrário dos métodos latentes determinísticos anteriores, o CoLaR emprega uma Cabeça Latente probabilística. Combinado com RL, isso permite que o modelo explore diversos caminhos de raciocínio e explote soluções mais compactas, melhorando significativamente a eficiência sem sacrificar a precisão.
3. Novos Objetivos de Treinamento: A introdução da tarefa de predição de próximo embedding comprimido e a perda de MSE suave permite que o modelo aprenda efetivamente representações densas e equilibre os trade-offs de exploração-explotação.

Resultados Experimentais

Avaliações extensas foram conduzidas em quatro conjuntos de dados de matemática de nível escolar (GSM8k, GSM8k-hard, SVAMP, MultiArith) e no conjunto de dados mais desafiador MATH.

• Desempenho vs. Baselines de Latentes: Nos conjuntos de dados de nível escolar, o CoLaR alcançou uma precisão 14,1% maior do que os baselines de latentes de última geração (ex: Coconut) em proporções de compressão comparáveis.
• Eficiência vs. CoT Explícito: Comparado ao CoT explícito, o CoLaR reduziu o comprimento da cadeia de raciocínio em 53,3% com apenas uma degradação de desempenho de 4,8%.
• Melhorias de RL em Tarefas Desafiadoras: No conjunto de dados MATH, aplicar RL ao CoLaR resultou em um ganho de precisão de 5,36% sobre a versão sem RL, enquanto reduzia o comprimento da cadeia de raciocínio em 82,8%.
• Escalabilidade: Experimentos escalando de 1B a 8B parâmetros confirmaram que os ganhos de eficiência e precisão do CoLaR persistem com maiores capacidades de modelo. O CoLaR-8B-RL superou seu professor de CoT no benchmark GPQA (37,5% vs 35,7%) enquanto encurtava a cadeia de raciocínio em 69%.
• Generalização: O modelo demonstrou uma robusta generalização para fatores de compressão não vistos e conjuntos de dados fora do domínio (ex: MultiArith, GPQA).

Significância

O artigo postula que o CoLaR aborda o gargalo crítico do custo computacional no raciocínio de LLMs ao mudar do processamento esparso de nível de token para o raciocínio em espaço latente denso. Sua significância reside em:

• Adaptabilidade Dinâmica: A capacidade de ajustar a velocidade de raciocínio (fator de compressão) durante a inferência oferece um equilíbrio flexível entre latência e precisão, adequado para variadas restrições de aplicações do mundo real.
• Eficiência sem Sacrifício: Ao alavancar o raciocínio latente probabilístico e RL, o CoLaR alcança reduções substanciais no comprimento da cadeia de raciocínio (até 82,8%) mantendo ou até melhorando a precisão em tarefas complexas, desafiando a noção de que a eficiência deve vir às custas do desempenho.
• Exploração-Explotação no Espaço Latente: O trabalho demonstra que introduzir estocasticidade e aprendizado por reforço nos processos de raciocínio latente permite que os modelos descubram caminhos de raciocínio mais eficientes que as abordagens determinísticas de comprimento fixo deixam passar.

Os autores reconhecem limitações, incluindo o fato de o desempenho atual aproximar (em vez de consistentemente superar) o CoT explícito em alguns benchmarks, e os desafios em generalizar para fatores de compressão não inteiros devido às restrições de tokenização discreta. No entanto, o framework estabelece uma direção promissora para o raciocínio de LLMs eficiente, escalável e adaptável.