Resumo Técnico: Transformers Provavelmente Aprendem a Internalizar Cadeia de Pensamento

Declaração do Problema

O prompting de Cadeia de Pensamento (CoT) melhora significativamente a eficiência de amostragem de transformers para tarefas complexas como o aprendizado de paridade $k$, reduzindo a complexidade de exponencial para polinomial no comprimento da entrada. No entanto, a geração explícita de CoT incorre em latência substancial de inferência e sobrecarga computacional devido à geração sequencial de tokens de raciocínio intermediários. A Cadeia de Pensamento Implícita (ICoT) emergiu como uma solução empírica, treinando modelos para internalizar etapas de raciocínio dentro de estados ocultos, removendo progressivamente tokens intermediários durante o ajuste fino. Apesar do sucesso empírico, os fundamentos teóricos da ICoT permanecem obscuros: especificamente, se os modelos podem alcançar a eficiência de amostragem do CoT explícito enquanto eliminam os custos de inferência, e como arquiteturas de múltiplas camadas facilitam essa internalização.

Metodologia

Configuração da Tarefa

Os autores analisam o problema de aprendizado de paridade $k$, um banco de testes canônico para supervisão intermediária. Dado um conjunto de índices secretos $S \subset [n]$ e bits de entrada $b \in \{\pm 1\}^n$, o objetivo é prever a paridade $y = \prod_{j \in S} b_j$. Sem supervisão intermediária, esta tarefa é provadamente difícil para métodos baseados em gradiente de precisão finita usando amostras polinomiais. A tarefa é decomposta em uma árvore binária hierárquica de computações de paridade dupla com profundidade $L = \log_2 k$.

Arquitetura do Modelo

O estudo emprega um transformer simplificado de $L$ camadas com escolhas de design específicas para garantir tratabilidade teórica:

• Atenção: Atenção de cabeça única com uma matriz chave-pergunta reparametrizada $W^{KQ}$, onde apenas o bloco superior $T \times T$ é treinável.
• Máscara Causal: Uma máscara causal personalizada restrita por nível garante que estados intermediários no nível $\ell$ dependam apenas de tokens de níveis estritamente inferiores da árvore de paridade.
• Função de Ligação: Uma função suave e simétrica $\phi$ (por exemplo, $-\cos(\pi t)$) satisfazendo $\phi(0)=-1, \phi(\pm 1)=1$, e derivadas nulas nestes pontos, permitindo expansões de Taylor locais.
• Conexões Portadas: Em vez de conexões residuais padrão, o modelo usa conexões portadas onde vetores de porta $g^{(\ell)}$ são prescritos pela estrutura da árvore de paridade. Isso isola sinais de gradiente para níveis específicos da árvore, prevenindo o colapso de representação onde estados de entrada convergem para valores uniformes.
• Quantização: Pesos de atenção são quantizados para o inteiro mais próximo após cada atualização de gradiente para controlar a propagação de erro e "congelar" camadas previamente treinadas.

Currículo de Treinamento: Log-ICoT

O artigo introduz Log-ICoT, um currículo que remove etapas intermediárias de CoT em blocos geométricos alinhados com os níveis da árvore de paridade, em vez de um token de cada vez (como na ICoT padrão).

• Estágio 1: Supervisão completa de CoT (todos os nós intermediários visíveis).
• Estágio $t$ ($2 \le t \le L$): Os primeiros $k(1 - 2^{-(t-1)})$ tokens intermediários são substituídos por preenchimento (zeros), enquanto nós de nível superior permanecem visíveis.
• Estágio $L$: Todos os tokens intermediários são preenchidos; o modelo deve prever a saída final usando apenas bits de entrada e estados ocultos internalizados.
  Isso reduz o número de estágios de treinamento de linear em $k$ (ICoT padrão) para logarítmico, $L = \log_2 k$.

Contribuições Principais

1. Primeira Análise Teórica da ICoT: O artigo fornece a primeira prova rigorosa de que a ICoT preserva a eficiência de amostragem do CoT explícito. Demonstra que um transformer de $L$ camadas treinado sob o currículo Log-ICoT resolve paridade $k$ com $\text{poly}(n)$ amostras e $\log_2 k$ estágios de treinamento.
2. Currículo Log-ICoT: Os autores propõem um currículo que remove tokens de raciocínio em incrementos geométricos. Isso corresponde à estrutura recursiva da árvore de paridade, reduzindo estágios de treinamento de $O(k)$ para $O(\log k)$ enquanto mantém eficiência de inferência (passada direta única).
3. Garantia de Convergência de Múltiplas Camadas: O Teorema 1 prova que um transformer de múltiplas camadas pode aprender paridade eficientemente, estendendo garantias anteriores de camada única. A prova depende de:
   • Conexões Portadas: Para prevenir colapso de representação e isolar sinais de gradiente para níveis específicos da árvore.
   • Quantização Inteira: Para travar pesos de camadas treinadas, permitindo análise por estágio tratando camadas anteriores como fixas.
   • Máscara Causal Personalizada: Para garantir que sinais de gradiente dependam apenas dos nós filhos corretos.
4. Controle de Propagação de Erro: A análise mostra que erros de aproximação em estados ocultos decaem exponencialmente, garantindo que o erro de previsão final permaneça negligenciável ($\exp(-\Omega(n^{\epsilon/16}))$) apesar da natureza sequencial dos estágios de treinamento.

Resultados

• Teórico: O teorema principal estabelece que, com um tamanho de lote por estágio $B = \Omega(n^{2+\epsilon})$, o modelo treinado alcança precisão quase perfeita no conjunto de teste com alta probabilidade. A complexidade de amostragem é polinomial em $n$, correspondendo ao CoT explícito, enquanto a inferência requer apenas uma passada direta (vs. $O(k)$ etapas sequenciais para CoT autoregressivo).
• Empírico: Experimentos em um transformer de 4 camadas com $n=30$ e $k=16$ ($L=4$) confirmam a teoria.
  • Perda de Validação: A perda cai rapidamente em cada transição de estágio, com picos transitórios à medida que o modelo internaliza novos níveis, eventualmente alcançando 100% de precisão com todos os tokens de CoT preenchidos.
  • Mapas de Atenção: A visualização mostra que, no estágio final, a atenção de cada camada concentra-se nitidamente em exatamente duas posições-chave (os filhos do nó de paridade consultado), confirmando a previsão teórica de concentração de softmax.

Significado e Afirmações

O artigo afirma fornecer uma base teórica rigorosa para a Cadeia de Pensamento Implícita, preenchendo a lacuna entre a expressividade do raciocínio explícito e a aprendibilidade do raciocínio internalizado. Ao provar que transformers de múltiplas camadas podem internalizar processos complexos de raciocínio sem sacrificar a eficiência de amostragem, o trabalho aborda os altos custos de inferência associados ao CoT explícito.

Os autores enfatizam que sua análise depende de simplificações arquitetônicas específicas (portas prescritas, matrizes de valor fixas, quantização inteira) para tornar a dinâmica de treinamento de múltiplas camadas tratável. Eles reconhecem que aplicar Log-ICoT a Grandes Modelos de Linguagem (LLMs) práticos sem a estrutura hierárquica explícita de tarefas de paridade permanece um desafio, exigindo heurísticas para o escalonamento do currículo. O trabalho se posiciona como um passo em direção à compreensão dos mecanismos de auto-distilação e internalização de raciocínio em arquiteturas profundas.