LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation

1. O Problema

Os Transformers em Loop (ou recorrentes) são uma classe de modelos eficiente que compartilha parâmetros através de múltiplas iterações (loops) de uma mesma pilha de blocos. Estudos recentes indicam que essas arquiteturas possuem um viés indutivo para raciocínio latente, onde o modelo internaliza habilidades de raciocínio através de iterações ocultas, sem necessidade de prompts explícitos de "cadeia de pensamento".

No entanto, a abordagem atual enfrenta um limitação fundamental:

• Rigidez de Profundidade: Os modelos existentes são treinados e avaliados com um número fixo de iterações (loops).
• Colapso de Representação: Quando esses modelos são avaliados com um número de loops diferente do treinado (seja menor ou maior), suas representações internas tendem a colapsar ou estagnar, resultando em degradação severa de desempenho.
• Ineficiência Computacional: Isso impede o uso flexível do orçamento computacional. O usuário não pode simplesmente "pedir" mais ou menos processamento para uma tarefa específica sem re-treinar o modelo, desperdiçando o potencial de adaptação dinâmica que os loops oferecem.

O objetivo central é criar um modelo de linguagem em loop que seja elástico em profundidade, capaz de operar eficientemente sob orçamentos computacionais variáveis (número de loops $M$) sem re-treinamento, mantendo representações informativas mesmo em profundidades reduzidas.

2. Metodologia: LoopFormer

O LoopFormer é uma arquitetura de Transformer em loop projetada para suportar inferência com profundidade elástica. A abordagem é inspirada em modelos de difusão e Equações Diferenciais Ordinárias Neurais (Neural ODEs), tratando a refinamento iterativo de representações como uma trajetória no espaço de estados.

Arquitetura e Condicionamento

Diferente de modelos anteriores que usam apenas o índice do loop como tempo, o LoopFormer condiciona cada iteração em duas variáveis:

1. Tempo Normalizado ($t$): A posição cumulativa na trajetória (de 0 a 1).
2. Tamanho do Passo ($\Delta t$): O "salto" ou granularidade da iteração atual.

Esses parâmetros são codificados via embeddings de frequência senoidal-cosseno e projetados por pequenos MLPs para gerar um vetor de condicionamento $c_i$. Este vetor modula os blocos do Transformer através de:

• Escalas (Scaling): Ajusta as camadas de normalização RMSNorm.
• Portões (Gating): Controla a contribuição dos resíduos das camadas de Atenção Multi-cabeça (MHSA) e Feed-Forward (FFN).

Isso permite que o modelo saiba "onde" está na trajetória e "quão fino" é o passo atual, permitindo que trajetórias mais curtas (passos maiores) aproximem a trajetória completa (passos menores).

Treinamento: Consistência de Shortcut

Para garantir que trajetórias de diferentes comprimentos sejam consistentes, o LoopFormer utiliza um esquema de treinamento inovador chamado Shortcut-Consistency:

• Objetivo: Durante o treinamento, além de calcular a perda para a trajetória máxima (com $L$ loops), o modelo amostra aleatoriamente uma trajetória "shortcut" mais curta (com $S < L$ loops).
• Loss de Consistência: Uma perda adicional é aplicada para alinhar as logits finais da trajetória curta com as da trajetória longa completa. Isso é feito usando um stop-gradient na trajetória longa, efetivamente realizando uma auto-distilação dentro do próprio loop.
• Resultado: O modelo aprende que, independentemente de quantos passos sejam dados, o estado final deve convergir para a mesma representação de alta qualidade.

3. Contribuições Principais

1. Formulação de Trajetórias Moduladas: Introdução de uma classe de Transformers em loop que condiciona cada passagem no tempo interno e no tamanho do passo, permitindo que o modelo evolua consistentemente em trajetórias de diferentes comprimentos.
2. Protocolo de Treinamento de Consistência: Um método que alinha representações de rotas curtas e longas, habilitando a inferência com orçamento computacional variável (profundidade elástica) sem re-treinamento.
3. Análise de Dinâmica de Representação: Demonstração de que, ao contrário de modelos de "saída antecipada" (early-exit) que sofrem de estagnação, o LoopFormer mantém uma evolução dinâmica das representações (curvatura, entropia e anisotropia) que se adapta à profundidade computacional.
4. Desempenho Adaptativo: Evidência empírica de que o modelo escala suavemente com o orçamento computacional, mantendo alta utilidade mesmo sob orçamentos agressivos.

4. Resultados Experimentais

Os experimentos foram conduzidos comparando o LoopFormer com Transformers não-loopados (Base), Transformers em loop padrão e variantes com saída antecipada (Early-Exit), utilizando orçamentos de FLOPs equivalentes.

• Desempenho em Perplexidade e Raciocínio:

  • O LoopFormer supera significativamente as variantes de loop padrão e os modelos de saída antecipada em tarefas de raciocínio zero-shot (como COPA, HellaSwag, PIQA, ARC).
  • Sob orçamentos reduzidos (ex: 6x ou 12x FLOPs em vez de 24x), o LoopFormer mantém um desempenho próximo ao de modelos não-loopados de profundidade equivalente, enquanto os baselines de loop colapsam.
  • À medida que o orçamento aumenta, o desempenho melhora suavemente, sem o colapso observado em outros métodos.

• Análise de Representação:

  • Métricas como Curvatura, Anisotropia e Entropia de Prompt mostram que o LoopFormer continua a refinar suas representações ao longo dos loops.
  • Em contraste, modelos baselines (como TMLT-EE ou Naive-Loop-EE) exibem trajetórias planas nessas métricas, indicando que as representações param de evoluir (estagnação) após poucos passos.
  • O CKA (Centered Kernel Alignment) entre passos de loop mostra que o LoopFormer tem uma divergência progressiva (drift), enquanto os baselines permanecem altamente similares (colapso).

• Escolha de Trajetórias:

  • O estudo mostrou que a escolha do agendamento de passos ($\Delta t$) importa. Trajetórias que alocam passos maiores no início e passos mais finos no final tendem a performar melhor, sugerindo que o modelo aprende a refinar detalhes finais com maior precisão.

5. Significado e Conclusão

O LoopFormer representa um avanço significativo na viabilidade prática de Transformers em Loop para modelagem de linguagem. Ao resolver o problema do colapso de representação em profundidades variáveis, o trabalho abre caminho para:

• Modelos de Linguagem Conscientes de Orçamento: Sistemas que podem ajustar dinamicamente sua profundidade de inferência com base na complexidade da tarefa ou nas restrições de hardware, sem perda de qualidade.
• Raciocínio Latente Eficiente: A confirmação de que o raciocínio em modelos de linguagem pode ser internalizado e refinado iterativamente de forma robusta, independentemente do número de passos computacionais disponíveis.
• Eficiência de Treinamento: A capacidade de treinar um único modelo que serve para múltiplos cenários de inferência, eliminando a necessidade de múltiplos modelos para diferentes níveis de profundidade.

Em suma, o LoopFormer transforma a arquitetura em loop de uma curiosidade teórica com limitações práticas em uma ferramenta robusta para raciocínio adaptativo e eficiente.