Resumo Técnico: Transformer de Largura de Banda Total (Full-Bandwidth Transformer)

Declaração do Problema

Transformers autorregressivos computam ao longo de dois eixos: horizontalmente através dos tokens gerados e verticalmente através da profundidade do modelo. Embora a atenção densa forneça um amplo acesso horizontal aos tokens passados, o canal de feedback vertical entre as etapas de decodificação permanece estreito. Na decodificação padrão, o modelo comprime todo o seu estado oculto de $D$ dimensões da camada superior em um único token discreto (carregando no máximo $\log_2 |V|$ bits) para servir como entrada para o próximo passo. Consequentemente, a computação não verbalizada (ativações intermediárias, incerteza, planos parciais) permanece "congelada na profundidade" no cache KV; ela permanece acessível às camadas acima de seu ponto de produção, mas não pode retornar à base da pilha para processamento adicional. Isso força o modelo a verbalizar estados intermediários token por token ou a recomputá-los do zero, limitando a profundidade efetiva de raciocínio por token.

Metodologia: Decodificação de Feedback Latente

Os autores propõem o Full-Bandwidth Transformer, que alarga o canal de feedback inter-etapas alimentando o estado oculto da camada superior anterior de volta à entrada junto com o embedding do token amostrado.

1. Mecanismo de Feedback Latente

Em cada etapa de decodificação $t$, a entrada $u_t$ é formada pela fusão do embedding do token amostrado $e_t$ com o estado oculto da camada superior anterior $h^L_{t-1}$:
$$u_t = e_t \otimes h^L_{t-1}$$
O operador de fusão $\otimes$ é uma Unidade Linear de Portão (GLU):
$$e_t \otimes h^L_{t-1} = (W_U h^L_{t-1}) \odot \sigma(W_G e_t)$$

• Assimetria: O estado oculto ocupa o caminho do valor, enquanto o token atua como um portão multiplicativo. Este design impede que o modelo aprenda a ignorar o caminho do estado (um potencial atalho em uma fusão simétrica), forçando-o a utilizar todo o estado oculto.
• Arquitetura: Isso adiciona apenas duas multiplicações de matriz $D \times D$ por token. A arquitetura padrão do transformer, o layout do cache KV e a pilha de serviço permanecem intactos.
• Efeito: Isso permite que estados não verbalizados reentrem na base da pilha com um orçamento de profundidade renovado. Cada camada, incluindo a mais rasa, pode agora acessar todo o histórico processado por toda a pilha, em vez de apenas pelas camadas abaixo dela.

2. Estratégia de Treinamento Paralelo

Treinar esta recorrência diretamente sacrificaria o teacher forcing paralelo, pois a entrada na posição $t$ depende da saída em $t-1$. Os autores introduzem um objetivo de múltiplas passagens para resolver isso:

• Paralelismo Temporal: Em vez de desenrolar sequencialmente sobre os tokens, o modelo executa $K$ passagens de ida paralelas sobre a sequência.
  • Passagem 1: Passagem direta padrão com embeddings de tokens simples.
  • Passagem $k$ ($k > 1$): Os estados ocultos da Passagem $k-1$ são deslocados para a direita em uma posição, fundidos com os embeddings dos tokens e alimentados de volta na pilha em paralelo.
• Função de Perda: O objetivo combina a perda padrão de Predição do Próximo Token (NTP) da primeira passagem com perdas NTP auxiliares das passagens subsequentes. Os gradientes das passagens posteriores retropropagam para os estados latentes anteriores, incentivando-os a serem entradas reutilizáveis.
• Escalonamento: Para garantir a estabilidade e gerenciar o custo computacional, os autores utilizam um escalonamento progressivo:
  • O treinamento começa com objetivos de passagem única padrão.
  • O feedback latente é introduzido no meio do treinamento (ex: lotes de 2 passagens).
  • Uma pequena fração (ex: 3%) de lotes com mais passagens (ex: 3 passagens) é misturada. Empiricamente, essa pequena fração estabiliza o mapa de feedback aprendido, tornando-o uma contração em direção a um ponto fixo que permanece estável mesmo quando extrapolado para horizontes longos (centenas de passagens) durante a inferência.
