Comment les têtes d'attention partagent leur mémoire — et la réduisent plusieurs fois. · Source

Sa mémoire du fil enfle — alors les têtes la partagent.

Chaque tête d'attention tient ses propres notes sur chaque mot jusqu'ici — ses Keys et ses Values. Empilez assez de têtes, de couches et de mots et cette pile, relue à chaque étape, devient le vrai goulot d'étranglement. La solution est presque effrontée : la plupart des têtes n'ont pas besoin de leur propre copie. Comme plusieurs pailles dans quelques bocaux partagés : chaque tête sirote encore à sa façon, mais dans une mémoire qu'elles ont en commun.

Chaque tête thésaurise sa propre copie du passé.

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

Voyez le cache comme la mémoire qu'a le modèle du fil. Dans l'attention standard, chaque tête garde ses propres Keys et Values — pour chaque couche, chaque mot. Comme une longue table de banquet avec une bouteille identique à chaque place : une seule forme, recopiée tout du long. En clair : la mémoire croît comme têtes × couches × mots, et on la relit en entier rien que pour ajouter le mot suivant.

Et si chaque tête partageait une seule mémoire ?

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

Voici le coup audacieux. Gardez le Query de chaque tête — sa propre question — mais laissez toutes les têtes lire dans un seul Key et Value partagé. Comme une seule machine à espresso servant toute une rangée de baristas : chacun tire encore sa propre tasse, d'une source unique. En clair : la question reste propre à chaque tête ; seule la mémoire est partagée — elle ne croît donc plus avec le nombre de têtes.

Mais une seule mémoire partagée est trop grossière.

Partager une mémoire entre toutes les têtes va trop loin. Chaque tête tenait des notes taillées pour sa propre question ; forcées dans un seul jeu, les plus fines se brouillent et l'entraînement peut vaciller. Comme une chorale entière serrée autour d'un seul microphone : on entend encore le chant, mais chaque voix distincte s'aplatit dans le même canal. On a économisé de la mémoire — et discrètement perdu une part de l'étendue du modèle.

Alors n'en partagez pas une — partagez par petits groupes.

\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

La solution coupe la poire en deux. Rangez les têtes en quelques groupes ; chaque tête d'un groupe partage un Key et un Value, mais les groupes restent distincts. Comme des navettes en covoiturage à l'aube : ni une voiture chacun, ni un seul bus pour tous — une poignée de navettes, chacune avec sa petite foule. Un seul réglage le fixe — appelez-le G mémoires partagées pour H têtes : G = H est l'attention ordinaire, G = 1 l'extrême grossier, et quelques groupes au milieu, le juste équilibre.

Huit fois moins à porter, à chaque mot.

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

Maintenant le gain. Avec G mémoires partagées au lieu de H, le cache rétrécit du facteur H ∕ G — souvent huit fois. Et ça paie deux fois : le contexte long tient, et à chaque mot généré le modèle recharge toute cette mémoire depuis un stockage lent. Comme une camionnette de livraison portant une seule caisse partagée au lieu de huit identiques : chaque trajet est plus léger. Mémoire réduite, trafic allégé, la même réponse.

Et vous pouvez convertir un modèle déjà entraîné.

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

Le meilleur : pas besoin de repartir de zéro. Prenez un modèle fini où chaque tête a sa propre mémoire, et moyennez les mémoires de chaque groupe en une seule — puis affinez brièvement pour la stabiliser. Comme plusieurs ruisselets se rejoignant en un seul bassin : beaucoup de filets quasi identiques n'en font qu'un. Ça marche parce que, au sein d'un groupe, ces mémoires étaient presque les mêmes depuis le début — et c'est toute la raison pour laquelle le partage ne coûte presque rien.

Cent têtes. Une poignée de mémoires.

L'attention groupée marche parce que des dizaines de têtes gardaient en fait presque les mêmes notes — assez différentes pour poser leurs propres questions, assez semblables pour partager les réponses. Alors, à quel point étaient-elles distinctes, vraiment ? Peut-être que l'identité d'une tête n'a jamais tenu à ce qu'elle retient, mais seulement à ce qu'elle demande. 🌱 Où vit donc, au juste, un point de vue ?