Resumo Técnico: Resíduos de Atenção com Desacoplamento de Papéis

Declaração do Problema

Nas arquiteturas Transformer padrão, o fluxo residual atua como uma memória fixa onde cada camada herda o estado imediatamente anterior via adição serial. Arquiteturas recentes de roteamento de profundidade, como o Block Attention Residuals (Block AttnRes), melhoram isso ao permitir que as subcamadas recuperem uma mistura dependente de conteúdo de representações anteriores, em vez de apenas o predecessor imediato. Isso transforma a profundidade de uma ordem computacional fixa em um eixo de recuperação aprendido.

No entanto, o Block AttnRes introduz uma restrição de design específica: ele utiliza uma única distribuição de profundidade para construir a entrada compartilhada para todas as três projeções de atenção (Queries, Keys e Values). Isso acopla dois papéis funcionalmente distintos:

1. Geometria de Correspondência (Matching Geometry): Determinada por Queries e Keys, que decidem onde olhar (quais posições de tokens coincidem).
2. Recuperação de Conteúdo (Content Retrieval): Determinada por Values, que decidem o que conteúdo é recuperado.

O artigo postula que forçar a representação que é melhor para decidir "onde olhar" a ser também a melhor para decidir "o que ler" é um acoplamento desnecessário. A questão central é se o desacoplamento desses papéis melhora o pré-treinamento de modelos de linguagem quando todos os outros fatores são mantidos constantes.

Metodologia: Resíduos de Atenção com Desacoplamento de Papéis (RD-AttnRes)

O autor introduz o Role-Decoupled Attention Residuals (RD-AttnRes), uma relaxação arquitetural mínima do Block AttnRes.

• Arquitetura: O RD-AttnRes mantém o conjunto de origem da arquitetura pai, a topologia de blocos, a normalização, as rotas MLP e os mecanismos de escrita (writeback). A única modificação ocorre na leitura de profundidade do lado da atenção.
• Mecanismo de Desacoplamento:
  • Queries e Keys (Q/K): Compartilham uma única consulta de rota aprendida ($q^{QK}_\ell$) para computar uma mistura ($h^{QK}_\ell$) a partir da memória de profundidade residual. Isso preserva uma geometria de correspondência coerente.
  • Values (V): Aprendem uma consulta de rota independente ($q^V_\ell$) para computar uma mistura separada ($h^V_\ell$) sobre o exato mesmo conjunto de fontes residuais.
• Eficiência de Parâmetros: A modificação adiciona apenas um vetor de dimensão $d$ (largura do modelo) por camada. Para as escalas testadas (120M e 343M de parâmetros), isso representa um aumento insignificante de parâmetros (<0,01%).
• Recuperação Exata: Se a consulta de rota de valor for vinculada à consulta de Q/K ($q^V_\ell = q^{QK}_\ell$), o RD-AttnRes reduz-se exatamente à arquitetura Block AttnRes original. Isso garante que a proposta seja uma extensão estrita e auditável, em vez de um sistema com múltiplas mudanças simultâneas.
• Custo Computacional: O método adiciona $O(BTK_\ell d)$ de trabalho por camada de atenção (onde $B$ é o tamanho do lote, $T$ é o comprimento da sequência e $K_\ell$ é o número de fontes), mas não introduz novos termos de atenção token-para-token de $O(T^2d)$.

Desenho Experimental

O estudo emprega um protocolo experimental pareado rigoroso para isolar o efeito do desacoplamento de papéis:

• Datasets: FineWeb-Edu (amostra de 10BT), com uma divisão determinística para treinamento, validação e teste.
• Modelos: Dois modelos Transformer Pre-Norm de escala decoder-only:
  • S: 120M de parâmetros (16 camadas, largura 576).
  • M: 343M de parâmetros (24 camadas, largura 832).
