Resumo Técnico: MUX - Raciocínio Contínuo via Tokens Multiplexados

1. Definição do Problema

Modelos de linguagem autorregressivos modernos resolvem problemas complexos (matemática, codificação, senso comum) articulando etapas intermediárias de raciocínio em linguagem natural (Cadeia de Pensamento ou Chain-of-Thought). Embora eficaz, essa abordagem discreta é computacionalmente ineficiente. Cada etapa de raciocínio transmite apenas um subtoken, criando um gargalo estrito de largura de banda. Além disso, muitos tokens de raciocínio são redundantes, espelhando padrões de comunicação humana em vez de computação otimizada.

O raciocínio latente contínuo oferece uma alternativa onde os modelos preveem vetores contínuos em vez de subtokens, permitindo teoricamente uma maior largura de banda ao codificar múltiplos subtokens simultaneamente via superposição. No entanto, os métodos de raciocínio latente existentes enfrentam obstáculos significativos:

• Métodos globais (que supervisionam apenas a resposta final ou o ponto final da trajetória) frequentemente levam a "comportamentos de atalho", onde os tokens latentes tornam-se placeholders informativos.
• Métodos locais (que supervisionam cada token latente contra um intervalo de raciocínio discreto) frequentemente exigem arquiteturas auxiliares complexas (ex: decodificadores autorregressivos) ou falham em manter a diversidade de hipóteses necessária para tarefas de busca.

O desafio central é definir um alvo de supervisão para o raciocínio latente contínuo que seja fiel ao raciocínio discreto, evite o colapso latente e suporte a exploração paralela sem complexidade arquitetônica excessiva.

2. Metodologia: MUX

Os autores propõem o MUX, um método de destilação local que treina modelos de linguagem para realizar raciocínio latente contínuo usando tokens multiplexados.

Conceito Central

O MUX trata a tarefa de destilação local como o aprendizado de um sinal contínuo de dimensão fixa que codifica um sinal categórico de comprimento variável (um intervalo de etapas de raciocínio discretas). Inspirando-se na multiplexação por divisão de código (code-division multiplexing) de sistemas de comunicação, o MUX representa um intervalo de subtokens discretos como uma superposição linear ponderada de suas codificações one-hot no espaço do vocabulário.

Implementação Técnica

1. Alvos Multiplexados: Para um intervalo de raciocínio discreto $r_i = (r^1_i, \dots, r^{S_i}_i)$, o alvo de aprendizado é construído como:
   $$\text{mux}(r_i) := \sum_{j=1}^{S_i} \alpha^{(i)}_j \cdot \text{onehot}(r^j_i)$$
   onde $\alpha^{(i)}_j$ são pesos posicionais positivos normalizados.
2. Ponderação Posicional: Para garantir que a superposição seja sem perdas (ou seja, o intervalo original possa ser recuperado de forma única), os autores propõem esquemas de ponderação específicos:
   • Geométrico: $w_j = \rho^{j-1}$ (decaimento exponencial).
   • Sinusoidal: Baseado em funções seno que atingem o pico próximo ao fim do intervalo.
   • Rotatório: Baseado em componentes de cosseno médios análogos a embeddings de posição rotatórios.
   • Garantia Teórica: Essas ponderações satisfazem uma "condição de separação de soma de subconjuntos", garantindo que diferentes conjuntos de posições produzam massas totais distintas para qualquer subtoken, evitando ambiguidade mesmo se um subtoken aparecer múltiplas vezes.
3. Objetivo de Treinamento: O modelo é treinado para corresponder a esses alvos multiplexados usando um cabeçalho linear-softmax e uma perda de divergência KL ($L_{local}$). Isso é combinado com uma perda padrão de previsão de resposta ($L_{answer}$) e uma perda opcional de nível de trajetória ($L_{global}$) para alinhar as características ocultas com modelos de raciocínio discretos.
4. Inferência: Durante a inferência, o modelo realiza a autorregressão de tokens latentes no espaço oculto. A projeção do vocabulário é usada apenas para supervisão durante o treinamento, não para gerar os tokens de saída finais.

3. Principais Contribuições

(1) Raciocínio Latente via Tokens Multiplexados

