Resumo Técnico: Revisitando a Convenção de Forma de Modelos de Linguagem Transformer

Declaração do Problema

Modelos de linguagem Transformer densos têm historicamente aderido a uma convenção arquitetônica rígida: a Rede de Alimentação Direta (FFN - Feed-Forward Network) utiliza uma forma de Perceptron Multicamadas (MLP) "estreita–larga–estreita". Nesse design, a dimensão do modelo ($d_{model}$) é expandida para uma dimensão intermediária mais larga ($d_h$) com uma razão de expansão tipicamente entre 2 e 4, antes de ser projetada de volta. Essa estrutura domina a alocação de parâmetros, frequentemente representando a maior parte dos parâmetros e das operações de ponto flutuante do modelo.

Os autores questionam se essa forma específica é essencial ou meramente uma instanciação conveniente. Trabalhos teóricos recentes sugerem que MLPs residuais "larga–estreita–larga" (ampulheta/hourglass) podem servir como aproximadores de função ótimos, potencialmente oferecendo expressividade superior com menos parâmetros. Além disso, evidências empíricas em modelos de visão generativos e Transformers de Mistura de Especialistas (MoE) mostraram que estruturas de ampulheta podem superar os designs padrão. O problema central abordado é se a convencional FFN estreita–larga–estreita é uma restrição necessária ou se uma FFN de ampulheta mais profunda e estreita pode oferecer uma alocação mais eficiente de capacidade entre os módulos de atenção e FFN.

Metodologia

O artigo propõe uma variante de Transformer que substitui a FFN convencional por uma FFN de Ampulheta (Hourglass FFN). A metodologia envolve:

1. Redesenho Arquitetônico:

   • A FFN de Ampulheta substitui a camada única de expansão larga por uma pilha de $K$ sub-MLPs em formato de ampulheta conectados por caminhos residuais.
   • Cada sub-MLP ($MLP_i$) segue uma estrutura de compressão-expansão: ela projeta a entrada da $d_{model}$ para uma dimensão de gargalo $d_h$ (onde $d_h < d_{model}$), aplica uma ativação não linear (SwiGLU) e projeta de volta para $d_{model}$.
   • A saída da camada é o resultado da passagem da entrada por $K$ desses sub-blocos residuais.

2. Estratégia de Redistribuição de Parâmetros:

   • Ao reduzir a dimensão intermediária $d_h$ (tornando a FFN "mais estreita"), os autores economizam parâmetros que são tipicamente alocados para a largura da FFN.
   • Esses parâmetros economizados são realocados para:
     • Aumentar a dimensão do modelo ($d_{model}$) e a capacidade de atenção.
     • Aumentar a profundidade da pilha FFN ($K$).
     • Aumentar o número de camadas Transformer ($L$).

3. Configuração Experimental:

   • Baselines: Comparações são feitas contra Transformers do estilo LLaMA e a arquitetura de última geração OLMo-2.
   • Escalas: Os experimentos são conduzidos em tamanhos de modelo variando de 113M a 1B de parâmetros.
   • Restrições: Os experimentos são realizados sob orçamentos de parâmetros totais correspondentes (dentro de uma diferença de 0,001%) e tamanhos de atenção fixos para isolar os efeitos da forma da FFN.
   • Espaço de Busca: Os autores realizam uma busca arquitetônica variando $K$ (profundidade da pilha de ampulheta), $L$ (número de camadas Transformer), $d_{model}$ e a razão $d_h/d_{model}$.

Principais Contribuições

O artigo apresenta quatro contribuições primárias:

