Differential Transformer: Transformer que cancela ruído de atenção

• Embora o Transformer decoder-only tenha se tornado a arquitetura padrão dos LLMs, em contextos longos ele deixa passar informações essenciais e distribui atenção a tokens irrelevantes; esse ruído de atenção abala o desempenho
• A atenção diferencial divide query e key em dois grupos, cria dois mapas de atenção softmax e, em seguida, calcula o score final de atenção subtraindo o segundo mapa com a aplicação de um λ treinável
• O DIFF Transformer mostrou resultados melhores que o Transformer em experimentos que aumentam o tamanho do modelo, os tokens de treinamento e o comprimento de contexto; para desempenho semelhante, precisa de cerca de 65% do tamanho do modelo ou dos tokens de treinamento
• Apresenta vantagens em modelagem de contexto longo, recuperação de informações essenciais, mitigação de alucinações, in-context learning, raciocínio matemático e redução de activation outliers; em QA e sumarização, reduz o impacto de contextos irrelevantes
• Como mantém o layout geral igual ao do Transformer e substitui apenas a atenção softmax, consegue igualar número de parâmetros e volume de computação, além de reutilizar FlashAttention

---

O problema de ruído de atenção no Transformer

• O Transformer decoder-only é a arquitetura de fato padrão dos LLMs, e em seu núcleo há o mecanismo de atenção, que pondera via softmax a importância dos tokens dentro de uma sequência
• LLMs têm dificuldade para localizar com precisão informações essenciais no contexto, especialmente quando há muitos contextos irrelevantes, pois as pistas para a resposta tendem a ficar soterradas
• Em um exemplo em que é preciso encontrar uma resposta inserida no meio de um monte de documentos, o Transformer tende a atribuir apenas um pequeno attention score à resposta e a distribuir scores excessivos a contextos irrelevantes
• Esses attention scores não desprezíveis distribuídos a contextos irrelevantes atuam como attention noise
• O exemplo de Multi-Needle Retrieval da Figure 1 mostra a acurácia do Transformer e do Differential Transformer como 55% e 85%, respectivamente

Como a atenção diferencial funciona

• O DIFF Transformer é uma arquitetura-base para sequence modeling e LLMs; ele mantém o macro layout do Transformer existente e substitui a atenção softmax comum por atenção diferencial
• A partir da entrada X, projeta query, key e value, mas divide query e key em dois grupos, Q1, Q2, K1, K2, e mantém value como V
• A saída da atenção é calculada como a diferença entre dois mapas de atenção softmax
  • DiffAttn(X) = (softmax(Q1K1^T / √d) − λ softmax(Q2K2^T / √d))V
  • É uma estrutura que remove o ruído comum subtraindo o segundo mapa de atenção do primeiro
• λ é um escalar treinável e é reparametrizado da seguinte forma para ajustar a dinâmica de treinamento
  • λ = exp(λq1 · λk1) − exp(λq2 · λk2) + λinit
  • Nos experimentos básicos, usa-se λinit = 0.8 − 0.6 × exp(−0.3 · (l − 1))
  • Também foi examinada a abordagem de usar o mesmo λinit em todas as layers, por exemplo 0.8, e a ablation mostrou desempenho relativamente robusto a diferenças na estratégia de inicialização
• Esse método é semelhante à ideia de um amplificador diferencial, que remove common-mode noise pela diferença entre dois sinais, e à de fones com cancelamento de ruído
• Naderi et al. provaram que a atenção diferencial torna a distribuição espectral da matriz de atenção mais equilibrada e resolve de forma eficaz o rank collapse

Estrutura multi-head e arquitetura completa

• A atenção diferencial multi-head usa matrizes de projeção diferentes para cada head e, dentro da mesma layer, compartilha o escalar λ entre as heads
• A saída de cada head recebe RMSNorm de forma independente, depois é multiplicada por (1 − λinit); em seguida, as heads são concatenadas na dimensão de canais e passam pela projeção de saída WO
• A notação GroupNorm da Figure 2 enfatiza que a normalização é aplicada independentemente a cada head
  • A atenção diferencial tende a apresentar padrões mais esparsos, portanto as informações estatísticas entre heads são mais diversas
  • A normalização por head normaliza cada head antes da concatenação e melhora as estatísticas de gradiente
• A layer completa do DIFF Transformer é composta por dois módulos
  • MultiHead(LN(Xl)) + Xl
  • SwiGLU(LN(Yl)) + Yl
• Estruturalmente, usa pre-RMSNorm e SwiGLU, seguindo as melhorias da família LLaMA

