Pretraining with Token-Level Adaptive Latent Chain-of-Thought

Este trabalho propõe o "Pretraining com CoT Latente Adaptativa em Nível de Token", uma abordagem que internaliza cadeias de pensamento latentes de comprimento variável durante o pré-treinamento para melhorar a eficiência e o desempenho dos modelos de linguagem sem aumentar o número de parâmetros.

O essencial

Imagine que você está ensinando um estudante muito inteligente, mas com um orçamento limitado de tempo e energia, a escrever um livro.

Normalmente, para fazer esse estudante ficar mais esperto, os cientistas tentam duas coisas:

1. Dar mais livros para ele ler (mais dados de treinamento).
2. Adicionar mais cérebros ao time (aumentar o tamanho do modelo).

O problema é que já estamos ficando sem bons livros (dados de alta qualidade) e adicionar mais cérebros é muito caro e demorado.

A LUMIA Lab (da Universidade Jiao Tong de Xangai) propôs uma ideia brilhante: e se, em vez de dar mais cérebros, nós ensinássemos o mesmo cérebro a pensar mais vezes antes de falar?

Aqui está a explicação da técnica deles, chamada "Cadeia de Pensamento Latente Adaptativa", usando uma analogia do dia a dia:

1. O Problema: O Estudante que Fala Demais ou depressa Demais

Imagine que você está pedindo para esse estudante completar uma frase.

• Se a frase for: "O céu é...", ele sabe que a resposta é "azul" instantaneamente. Se ele gastar 10 minutos "pensando" antes de dizer "azul", ele está desperdiçando energia.
• Se a frase for: "A solução para a crise climática global envolve...", ele precisa pensar muito, considerar várias opções, antes de dar uma resposta boa.

Antes, os modelos de IA tratavam todas as palavras da mesma forma: ou pensavam o mesmo tempo para tudo (desperdício nas fáceis) ou usavam um método rígido que era lento.

2. A Solução: O "Roteiro de Pensamento Invisível"

Os autores criaram um sistema onde o modelo aprende a pensar em silêncio (no "latente", ou seja, dentro da sua mente, sem escrever nada) antes de soltar a próxima palavra.

Eles chamam isso de Cadeia de Pensamento Latente (Latent CoT). É como se o modelo tivesse um bloco de rascunho mental onde ele pode fazer cálculos, revisar ideias e planejar antes de escrever a palavra final.

3. A Mágica: A "Adaptabilidade"

Aqui está o grande diferencial. O modelo não é burro nem teimoso; ele é adaptativo.

• Para palavras fáceis: O modelo olha para a frase, pensa: "Ah, isso é óbvio!" e decide: "Vou pular o bloco de rascunho e escrever direto". (Zero ou um passo de pensamento).
• Para palavras difíceis: O modelo olha, pensa: "Isso é complicado..." e decide: "Preciso fazer 3 ou 4 passos de raciocínio antes de escrever".

Isso é feito por um pequeno "gerente" dentro do modelo (chamado de Router ou Roteador) que decide, a cada palavra, se vale a pena continuar pensando ou se já é hora de falar.

4. A Eficiência: O "Pulo do Gato" (Paralelismo)

O desafio técnico era: se o modelo precisa pensar passo a passo, isso seria muito lento, certo? Como pensar em 1000 palavras, uma por uma?

Os autores inventaram uma forma de fazer isso em paralelo. Imagine uma sala de aula onde, em vez de um aluno pensar de cada vez, todos os alunos pensam ao mesmo tempo, mas cada um para quando acha que já sabe a resposta.

• O aluno que sabe a resposta para rápido.
• O aluno que precisa de ajuda continua pensando.
• O professor (o computador) consegue processar todos eles juntos, economizando tempo e energia.

5. O Resultado: Mais Esperto, Mais Rápido e Mais Barato

Ao treinar o modelo dessa forma (em apenas uma etapa, sem precisar de professores humanos para corrigir cada pensamento), eles conseguiram:

• Melhor desempenho: O modelo ficou mais preciso em tarefas difíceis e em escrever textos.
• Menos custo: Como ele para de pensar nas coisas fáceis, ele gasta menos energia (computação) no total, mesmo tendo a capacidade de pensar muito nas coisas difíceis.
• Sem "gordura": O modelo aprendeu sozinho a saber quando parar de pensar, sem precisar de regras rígidas.

Resumo em uma frase

Em vez de tentar construir um cérebro gigante que é lento e caro, os pesquisadores ensinaram um cérebro normal a pensar rápido nas coisas fáceis e pensar devagar nas coisas difíceis, tudo isso de forma automática e eficiente. É como ter um assistente que sabe exatamente quando precisa de um café extra para resolver um problema complexo e quando pode resolver tudo com um gole de água.

Resumo técnico

Título: Pré-treinamento com Cadeia de Pensamento Latente Adaptativa em Nível de Token (Adaptive Latent CoT)

Autores: Boyi Zeng, Yiqin Hao, He Li, et al. (LUMIA Lab, Shanghai Jiao Tong University, etc.)
Data: Fevereiro de 2026

---

1. O Problema

O avanço dos Grandes Modelos de Linguagem (LLMs) tem sido impulsionado principalmente pela escalagem de parâmetros e dados de treinamento. No entanto, essa abordagem enfrenta dois gargalos críticos:

1. Escassez de Dados: A exaustão de corpora públicos de alta qualidade.
2. Custos de Comunicação e Computação: O aumento do tamanho do modelo gera sobrecargas significativas de comunicação e custos computacionais.

