Como as cabeças de atenção compartilham sua memória — e a encolhem muitas vezes. · Source

Sua memória do chat dispara — então as cabeças a compartilham.

Cada cabeça de atenção guarda suas próprias anotações de cada palavra até aqui — suas Keys e Values. Empilhe cabeças, camadas e palavras suficientes e essa pilha, relida a cada passo, vira o verdadeiro gargalo. A solução é quase atrevida: a maioria das cabeças não precisa da própria cópia. Como vários canudos em poucos potes compartilhados: cada cabeça ainda sorve do seu jeito, mas de uma memória que têm em comum.

Cada cabeça acumula sua própria cópia do passado.

M = 2\,L\,n\,H\,d_{\text{head}}\,b \;=\; 2\,L\,n\,d_{\text{model}}\,b

Pense no cache como a memória que o modelo tem do chat. Na atenção padrão, cada cabeça guarda suas próprias Keys e Values — para cada camada, cada palavra. Como uma longa mesa de banquete com uma garrafa idêntica em cada lugar: uma forma só, copiada ao longo da fila. Em poucas palavras: a memória cresce como cabeças × camadas × palavras, e você relê tudo só para acrescentar a próxima palavra.

E se cada cabeça compartilhasse uma só memória?

\mathrm{head}_i = \mathrm{softmax}\!\left(\dfrac{Q_i K^\top}{\sqrt{d_{\text{head}}}}\right) V \qquad i = 1,\dots,H

Eis a jogada ousada. Mantenha o Query de cada cabeça — sua própria pergunta — mas deixe todas lerem de um único Key e Value compartilhado. Como uma só máquina de espresso servindo uma fila inteira de baristas: cada um ainda tira a própria xícara, de uma única fonte. Em poucas palavras: a pergunta continua única em cada cabeça; só a memória é compartilhada — então ela não cresce mais com o número de cabeças.

Mas uma só memória compartilhada é grosseira demais.

Compartilhar uma memória entre todas as cabeças vai longe demais. Cada cabeça mantinha anotações moldadas para a própria pergunta; forçadas por um único conjunto, as mais afiadas borram e o treino pode oscilar. Como um coro inteiro espremido em um só microfone: você ainda ouve o canto, mas cada voz distinta se achata no mesmo canal. Economizamos memória — e perdemos, em silêncio, parte do alcance do modelo.

Então não compartilhe uma — compartilhe em grupos pequenos.

\mathrm{head}_i = \mathrm{softmax}\!\left(\dfrac{Q_i\,K_{g(i)}^\top}{\sqrt{d_{\text{head}}}}\right) V_{g(i)},\qquad g(i)=\left\lceil iG/H\right\rceil,\quad 1\le G\le H

A solução divide a diferença. Separe as cabeças em poucos grupos; cada cabeça de um grupo compartilha um Key e um Value, mas os grupos seguem separados. Como vans de carona ao amanhecer: nem um carro para cada um, nem um único ônibus para todos — um punhado de vans, cada uma com sua pequena turma. Um botão ajusta tudo — chame de G memórias compartilhadas para H cabeças: G = H é a atenção comum, G = 1 é o extremo grosseiro, e alguns grupos no meio é o ponto ideal.

Oito vezes menos para carregar, a cada palavra.

M = 2\,L\,n\,n_{kv}\,d_{\text{head}}\,b,\qquad \text{shrink}=\dfrac{H}{n_{kv}}\quad (n_{kv}=H,\,G,\,1)

Agora o ganho. Com G memórias compartilhadas em vez de H, o cache encolhe pelo fator H ∕ G — muitas vezes oito vezes. E rende em dobro: o contexto longo cabe, e a cada palavra gerada o modelo recarrega essa memória inteira de um armazenamento lento. Como uma van de entrega levando uma só caixa compartilhada em vez de oito idênticas: cada viagem fica mais leve. Memória menor, tráfego mais leve, a mesma resposta.

E você pode adaptar um modelo que já treinou.

K^{(g)} = \dfrac{G}{H}\sum_{i\,\in\,g} K_i \qquad \big(\,V^{(g)}\text{ likewise},\ |g|=H/G\,\big)

O melhor: não precisa treinar do zero. Pegue um modelo pronto em que cada cabeça tem sua própria memória e tire a média das memórias de cada grupo até sobrar uma — depois ajuste de leve para assentar. Como vários riachos se juntando em uma única poça: muitos fios quase idênticos viram um. Funciona porque, dentro de um grupo, essas memórias já eram quase iguais desde o início — e é justamente por isso que compartilhar quase não dói.

Cem cabeças. Um punhado de memórias.

A atenção agrupada funciona porque dezenas de cabeças, no fim, guardavam quase as mesmas anotações — diferentes o bastante para fazer suas próprias perguntas, parecidas o bastante para compartilhar as respostas. Então, quão separadas elas eram, de verdade? Talvez a identidade de uma cabeça nunca estivesse no que ela lembra, só no que ela pergunta. 🌱 Onde vive, então, um ponto de vista?