Eficiência e estabilidade de treinamento

• A atenção diferencial pode reutilizar diretamente o FlashAttention, o que permite melhorar bastante a eficiência do modelo
• O número de heads é definido como h = dmodel / 2d, em que d é igual à dimensão de head do Transformer
• Essa configuração serve para igualar ao Transformer o número de parâmetros e a complexidade computacional
• Após a normalização das heads, usa-se um multiplicador fixo (1 − λinit) para alinhar o fluxo de gradientes com o do Transformer
• O Appendix G mostra que o fluxo geral de gradientes é mantido de forma semelhante ao do Transformer, permitindo herdar hiperparâmetros parecidos e garantir estabilidade de treinamento

Resultados experimentais e efeitos de aplicação

• Os experimentos de modelagem de linguagem escalam o DIFF Transformer aumentando parameter count, training tokens e context length
• Na scaling curve, o DIFF Transformer precisa de cerca de 65% do tamanho do modelo ou dos tokens de treinamento para alcançar desempenho de modelagem de linguagem semelhante ao do Transformer
• Também apresentou desempenho melhor que o Transformer em downstream tasks e, em avaliações de long-sequence, aproveita o contexto de forma eficaz à medida que ele aumenta
• Na recuperação de informações essenciais, mostra um padrão de atribuir attention scores mais altos ao span da resposta e scores mais baixos a contextos irrelevantes
• Em QA e text summarization, fica menos distraído por contextos irrelevantes, mitigando hallucination
• Em in-context learning, além de aumentar a acurácia, é mais robusto a order permutation, considerada um problema crônico de robustez
• O resultado de reduzir activation outliers oferece novas oportunidades para quantization

[GN⁺]

Opiniões do Hacker News

• Sinto que estão deixando passar a intuição central aqui. Entende-se o problema de que a atenção softmax comum tem dificuldade para atribuir atenção próxima de 0 a coisas irrelevantes, e também se entende que, com uma estrutura de subtração, é possível criar pesos de atenção exatamente, ou quase, 0 sem valores de ativação anômalos.
  Porém, essa estrutura também parece capaz de criar facilmente pesos de atenção negativos, o que parece equivalente a dar atenção positiva à negação do vetor de valores. Intuitivamente, parece difícil manter equilibradas todas as coisas sem interesse perto de 0. Ainda assim, a Figura 1 mostra que funciona bem, então não duvido da possibilidade; só não consigo visualizar bem como a rede consegue fazer isso concretamente

  • Há um erro no softmax e na atenção comuns. O softmax deveria ser exp()/1+∑exp()
    O ponto essencial é que se adiciona 1 ao denominador. No limite negativo, o softmax pode virar 0, e não algum épsilon. Colocar um valor 0 adicional em x produz o mesmo efeito. A desvantagem é que, para corrigir isso, seria preciso treinar o modelo de novo do zero
  • Em vez de permitir pesos negativos, não passá-los por algo como ReLU até soa como se atrapalhasse um pouco o modelo. Mas esse tratamento pode ser um problema mais fácil para o modelo do que parece.
    Olhando o primeiro gráfico de pesos de atenção, há de fato pontuações negativas nas regiões de ruído. Ainda assim, a atenção sobre essa região já é muito pequena de qualquer forma. O segundo mapa de atenção só precisa prever o ruído do primeiro e, como tem acesso completo à primeira entrada, é uma tarefa que pode executar com bastante precisão.
    Voltando à analogia do artigo com o mundo real, fones com cancelamento de ruído têm acesso, pelo microfone, ao som ouvido pelo ouvido, então conseguem emitir um sinal de cancelamento preciso. Da mesma forma, o segundo mapa de atenção sabe o que entra no primeiro, então pode produzir o sinal de cancelamento correspondente. Não é perfeito, mas fones com cancelamento de ruído também não são perfeitos e ainda assim chegam a 99%, o suficiente para melhorar o desempenho
  • Intuitivamente, parece muito fácil o modelo otimizar lambda para 0 durante o treinamento. Nesse caso, ele vira basicamente um Transformer comum com um mecanismo de poda de parâmetros excessivamente complexo acoplado.
    A poda já está bem estabelecida na literatura como uma forma surpreendentemente eficaz de reduzir o número de parâmetros, chegando a cerca de 40%. O modelo real provavelmente não funciona exatamente assim, mas não seria surpreendente se, no fim, ele apenas aproximasse um Transformer comum
  • Valores negativos podem aumentar a expressividade