A pesquisa atual busca melhorar a capacidade do modelo sob um orçamento fixo de parâmetros e dados, escalando a computação por token. Métodos existentes, como o compartilhamento recursivo de parâmetros (ex: PonderLM2), muitas vezes exigem iterações complexas (estilo Jacobi) que aumentam drasticamente os FLOPs de treinamento, ou distribuem a computação de forma uniforme, ignorando que a cognição humana é adaptativa (pensamento breve para conceitos simples, prolongado para raciocínio complexo). Além disso, métodos baseados em Chain-of-Thought (CoT) explícito dependem de supervisão externa e espaço de tokens discretos.

2. Metodologia Proposta

Os autores propõem o Adaptive Latent CoT, um framework de pré-treinamento que internaliza o raciocínio em um espaço contínuo latente, permitindo que o modelo gere uma trajetória de CoT de comprimento variável antes de emitir cada token observado.

Componentes Principais:

1. Máscara de Atenção Paralela (Parallel Masking):

   • Para resolver o gargalo de dependência sequencial (onde o estado de um token depende da resolução completa dos passos latentes do token anterior), o método expande a atenção para duas dimensões: posição da sequência ($t$) e passo latente ($k$).
   • Uma máscara de causalidade estrita permite que, para um passo latente fixo $k$, todos os tokens da sequência sejam computados em paralelo. Isso transforma a dependência de $O(L \times K)$ operações sequenciais para $O(K)$ passos sequenciais, aproveitando o paralelismo massivo das GPUs.

2. Mecanismo de Parada Probabilística (Probabilistic Halting):

   • Um Roteador (Router) leve prevê a probabilidade condicional de continuar para o próximo passo latente em cada etapa.
   • Calcula-se a probabilidade de "chegada" ($p_{reach}$) e de "saída" ($p_{exit}$) para cada token.
   • Poda por Limiar (Threshold Pruning): Se a probabilidade de chegar ao próximo passo cair abaixo de um limiar $\tau$, o token é removido do lote para os passos subsequentes, economizando FLOPs.
   • Mistura Baseada em Expectativa: A representação final do token é uma média ponderada dos estados latentes executados, re-alocando a massa de probabilidade residual para o último passo executado para preservar a soma total.

3. Perda Adaptativa Consciente de Correção (Correctness-Aware Adaptive Loss):

   • Para evitar computação desnecessária, o modelo é penalizado se continuar gerando passos latentes quando já está altamente confiante no token alvo (alta probabilidade $p_{target}$).
   • A perda adiciona um termo que incentiva a parada precoce para tokens fáceis, reduzindo os FLOPs totais tanto no treinamento quanto na inferência.

4. Fluxo de Trabalho:

   • Treinamento: Um único estágio (one-stage) em texto geral. O Roteador e o LM são treinados conjuntamente.
   • Inferência: O modelo itera passos latentes até que a probabilidade de continuação caia abaixo do limiar, adaptando o "tempo de pensamento" por token.

3. Contribuições Chave

• Internalização do CoT: Move o raciocínio do espaço de tokens discretos (que consome tokens de saída) para um espaço latente contínuo, sem expandir o vocabulário ou o contexto.
• Computação Adaptativa Eficiente: Diferente de métodos anteriores que exigem treinamento multi-estágio ou supervisão adicional, este método emerge naturalmente do pré-treinamento em texto geral.
• Redução de FLOPs: Ao contrário de métodos recursivos que aumentam o custo de treinamento, a poda dinâmica e a máscara paralela reduzem a computação efetiva durante o treinamento e a inferência.
• Escalabilidade: Funciona com arquiteturas padrão (LLaMA) sem necessidade de mudanças estruturais complexas nos blocos de atenção.

4. Resultados Experimentais

Os experimentos foram realizados em modelos LLaMA (410M e 1.4B parâmetros) pré-treinados no dataset The Pile.

• Desempenho em Modelagem de Linguagem:
  • O modelo com Adaptive Latent CoT alcançou consistentemente a menor perplexidade em todos os conjuntos de dados (The Pile, WikiText, LAMBADA).
  • Eficiência Computacional: O modelo LLaMA-1.4B com $\ell_{max}=3$ superou o baseline mais forte (PonderLM-2) utilizando menos da metade dos FLOPs de treinamento (7.47 vs 17.47 $\times 10^{20}$ FLOPs).
• Tarefas de Downstream:
  • O método obteve a melhor precisão média em tarefas de 0-shot e 5-shot (incluindo ARC, HellaSwag, PIQA, RACE).
  • Um modelo de 410M com o método proposto superou um modelo LLaMA vanilla de 1.4B (treinado do zero) em precisão média, demonstrando que a computação adaptativa é mais eficiente do que simplesmente aumentar o número de parâmetros.
• Análise de Adaptação:
  • A análise mostrou uma correlação clara: tokens fáceis (alta probabilidade no alvo) executam 0-1 passos latentes, enquanto tokens difíceis (baixa probabilidade) ativam passos latentes mais longos, aproximando-se do limite máximo ($\ell_{max}$).

5. Significado e Conclusão

Este trabalho demonstra que é possível escalar a inteligência dos modelos de linguagem aumentando a computação por token de forma adaptativa e eficiente, sem os custos proibitivos de métodos recursivos anteriores ou a necessidade de dados anotados de CoT.

A principal inovação é a capacidade de aprender uma política de alocação de computação end-to-end durante o pré-treinamento padrão, onde o modelo decide "quanto pensar" para cada token. Isso oferece um caminho promissor para superar os limites de dados e custos de comunicação, permitindo modelos mais capazes e eficientes sob orçamentos fixos de parâmetros.