• Protocolo: Cinco sementes (seeds) correspondentes por escala. Tanto o Block AttnRes quanto o RD-AttnRes foram treinados com inicialização, ordem de dados, cronogramas de otimização e orçamento de 2.0B de tokens idênticos.
• Controles Direcionais aos 0.5B de tokens:
  • Tied (Vinculado): Força a rota de valor a corresponder à rota Q/K (recuperando o modelo pai).
  • Average (Média): Executa duas rotas, mas calcula a média de suas misturas (testando se a execução dupla por si só ajuda).
  • Frozen-V (V Congelado): Mantém a rota de valor uniforme (testando se uma rota fixa ajuda).
• Diagnósticos: A divergência de Jensen–Shannon (JS) é medida entre as distribuições de profundidade de Q/K e V para verificar se as rotas aprendem estratégias distintas.

Principais Resultados

1. Melhoria na Perda de Validação: O RD-AttnRes melhora a perda de log-verossimilhança negativa (NLL) de validação em todas as 10 comparações pareadas (5 sementes × 2 escalas).
   • Escala 120M: Redução média de NLL de 0,0301, correspondendo a uma redução de 2,97% na perplexidade.
   • Escala 343M: Redução média de NLL de 0,0247, correspondendo a uma redução de 2,43% na perplexidade.
2. Persistência do Ganho: A melhoria é visível em todos os marcos de avaliação (0.25B, 0.5B, 1.0B e 2.0B de tokens), embora a magnitude absoluta da diferença diminua ligeiramente conforme o treinamento avança.
3. Validação dos Controles:
   • Os controles Tied e Average performam de forma semelhante ao Block AttnRes original, indicando que o ganho não se deve à contagem de parâmetros ou à duplicação de computação.
   • O controle Frozen-V performa mal, confirmando que a rota de valor deve ser aprendida, não fixa.
4. Divergência de Rota: As distribuições aprendidas de Q/K e V não colapsam para a mesma solução. A divergência JS média entre as camadas é de aproximadamente 0,218 (120M) e 0,215 (343M), confirmando que o modelo aprende leituras de profundidade distintas para correspondência versus recuperação de conteúdo.

Significância e Alegações

O artigo faz uma alegação estreita e delimitada em vez de uma universal. O autor não afirma que o RD-AttnRes é superior a todas as arquiteturas de roteamento residual ou que estabelece um comportamento de escala universal.

Em vez disso, os resultados apoiam a conclusão específica de que dentro do regime de treinamento testado (FineWeb-Edu, 2B de tokens, escalas de modelo específicas), a leitura de profundidade QKV compartilhada no Block AttnRes é uma restrição consequencial. As evidências sugerem que a correspondência (Q/K) e a recuperação de conteúdo (V) beneficiam-se de leituras de profundidade diferentes.

Limitações Reconhecidas:

• O estudo é limitado a um corpus web em inglês e modelos com menos de um bilhão de parâmetros.
• A lacuna de desempenho absoluto diminui em horizontes de treinamento mais longos; trabalhos futuros são necessários para determinar se a lacuna persiste ou cresce em 4B–10B de tokens.
• Há uma ressalva de sistemas: a implementação atual incorre em uma penalidade de throughput (~9–10%) e maior uso de memória em comparação ao Block AttnRes de referência, embora kernels fundidos possam mitigar isso.
• Comparações com outras arquiteturas recentes (ex: MUDDFormer, Delta AttnRes) são conceituais, não baselines empíricas pareadas.

Conclusão:
O artigo estabelece que, quando a profundidade se torna um eixo de recuperação aprendido, o compartilhamento deve seguir a função subsequente (papel) em vez da conveniência de implementação. Ao desacoplar a rota de valor com um overhead mínimo de parâmetros, o RD-AttnRes alcança melhorias consistentes na perda de pré-treinamento, validando a hipótesia de que a geometria de correspondência e a recuperação de conteúdo são tarefas distintas que exigem representações de profundidade distintas.