• Muito inteligente. Gosto desse tipo de trabalho detalhado, e a mudança é pequena, então parece fácil para outras pessoas aplicarem. Excelente.
  No entanto, a última frase da introdução da seção "2 Differential Transformer" me preocupa um pouco. Eles dizem que usam melhorias de trabalhos anteriores, mas, pelo contexto gramatical, não fica claro se essas melhorias foram incluídas tanto no Transformer comum quanto no diff Transformer. Se não foram, a comparação fica confusa. A expressão "main difference" na frase imediatamente anterior acendeu um alerta.
  Claro, pesquisadores bem-intencionados podem estar cientes disso e simplesmente não ter sentido necessidade de deixar explícito. Mas, em algumas publicações dessa área, todo cuidado ainda é pouco

  • Sim. Parece realmente bom. Há melhorias generalizadas de perplexity por tempo de treinamento, por token de treinamento e por tamanho de modelo.
    Isso me lembra arquiteturas MoE, em que se escolhe o melhor modelo pequeno para lidar com parte ou com todo o trabalho de inferência. Fico curioso se o MoE também ganhou benefícios semelhantes ao forçar o Transformer a distinguir possibilidades alternativas.
    De qualquer forma, se os números se sustentarem, acho que será amplamente adotado. Como você disse, praticamente não parece ter desvantagens e parece fácil de reproduzir
  • As outras duas mudanças que eles mencionaram já foram amplamente adotadas e estão incluídas também em alguns dos modelos usados como comparação. Parece que eles as listaram por completude, como diferenças em relação à arquitetura Transformer original

