JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation

Resumo Técnico: JEPA-Reasoner: Desacoplando o Raciocínio Latente da Geração de Tokens

Definição do Problema
Os Grandes Modelos de Linguagem (LLMs) atuais operam sob um paradigma autorregressivo unificado onde o raciocínio de alto nível e a geração de tokens de baixo nível estão inextricavelmente acoplados. Neste framework, o modelo prevê o próximo token com base em um histórico de tokens discretos previamente amostrados. Este acoplamento cria um gargalo fundamental: qualquer erro localizado na seleção de tokens (por exemplo, um deslize gramatical ou uma escolha de palavra subótima) contamina a janela de contexto, corrompendo diretamente a trajetória de todos os passos de raciocínio subsequentes. Isso leva a erros cumulativos em tarefas de longo horizonte. Além disso, as tentativas existentes de raciocínio no espaço latente (ex: COCONUT, modelos de Profundidade Recorrente) permanecem limitadas por este acoplamento, muitas vezes intercalando vetores latentes com tokens ou dependendo de uma máscara causal que impede que os tokens iniciais sejam informados por estados de raciocínio futuros. Esses modelos também têm dificuldade em representar incerteza, frequentemente forçando um colapso em uma única trajetória de raciocínio devido à natureza discreta do acoplamento ao nível do token.

Metodologia
Os autores propõem o JEPA-Reasoner, um novo paradigma arquitetural que desacopla totalmente o raciocínio no espaço latente da geração de tokens. A arquitetura consiste em dois módulos distintos:

1. JEPA-Reasoner (O Raciocinador): Este módulo opera inteiramente dentro de um espaço latente contínuo e normalizado. Ele utiliza uma Arquitetura Preditiva de Embedding Conjunto (JEPA) adaptada para geração autorregressiva no espaço latente.

   • Mecanismo: Ele recebe um input, codifica-o e utiliza blocos Transformer modificados para prever a próxima matriz latente que representa os passos de raciocínio subsequentes.
   • Restrição: Crucialmente, o Raciocinador não projeta saídas em probabilidades de vocabulário. Em vez disso, a matriz latente gerada é normalizada (via uma camada híbrida de normalização RMS e L2) para uma hiperesfera unitária e retorna como input para o próximo passo.
   • Treinamento: É treinado via Treinamento de Auto-Supervisão (SST) para prever representações latentes fornecidas por um codificador alvo (cujos pesos são atualizados via Média Móvel Exponencial). O objetivo é uma perda de similaridade de cosseno escalonada, incentivando o modelo a aprender diferenças angulares e dinâmica de características consistentes, em vez de predição de tokens discretos.

2. Talker (O Gerador): Este é um módulo separado, baseado em Transformer e leve, treinado independentemente para reconstruir texto legível por humanos.

   • Mecanismo: Ele recebe toda a cadeia de raciocínio latente gerada pelo JEPA-Reasoner e reconstrói a sequência de tokens.
   • Variantes: O artigo introduz o Mono-Talker (decoders apenas, para tarefas sem contexto) e o Dual-Talker (encoders e decoders, para tarefas com consciência de contexto).
   • Dependência: O Talker é treinado exclusivamente para reconstrução baseada na saída do Reasoner; ele não pode fazer predições de forma independente.

Principais Contribuições
O artigo identifica três propriedades emergentes desta arquitetura desacoplada:

1. Contenção de Erro Matemático: Os autores fornecem uma fatoração probabilística mostrando que, como a cadeia de raciocínio ($R$) é gerada independentemente da amostragem de tokens ($X$), erros ao nível do token não possuem caminho matemático para influenciar a lógica do Reasoner ($\frac{\partial r_{t+1}}{\partial \hat{x}_t} = 0$). Isso isola o motor de raciocínio de falhas de amostragem discreta.
2. Guia Contínuo: O gerador (Talker) tem acesso a toda a trajetória de raciocínio sem perdas antes de gerar qualquer token, ao contrário dos modelos acoplados onde a geração é intercalada com o raciocínio.
3. Representação de Incerteza: O Reasoner pode produzir "vetores latentes mistos" que codificam múltiplas trajetórias de raciocínio simultaneamente. Esses vetores atuam como combinações lineares de embeddings de vocabulário, permitindo que o modelo mantenha múltiplas hipóteses sem o overhead de beam search ou múltiplos decodes.

Resultados Experimentais
Os autores conduziram experimentos controlados em tarefas de linguagem sintética e natural:

• Tarefas Sintéticas (Busca em Árvore & CFG):
  • Em uma tarefa de busca em árvore binária, o modelo alcançou 99,87% de precisão. A análise de PCA confirmou que os vetores latentes preditos formaram uma nuvem contínua dentro do espaço ocupado pelos embeddings de vocabulário discretos, validando a hipótese do "estado latente misto".
  • Em tarefas de geração de Gramática Livre de Contexto (CFG), o JEPA-Reasoner demonstrou robustez superior tanto ao ruído de tokens de entrada (substituição de 0–30%) quanto ao ruído Gaussiano no espaço latente, comparado a um baseline Transformer padrão e ao modelo acoplado COCONUT.
• Tarefas de Linguagem Natural (GSM8K):
  • Um JEPA-Reasoner de 0.9B de parâmetros (combinado com um Talker de 198M) foi treinado com os mesmos dados e hiperparâmetros que um baseline Transformer acoplado de 0.9B.
  • Desempenho: O JEPA-Reasoner alcançou 51,9% de precisão no benchmark GSM8K de 8-shot, uma melhoria de 149,5% sobre o baseline (20,8%).
  • Escalabilidade: Enquanto o baseline apresentou um "platô de raciocínio" (melhoria de apenas 0,1% de 5-shot para 8-shot), o JEPA-Reasoner utilizou exemplos adicionais para alcançar uma melhoria de 131,7% no desempenho, indicando escalonamento de contexto eficaz.
  • Comparação: O modelo de 0.9B superou um modelo CoT de 1B (Llama 3.2) e um modelo de Profundidade Recorrente de 3.5B (Huginn-0125) no mesmo benchmark.

Significância e Alegações
O artigo posiciona o JEPA-Reasoner não meramente como um aumento de desempenho através de escala, mas como uma investigação arquitetural sobre os benefícios do desacoplamento. Os autores alegam que a transição de um paradigma acoplado para um desacoplado gera ganhos significativos em robustez e profundidade de raciocínio sem exigir aprendizado por reforço sofisticado ou receitas de treinamento industriais.

O trabalho sugere que o "platô de raciocínio" observado em modelos de pequeno porte baseados em tokens pode ser superado ao mover o processo de raciocínio para um espaço latente contínuo e abstrato. Ao isolar o motor de raciocínio do processo de amostragem de tokens discretos, a arquitetura fornece uma base mais robusta para o raciocínio complexo, permitendo a contenção de erros e a representação de incerteza que são difíceis de alcançar em modelos autorregressivos padrão. Os autores concluem que esta abordagem desacoplada representa um novo caminho para o design de modelos de fundação.