• Mix-in de Prefixo: Para lidar com a disparidade de distribuição entre o treinamento (entradas totalmente fundidas) e a inferência (prompt simples + geração fundida), o processo de treinamento randomiza o comprimento do prefixo onde as entradas fundidas começam, correspondendo à estrutura de prompt-então-geração da inferência.

3. Técnicas de Estabilidade

Para evitar a divergência durante a inferência de longo horizonte, os autores empregam:

• Escalonamento de Profundidade: Garantir que a norma do estado da camada superior permaneça $O(1)$ em vez de crescer com a profundidade.
• Vinculação de Pesos (Weight Tying): Vincular os pesos da camada de embedding e da camada de leitura para garantir compatibilidade entre as entradas de token simples e de estado fundido.
• Regularização de Ruído: Adicionar ruído de oscilação (jitter noise) ao estado oculto carregado durante o treinamento para melhorar a robustez.

Principais Resultados

Os autores treinaram modelos de 1 bilhão de parâmetros até 400 bilhões de tokens e os avaliaram contra transformers padrão treinados com 100 bilhões a 1 trilhão de tokens.

• Eficiência de Dados: Transformers de largura de banda total igualam ou aproximam-se do desempenho de transformers padrão treinados com aproximadamente 1,5× a 2× mais tokens. Por exemplo, um modelo de 200 bilhões de tokens de largura de banda total aproxima-se do desempenho de um baseline padrão de 400 bilhões de tokens.
• Ganhos de Desempenho:
  • Perda de Validação: Melhorou em todas as escalas.
  • Raciocínio e Codificação: Em tarefas de geração livre (GSM8K, Math500, HumanEval, MBPP), a decodificação de feedback latente (regimes SOFT e FUSED) supera consistentemente a decodificação padrão com os mesmos pesos.
    • SOFT (Feedback Latente): Melhor para tarefas de matemática, sugerindo que o estado oculto carrega melhor o raciocínio do que os tokens verbalizados.
    • FUSED (Prefill de Múltiplas Passagens): Melhor para tarefas de codificação, sugerindo que representações de prompt refinadas são cruciais.
  • Ajuste de Instrução (Instruction Tuning): Os ganhos persistem após a extensão de contexto longo e o ajuste de instrução.
• Concisão: Em modelos base (pré-ajuste de instrução), o feedback latente produz traços de raciocínio marcadamente mais curtos com precisão igual ou superior. O modelo realiza a computação "silenciosamente" no estado oculto em vez de narrá-la token por token. Esse efeito diminui após o ajuste de instrução, provavelmente porque os dados de ajuste imitam o raciocínio verboso padrão.
• Acessibilidade do Estado: Experimentos de sondagem confirmam que o preenchimento recorrente expõe informações de estado global às camadas rasas (Camada 0) de forma muito mais eficaz do que o preenchimento padrão, validando a afirmação de "largura de banda total".

Significância e Alegações

O artigo afirma que o Full-Bandwidth Transformer oferece um novo eixo de escala para Grandes Modelos de Linguagem: alocar mais computação por token, em vez de apenas mais tokens ou parâmetros.

• Revisitando a Escala: À medida que a disponibilidade de dados únicos de alta qualidade se torna uma restrição, os autores argumentam para extrair mais sinal de aprendizado de cada token. O feedback latente alcança isso ao permitir que o modelo itere sobre seu estado interno sem o gargalo da verbalização.
• Eficiência de Inferência: O método fornece ganhos de desempenho com overhead de decodificação por token desprezível (apenas duas multiplicações de matriz). O computação adicional é incurrida principalmente durante o preenchimento opcional de múltiplas passagens ou no treinamento, não durante a geração sequencial de cada token.
• Minimalismo Arquitetural: O método preserva a arquitetura padrão do transformer e o cache KV, exigindo apenas uma mudança na construção da entrada e um cronograma de treinamento específico.
• Estabilidade: O trabalho demonstra que loops de feedback recorrentes podem ser treinados efetivamente em transformers usando objetivos de múltiplas passagens paralelos, superando a instabilidade frequentemente associada ao treinamento de redes recorrentes.

Os autores concluem que os transformers de largura de banda total representam uma mudança de depender apenas da escala de dados para alavancar a computação baseada em feedback recorrente para melhorar a qualidade da representação e a eficiência do raciocínio.