• Como quase tudo neste novo mundo de machine learning, é realmente confuso entender por que isso funciona.
  A analogia com fones de ouvido com cancelamento de ruído ajuda, mas, nesse caso, sabemos claramente o que é sinal e o que é ruído. Se soubéssemos isso aqui também, não sei por que precisaríamos fazer cancelamento de ruído em primeiro lugar.

  • Um único softmax não consegue prever exatamente 0, apenas um número muito pequeno. Quando há muitos valores a somar, esses valores pequenos misturam muitas coisas irrelevantes na saída, contaminando-a com o ruído mencionado no artigo.
    Pior ainda: o gradiente dos valores baixos de atenção fica muito pequeno, de modo que são necessárias muitas atualizações de pesos para desfazer esse tipo de erro. Por outro lado, ao subtrair as saídas de dois softmaxes, o modelo consegue prever pesos exatamente iguais a 0 para alguns valores e ainda manter um fluxo de gradiente razoável.
    Ou seja, o modelo já sabe o que é ruído, mas um único softmax dificulta excluí-lo. Além disso, em um softmax único, a saída de todos os heads é forçada a permanecer dentro do fecho convexo dos vetores de valor; já nessa variação, cada head pode escolher seu próprio lambda e deslocar o intervalo de saída para fora do fecho convexo pré-definido pelos valores. Com isso, aumenta a expressividade do modelo como um todo.
  • Fones de ouvido com cancelamento de ruído provavelmente são uma analogia ruim aqui.
    Um exemplo melhor é a sinalização diferencial, usada em áudio profissional e em muitos protocolos de sinais digitais como Ethernet, HDMI e USB. Em vez de usar um único fio referenciado ao terra, o sinal é enviado como a diferença entre dois fios. Os dois fios carregam o mesmo sinal em polaridades opostas e passam lado a lado, de modo que o ruído externo é aplicado igualmente a ambos.
    A tensão pode variar, mas a diferença de tensão entre os dois fios permanece igual. No receptor, ao subtrair as duas tensões, o ruído simplesmente se cancela.
  • Não tente encontrar uma analogia; pense nisso apenas como a adição de uma nova capacidade matemática. Ao permitir atenção negativa, a rede pode dizer, no cálculo de atenção: "quero subtrair a contribuição deste token". Antes, ela só podia reduzir o quanto seria somado.
    Uma forma simples de fazer isso seria remover o softmax ou usar sigmoid, mas, na prática, parece que o softmax funciona melhor.
  • Uma hipótese para explicar por que isso funciona é que ele atenua uma desvantagem do RoPE.
    Em termos simples, RoPE é uma estratégia moderna que dá ao modelo informações sobre a distância entre a query e a key ao calcular atenção. É a melhor estratégia disponível hoje, mas tem uma grande desvantagem: torna algumas conexões entre tokens distantes muito mais fortes do que gostaríamos. Xpos (https://arxiv.org/pdf/2212.10554) também é um artigo da Microsoft que trata do problema do RoPE, e a Figure 1 na página 4 mostra uma interpretação visual da força de atenção em forma de onda senoidal. O ideal seria que ela fosse suave.
    Acredito que um dos grandes motivos pelos quais o Differential Transformer funciona especialmente bem em sequências longas é que, mesmo quando q1 e q2 não combinam com nenhum token, a força relativa do RoPE assume o mesmo valor, fazendo o ruído se cancelar. Ficam apenas as correspondências pretendidas, embora ao custo de enfraquecer em certa medida os valores que o RoPE originalmente trazia.
    Claro, é apenas uma hipótese. Seria fácil verificar isso experimentalmente comparando ambos com uma baseline que use alibi attention (https://arxiv.org/pdf/2108.12409). O alibi tem outros compromissos que esse método não conseguiria mitigar, mas ainda assim seria um resultado realmente interessante.
  • Parte dos trabalhos anteriores relacionados aqui são as ladder networks e, de forma um tanto aproximada, as residual nets. Ambas podem ser interpretadas como treinar o modelo para reduzir o erro de previsões anteriores, em vez de prever diretamente o resultado final.
    A intuição sobre por que isso funciona parece estar em tornar a paisagem do gradiente descendente um pouco mais amigável, facilitando o aprendizado em pequenos passos. Isso porque agora a própria rede é explicitamente projetada em torno da ideia de que, no início, suas previsões cometem muitos erros e melhoram com o tempo.

• Se eu entendi corretamente "Differential attention takes the difference between two softmax attention functions to eliminate attention noise", essa arquitetura parece ser um trade-off: usa o dobro de memória de atenção em troca de um modelo de maior qualidade, ou de qualidade semelhante com menos parâmetros.
  O trecho "6.8B-size DIFF Transformer achieves a validation loss comparable to 11B-size Transformer, requiring only 62.2% of parameters" levanta algumas perguntas. Se ele tem apenas 60% dos parâmetros, será que isso compensa o fato de o espaço de atenção ser duplicado e resulta em características de memória parecidas com as de um Transformer tradicional? E esse trade-off muda de forma perceptível entre treinamento e inferência?

  • Entendi que os parâmetros adicionais necessários para o segundo mecanismo de atenção também estão incluídos nesses 6.8B parâmetros. Ou seja, não é uma contagem hipotética de parâmetros que um Transformer padrão teria, mas o total de parâmetros do modelo. Por isso o resultado é duplamente impressionante.
    O artigo diz: "We set the number of heads h = dmodel/2d, where d is equal to the head dimension of Transformer. So we can align the parameter counts and computational complexity." Em outras palavras, eles compensam isso reduzindo pela metade o número de heads de atenção em cada camada.
  • Parece que eles reduziram pela metade o número total de heads e duplicaram V e O para mitigar a memória e o cálculo adicionais. Não conferi a matemática de fato, mas, excetuando operações baratas como multiplicações por constantes e subtrações, o número de operações de ponto flutuante deve ser equivalente.
  • A economia de RAM provavelmente será compensada, mas pode reduzir o espaço necessário para armazenamento e também o tempo de inicialização, dependendo da velocidade do armazenamento e do tamanho do modelo. Então pode ser uma boa opção para modelos de menor porte em dispositivos de consumo.
  • O tamanho do cache KV dobra e, com contextos grandes, isso pode chegar a vários GB.

• Fico curioso sobre qual deve ter sido a história por trás da fórmula "We empirically find that the setting λᵢₙᵢₜ = 0.8 − 0.6 × exp(−0.3 · (l − 1)) works well in practice"

  • 0,8 funciona bem; mas vamos tentar valores iniciais mais baixos nas camadas inferiores. Algo como 0,2. Ótimo, precisamos de uma fórmula que vá de 0,2 a 0,8 e se aproxime lentamente de 0,8. Parece que ficaram mexendo nos números por 20 minutos e decidiram que isso estava bom
  • Muita coisa mesmo acaba sendo ajustada por otimização do tipo girar um dial analógico ou ouvir até parecer aceitável
  • Essa fórmula parece, pelo menos no início do treinamento, tornar o termo de atenção negativo menor nas camadas da frente, ou seja, para l pequeno, do que nas camadas posteriores. Faz sentido. Antes de concluir quais são de fato as poucas posições a observar, você provavelmente quer dar um pouco de atenção a tudo
    Mas o autor não parece discutir separadamente essa escolha no artigo

• O ponto central que eu não entendia no começo era: o que acontece se os dois grupos de atenção aprenderem a mesma coisa? Como as máscaras de atenção são subtraídas uma da outra, se ambas produzirem valores parecidos, a atenção total cai para 0 e a perda aumenta
  Então a única forma de reduzir a perda é aprender a prestar atenção em coisas diferentes. Uma das estratégias mais simples que eles podem aprender é, como o artigo afirma, um grupo focar no contexto relevante e o outro no contexto irrelevante. Assim, um grupo aprende o ruído e o outro aprende o sinal. Na prática não é tão bem separado, mas é uma simplificação útil para entender

  • A parte interessante não é a simples subtração, mas subtrair apenas uma parte do segundo softmax
    Se as duas cópias forem idênticas, a saída do softmax também será idêntica e a diferença será 0 em todos os pontos, então isso faz sentido. Mas, ao subtrair uma cópia escalada, o processo de normalizar a diferença parece destacar os valores de sinal acima do ruído mais do que antes da normalização
  • Fico pensando se há alguma analogia com nossa própria experiência de surpresa e sua utilidade para explicar o que acontece quando os dois grupos de atenção aprendem a mesma coisa
    Seria como um head de atenção aumentar o peso quando se surpreende com o que outro head aprendeu e diminuir o peso quando ambos encontram a mesma coisa, por não considerar isso muito surpreendente
    Admito que "surpresa" ocupa uma área bem grande na minha base de conhecimento[1][2][3]. Como emoção subjetiva e função adaptativa da mente, é um dos sistemas adaptativos mais complexos que conhecemos
    [1] https://plus.maths.org/content/information-surprise
    [2] https://blakeelias.name/papers/Multi-Agent-Cooperation-Intri...
    [3] https://complexity.simplecast.com/episodes/81/transcript
  • Talvez exista uma pequena possibilidade de ambos aprenderem a mesma coisa, mas não parece provável o bastante para virar um problema importante
  • A função de perda também não poderia aplicar uma penalidade quando os dois aprendem a mesma coisa?

• Fico curioso sobre o que se perde aqui. Certamente deve haver algum trade-off
  Também me pergunto se isso afeta a criatividade ou a capacidade de interpolar entre conceitos. Alucinação e criatividade parecem bastante relacionadas. Entendo alucinação como algo desalinhado com o espaço de interpolação que humanos consideram apropriado

  • Não entendo por que alucinação e criatividade parecem relacionadas. Para mim é simplesmente um erro de amostragem
    Claro, erros podem às vezes servir de inspiração, mas criatividade é muito mais do que erro
    Esses modelos de linguagem são preditores do próximo token. O próximo token é previsto por amostragem no espaço de probabilidades gerado pelo modelo. Esse processo de amostragem pode ser não determinístico
    Alucinação ocorre quando o resultado dessa amostragem gera tokens que formam uma frase falsa ou não intencional. Também poderíamos considerar tudo que o modelo produz como alucinação, mas treinamos o modelo para gerar um espaço no qual é mais provável que ele alucine o que queremos. Caso contrário, ele só produziria ruído sem sentido
    "Alucinação" é uma palavra realmente péssima para aquilo que tenta explicar
  • Um dos trade-offs é velocidade e memória. Como há o dobro de pesos Q e K no bloco de atenção, a vazão nos H100 deles caiu cerca de 10%. Está na Tabela 7 do Apêndice A
  • Nem toda alucinação é criatividade. Imagine uma aplicação RAG: o modelo precisa seguir os documentos fornecidos

• Fico curioso sobre quanto do valor aqui vem de compensar o ruído posicional criado pelo RoPE. Além dos modelos RoPE daqui, eu também gostaria de ver uma tabela comparando uma versão com alibi e uma baseline com alibi
  Ainda assim, é uma melhoria enorme, e parabéns aos pesquisadores

• O que está acontecendo aqui é que o softmax não consegue empurrar valores para 0, mas, ao subtrair dois mapas de softmax, é possível criar uma saída 0?

  • A pergunta seguinte é se a probabilidade de produzir 0 não é extremamente baixa
  • Ou também pode ser negativa

• É um bom problema a resolver, mas acho que a abordagem está errada
  Para saber ao que se prestou atenção e o contexto completo, isso deveria ser feito de forma hierárquica. Se o vetor diferencial é calculado com as mesmas entradas do vetor de atenção, não sei como ele saberia modificar corretamente o vetor de atenção

  • No fim, tudo não é ajustado proporcionalmente à direção indicada pela derivada da retropropagação e ao seu gradiente? Em outras palavras, enquanto o sistema de retropropagação funcionar, imagino que não deva haver problema em saber em que direção ajustar os pesos