1. Desafiando o Padrão: Ele desafia a forma de fato "estreita–larga–estreita" da FFN, propondo uma FFN de Ampulheta "larga–estreita–larga" eficiente que mantém a expressividade através de sub-MLPs empilhados e conectados residualmente.
2. Realocação de Parâmetros: Demonstra uma estratégia onde redirecionar parâmetros da largura da FFN para a largura da atenção ou para a profundidade interna da FFN ($K > 1$) produz desempenho superior, particularmente em escalas menores (113M).
3. Insights Arquitetônicos: Identifica um trade-off de largura-profundidade em forma de U, encontrando razões ótimas para $d_{model}/L$ e $d_h/d_{model}$, e estabelece que estruturas de FFN de ampulheta mais profundas são robustas.
4. Validação de Escalabilidade: Valida a FFN de Ampulheta de 113M a 1B de parâmetros, mostrando que é uma alternativa competitiva aos baselines padrão, com configurações específicas superando modelos convencionais em escalas menores e igualando-os em escalas maiores.

Resultados

A validação empírica produz as seguintes descobertas:

• Desempenho com Atenção Fixa: Na escala de 113M com parâmetros de atenção fixos, as FFNs de Ampulheta alcançam menor perda de validação e perplexidade (PPL) do que as FFNs convencionais. Por exemplo, em $K=8$, a PPL caiu de 36,441 (convencional) para 36,179.
• Configurações Ótimas:
  • Profundidade vs. Largura: Uma razão de $d_h/d_{model} \approx 0,4$ com $K=4$ camadas rendeu a menor perda de validação.
  • Mudança de Parâmetros: Em 113M, deslocar parâmetros para tornar o módulo de atenção dominante (ex: 66M de atenção vs. 47M de FFN) resultou no melhor desempenho (PPL 35,101), sugerindo que o processamento contextual é mais valioso do que a recuperação factual estática para modelos pequenos.
  • Trade-off Largura-Profundidade: Uma curva distinta em forma de U foi observada para a razão $d_{model}/L$, com um ponto ideal (sweet spot) entre 100 e 250.
• Escalabilidade:
  • Escalas Pequena a Média (113M – 906M): Variantes de Ampulheta superaram consistentemente os baselines convencionais em perda de validação, perplexidade e tarefas de raciocínio downstream (ex: Arc Easy, HellaSwag, PIQA).
  • Escala Grande (1B): Na escala de 1B de parâmetros, a variante de Ampulheta alcançou desempenho comparável ao forte baseline convencional (Val PPL 20,082 vs. 20,002) e superou a configuração OLMo-2 (20,46 tinggi).
• Estudos de Ablação: Aumentar a profundidade da ampulheta $K$ melhorou consistentemente o desempenho, embora com retornos decrescentes além de $K=6$. Reduzir $d_h/d_{model}$ abaixo de 0,4 levou a uma degradação significativa no desempenho, indicando um limite inferior para a compressão da FFN.

Significância e Alegações

O artigo alega que a forma de MLP "estreita–larga–estreita" de longa data nos Transformers não é essencial e pode ser redundante. Ao adotar uma FFN de Ampulheta, os autores demonstram que:

1. Eficiência: Modelos podem alcançar desempenho competitivo ou superior com uma distribuição de parâmetros mais eficiente.
2. Realocação de Recursos: A estrutura de Ampulheta permite uma mudança fundamental onde o mecanismo de atenção pode se tornar o consumidor dominante de parâmetros, desafiando o orçamento tradicional dominado pela FFN.
3. Reavaliação de Design: As descobertas provocam uma reavaliação do equilíbrio entre atenção e FFN, sugerindo que, para modelos de linguagem modernos eficientes e expressivos, FFNs mais profundas e estreitas combinadas com mecanismos de atenção mais largos podem ser uma estratégia de design superior.

Os autores permanecem modestos quanto à escalabilidade para modelos de dezenas ou centenas de bilhões de parâmetros, observando que suas descobertas foram extrapoladas da escala de 113M e que a "maldição da profundidade" (problemas de propagação de sinal) requer investigação adicional em arquiteturas mais profundas. Eles também observam que a interação entre FFNs de Ampulheta e mecanismos de atenção avançados (ex: Group-Query Attention) permanece inexplorada.