Resumo Técnico: Transformers de Previsão de Próximo Latente Aprendem Modelos de Mundo Compactos

Declaração do Problema

As arquiteturas padrão de Transformers substituem a recorrência por auto-atenção e um mecanismo de memória que escala com o comprimento da sequência. Embora isso permita paralelização excepcional e poder preditivo, remove o viés indutivo inerente para comprimir o histórico em estados latentes compactos com regras de transição consistentes. Consequentemente, os Transformers frequentemente aprendem atalhos complexos e específicos de tarefas que se ajustam bem aos dados de treinamento, mas generalizam mal, falhando em formar "modelos de mundo" coerentes ou estados de crença necessários para planejamento e raciocínio de longo horizonte. O objetivo padrão de previsão do próximo token não garante que os estados ocultos internos do modelo evoluam como estatísticas suficientes para o histórico passado, nem assegura que esses estados possam prever recursivamente observações futuras.

Metodologia: Previsão de Próximo Latente (NextLat)

Os autores introduzem a Previsão de Próximo Latente (NextLat), um framework que aprimora o treinamento padrão de próximo token com previsões auto-supervisionadas no espaço latente. A ideia central é treinar o Transformer para aprender representações latentes (especificamente, os estados ocultos da camada final $h_t$) que sejam preditivos do próximo estado latente dado o próximo token.

Fundamentação Teórica

O artigo estabelece que, se um modelo otimizar tanto a Consistência do Próximo Token (prever $X_{t+1}$ a partir de $h_t$) quanto a Consistência de Transição (prever $h_{t+1}$ a partir de $h_t$ e $X_{t+1}$), os estados ocultos $h_t$ são teoricamente garantidos a convergir para estados de crença. Um estado de crença é definido como uma estatística suficiente do histórico $X_{1:t}$ para prever o futuro $X_{t+1:T}$. Ao contrário do treinamento autorregressivo padrão, que permite consultas ad-hoc de tokens passados sem comprimir informações, o NextLat impõe um viés indutivo recorrente onde o estado latente deve conter todas as informações necessárias para transicionar para o próximo estado e prever o próximo token.

Implementação

O NextLat opera adicionando um objetivo auxiliar leve ao loop de treinamento padrão do Transformer:

1. Modelo de Dinâmica Latente ($p_\psi$): Um simples Perceptron Multicamadas (MLP) é treinado para prever o próximo estado oculto $\hat{h}_{t+1}$ dado o estado oculto atual $h_t$ e o próximo token $X_{t+1}$.
2. Função de Perda: A perda total combina:
   • Perda de Próximo Token ($L_{next-token}$): Entropia cruzada padrão para previsão de token.
   • Perda de Próximo Estado Oculto ($L_{next-h}$): Uma perda Smooth L1 que regressa o estado previsto $\hat{h}_{t+1}$ contra o próximo estado verdadeiro $h_{t+1}$ (com gradiente interrompido aplicado ao alvo para evitar colapso representacional).
   • Perda de Divergência KL ($L_{KL}$): Um objetivo complementar que garante que a distribuição de tokens prevista a partir do estado latente previsto corresponda à distribuição do estado latente verdadeiro, atuando como um sinal de supervisão suave.
3. Supervisão Multi-Etapas: O modelo pode ser treinado para desenrolar recursivamente a dinâmica latente por $d$ etapas, fornecendo sinais de gradiente mais ricos para dependências de longo alcance.

Crucialmente, o NextLat deixa a arquitetura do Transformer, o procedimento de inferência e a eficiência do treinamento paralelo inalterados. O modelo de dinâmica latente é usado apenas durante o treinamento para moldar representações; durante a inferência, o modelo decodifica autorregressivamente como de costume.

Contribuições Principais