O MUX introduz um simples framework de destilação local onde tokens latentes representam intervalos de raciocínio discretos via superposições lineares ponderadas. Ao contrário de métodos locais anteriores, ele não requer um decodificador autorregressivo auxiliar ou alvos de cache comprimido, baseando-se, em vez disso, em projeções do espaço do vocabulário.

(2) Multiplexação Sem Perdas

O artigo fornece provas teóricas de que ponderações posicionais específicas (geométrica, sinusoidal, rotatória) garantem a multiplexação sem perdas. Isso significa que o intervalo de raciocínio discreto é totalmente preservado na representação contínua. Os autores provam que essa ausência de perdas evita o colapso latente (onde os tokens se tornam semanticamente homogeneizados), um modo de falha comum em métodos de raciocínio latente global.

(3) Busca Paralela via Superposição

Teoricamente, os autores demonstram que tokens multiplexados podem implementar naturalmente uma busca em largura (BFS) paralela. Como os tokens estão em superposição, um único token latente pode representar e atualizar um conjunto inteiro de hipóteses (uma fronteira de busca) simultaneamente. Isso permite que o modelo explore múltiplos caminhos de resolução de problemas em paralelo dentro de um framework de supervisão serial.

(4) Superioridade Empírica

Em 32 configurações de avaliação (abrangendo quatro modelos de linguagem, dois corpora de treinamento e quatro conjuntos de teste), o MUX supera as fortes linhas de base tanto em raciocínio discreto quanto contínuo. Notavelmente, o MUX frequentemente supera a linha de base de Cadeia de Pensamento SFT-CoT da qual é destilado, apesar de usar significativamente menos tokens (aproximadamente 2,4x a 5,9x menos).

4. Resultados Experimentais

• Raciocínio Matemático: No GSM8K-AUG, GSM8K-AUG-NL, SVAMP, GSM-Hard e MultiArith, o MUX alcançou o melhor desempenho de raciocínio latente em modelos GPT-2, LLaMA 3.2 1B, 3B e 3.1 8B. Ele frequentemente superou o SFT-CoT, demonstrando que o raciocínio superposto compacto pode generalizar melhor do que os traços discretos originais.
• Benchmarks de Busca: Nas tarefas Minimum Nonnegative Sum (MNNS) e Game of 24, o MUX superou significativamente as linhas de base (incluindo Coconut e CoT2), validando a afirmação teórica de que tokens multiplexados permitem capacidades de busca paralela.
• Estudos de Ablação:
  • Perda Local: Remover a perda global ($\gamma=0$) e depender exclusivamente da perda local multiplexada ainda produziu resultados superiores em comparação ao SIM-CoT, confirmando a eficácia do alvo local.
  • Ponderação: Ponderações sem perdas (Geométrica, Sinusoidal, Rotatória) superaram consistentemente a ponderação uniforme com perdas, embora a diferença tenha sido modesta em conjuntos de dados altamente estruturados.
  • Fragmentação (Chunking): A fragmentação aleatória de intervalos discretos em tokens latentes apresentou o melhor desempenho, atuando como um aumento de dados estruturado.

5. Significância e Alegações

O artigo afirma que a superposição sem perdas como um alvo de aprendizado local constitui uma condição suficiente para alcançar um raciocínio latente contínuo forte, eficiente e interpretável.

• Eficiência: O MUX alcança alto desempenho com menos etapas de raciocínio do que o CoT discreto, abordando o gargalo da largura de banda computacional.
• Simplicidade: O método evita a complexidade arquitetônica de abordagens de destilação local anteriores (ex: decodificadores auxiliares) enquanto mantém ou melhora o desempenho.
• Interpretabilidade: Análises de sondagem (probing) confirmam que os tokens latentes aprendidos codificam conteúdo de raciocínio fiel e interpretável, combatendo a noção de que o raciocínio latente é apenas um placeholder ininterpretável.
• Paralelismo: Os resultados sugerem que a dinâmica de busca paralela pode emergir naturalmente de supervisão serial ao usar alvos multiplexados, ofereando um novo caminho para melhorar o planejamento e a busca em modelos de linguagem.

Os autores concluem que alvos locais simples e interpretáveis podem tornar o raciocínio latente mais forte e fácil de treinar, diminuindo a lacuna entre a eficiência dos espaços contínuos e a fidelidade do raciocínio discreto.