Transformadores full-bandwidth reciccam estado oculto entre tokens

Problema e proposta

O artigo apresenta o full-bandwidth transformer, uma arquitetura que busca ampliar o fluxo de informação entre etapas consecutivas da geração autorregressiva. Em um transformador convencional, o cálculo ocorre horizontalmente entre tokens e verticalmente ao longo das camadas. Embora a atenção densa permita a cada token acessar amplamente o contexto anterior, o retorno vertical entre passos de decodificação é limitado: apenas o token amostrado volta à entrada da pilha, enquanto o estado oculto da camada superior é descartado.

A proposta adiciona um mecanismo chamado latent feedback. A cada passo, o estado oculto da camada superior produzido anteriormente é combinado com a representação do token amostrado por meio de uma unidade linear com portas (gated linear unit). Essa combinação é então usada como entrada do passo seguinte. Segundo os autores, o mecanismo permite que computações não verbalizadas retornem à pilha e recebam novamente um orçamento de profundidade, sem alterar a arquitetura padrão do transformador, o cache de chaves e valores (KV cache) ou o objetivo de modelagem de linguagem.

Treinamento e resultados

Para evitar a perda do treinamento paralelo por teacher forcing, os autores introduzem um objetivo de múltiplas passagens programadas. O latent feedback é incorporado apenas na parte final do pré-treinamento, enquanto uma pequena fração das passagens usa realimentação mais profunda, em uma estratégia apresentada como forma de promover estabilidade.

Foram treinados transformadores full-bandwidth de 1 bilhão de parâmetros com até 400 bilhões de tokens. De acordo com os resultados relatados, a realimentação latente melhora a perda de validação, a avaliação de modelos de linguagem em cinco exemplos (5-shot), a geração matemática e de código e o desempenho após ajuste para instruções.

Os autores também afirmam que o mecanismo acrescenta custo insignificante à decodificação por token. Nos experimentos descritos, os modelos com a nova arquitetura igualam ou se aproximam do desempenho de transformadores convencionais treinados com cerca de 1,5 vez mais tokens. Além disso, conseguem produzir rastros de raciocínio mais curtos mantendo precisão igual ou superior, segundo a avaliação apresentada.

Assim, a tese central do trabalho é que reutilizar o estado oculto de alto nível entre passos pode ampliar a capacidade efetiva de processamento vertical de transformadores autorregressivos, aproveitando informação que normalmente seria descartada. A fonte, porém, apresenta apenas os resultados e conclusões dos experimentos realizados pelos autores, sem detalhar neste resumo as condições completas de comparação ou possíveis limitações da abordagem.