1. Garantia Teórica: O artigo prova que otimizar para consistência de próximo token e próximo latente molda comprovadamente as representações do transformer em estados de crença, uma propriedade não garantida pela previsão de próximo token isoladamente.
2. Framework Prático: O NextLat introduz um objetivo auxiliar computacionalmente eficiente que injeta viés indutivo recorrente sem alterar o treinamento paralelo ou a estrutura de inferência do Transformer.
3. Decodificação Auto-Especulativa de Comprimento Variável: Ao aprender dinâmicas latentes coerentes, o NextLat permite previsão recursiva no espaço latente. Isso permite decodificação auto-especulativa de comprimento variável, onde o modelo pode rascunhar tokens além do horizonte fixo de treinamento, acelerando significativamente a inferência.

Resultados Experimentais

Os autores avaliam o NextLat em quatro domínios: modelagem de mundo, raciocínio, planejamento e modelagem de linguagem.

• Modelagem de Mundo (Passeios de Táxi em Manhattan): O NextLat aprende mapas internos significativamente mais consistentes com o verdadeiro modelo de mundo (topologia das ruas) em comparação com as linhas de base padrão GPT, MTP e JTP. Ele alcança maior compressão de sequência (0,71 vs. 0,65 para o GPT) e uma classificação latente efetiva muito menor (52,7 vs. 160,1 para o GPT), indicando representações mais compactas e eficientes.
• Raciocínio (Countdown): Na tarefa de raciocínio matemático, o NextLat supera substancialmente as linhas de base (54,8% de precisão em $d=1$ vs. 33,1% para o GPT). Ele demonstra planejamento de antecipação superior, reduzindo "compromissos arrependidos" onde os modelos forçam equações finais inválidas para corresponder a um alvo.
• Planejamento (Grafos Path-Star): O NextLat mantém taxas de resolução próximas a 100% em várias topologias de grafos, enquanto linhas de base como BST e JTP falham em grafos maiores ($G_{7,7}$). Isso sugere que o NextLat evita a aprendizagem de atalhos "Clever Hans" comum na previsão no espaço de tokens.
• Modelagem de Linguagem (FineWeb-Edu):
  • Precisão a Montante: O NextLat preserva a perplexidade do próximo token melhor do que MTP e JTP, enquanto alcança ganhos modestos na precisão zero-shot a montante (59,21% em média vs. 58,82% para o GPT).
  • Velocidade de Inferência: Aproveitando a decodificação auto-especulativa de comprimento variável, o NextLat alcança até 3,3× inferência mais rápida em comparação com a amostragem autorregressiva padrão. Essa aceleração persiste mesmo ao rascunhar além do horizonte de treinamento, uma capacidade que MTP e JTP não possuem devido aos seus horizontes fixos no espaço de tokens.

Significado e Alegações

O artigo posiciona o NextLat como um paradigma simples, porém eficaz, para aprender representações compactas, preditivas e generalizáveis em Transformers. Seu significado principal reside em unificar o viés indutivo recorrente de modelos de mundo compactos com a escalabilidade e paralelismo dos Transformers.

Os autores afirmam que o NextLat oferece uma solução para a natureza "miópica" da previsão de próximo token, incentivando os modelos a formar estados internos estruturados e preditivos, em vez de depender de correlações superficiais de tokens. Ao permitir decodificação especulativa de comprimento variável sem alterações arquiteturais, o NextLat fornece um caminho prático para inferência mais rápida e melhor raciocínio de longo horizonte. O artigo conclui que o NextLat é um passo em direção a objetivos de treinamento que dotam modelos autorregressivos de representações mais simples, mais compactas e, portanto, mais generalizáveis de sistemas complexos.

Os autores permanecem modestos quanto a direções futuras, observando que não exploraram arquiteturas complexas de dinâmica latente (usando apenas MLPs simples) e que o design ótimo do objetivo (por exemplo, a necessidade da perda KL em escalas maiores) requer estudo sistemático adicional. Eles também reconhecem que, embora o NextLat melhore as representações, a complexidade de circuito subjacente do Transformer de profundidade fixa permanece limitada pela classe TC0, embora o modelo de dinâmica latente co-treinado (RNN) pareça capaz de generalizar além desses limites.