Cómo los cabezales de atención comparten su memoria — y la reducen muchas veces. · Source

Su memoria del chat se dispara — así que los cabezales la comparten.

Cada cabezal de atención lleva sus propias notas de cada palabra hasta ahora — sus Keys y Values. Apila suficientes cabezales, capas y palabras y ese montón, releído en cada paso, se vuelve el verdadero cuello de botella. El arreglo es casi descarado: la mayoría de los cabezales no necesitan su propia copia. Como muchas pajitas en unos pocos tarros compartidos: cada cabezal sigue bebiendo a su manera, pero de una memoria que tienen en común.

Cada cabezal acapara su propia copia del pasado.

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

Piensa en la caché como la memoria del modelo sobre el chat. En la atención estándar, cada cabezal guarda sus propios Keys y Values — para cada capa, cada palabra. Como una larga mesa de banquete con una botella idéntica en cada asiento: una forma, copiada a lo largo de la fila. En pocas palabras: la memoria crece como cabezales × capas × palabras, y la relees entera solo para añadir la siguiente palabra.

¿Y si cada cabezal compartiera una sola memoria?

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

Aquí va la jugada audaz. Conserva el Query de cada cabezal — su propia pregunta — pero deja que todos lean de un único Key y Value compartido. Como una sola máquina de espresso sirviendo a toda una fila de baristas: cada uno saca su propia taza, de una sola fuente. En pocas palabras: la pregunta sigue siendo única en cada cabezal; solo la memoria se comparte — así ya no crece con el número de cabezales.

Pero una sola memoria compartida es demasiado tosca.

Compartir una memoria entre todos los cabezales va demasiado lejos. Cada cabezal solía llevar notas hechas a la medida de su propia pregunta; forzados a un único conjunto, los más finos se difuminan y el entrenamiento puede tambalearse. Como un coro entero agolpado en un solo micrófono: aún oyes el canto, pero cada voz distinta se aplana en el mismo canal. Ahorramos memoria — y perdimos en silencio parte del rango del modelo.

Así que no compartas una — comparte en grupos pequeños.

\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

El arreglo parte la diferencia. Ordena los cabezales en unos pocos grupos; cada cabezal de un grupo comparte un Key y un Value, pero los grupos siguen separados. Como furgonetas compartidas al amanecer: ni un coche por cabeza, ni un solo autobús para todos — un puñado de furgonetas, cada una con su pequeña multitud. Un dial lo fija — llámalo G memorias compartidas para H cabezales: G = H es la atención corriente, G = 1 es el extremo tosco, y unos pocos grupos en medio es el punto justo.

Ocho veces menos que cargar, en cada palabra.

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

Ahora la ganancia. Con G memorias compartidas en vez de H, la caché se encoge por el factor H ∕ G — a menudo ocho veces. Y rinde doble: el contexto largo cabe, y en cada palabra generada el modelo recarga esa memoria entera desde un almacenamiento lento. Como una furgoneta de reparto que lleva una sola caja compartida en vez de ocho idénticas: cada viaje es más ligero. Menos memoria, menos tráfico, la misma respuesta.

Y puedes adaptar un modelo que ya entrenaste.

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

Lo mejor: no hace falta entrenar desde cero. Toma un modelo terminado donde cada cabezal tiene su propia memoria, y promedia las memorias de cada grupo hasta dejar una — luego afina un poco para asentarla. Como varios riachuelos que confluyen en una sola poza: muchos hilos casi idénticos se vuelven uno. Funciona porque, dentro de un grupo, esas memorias eran casi iguales desde el principio — y por eso compartir apenas duele.

Cien cabezales. Un puñado de memorias.

La atención agrupada funciona porque resultó que docenas de cabezales guardaban casi las mismas notas — bastante distintas para hacer sus propias preguntas, bastante iguales para compartir las respuestas. Entonces, ¿qué tan separados estaban en realidad? Quizá la identidad de un cabezal nunca estuvo en lo que recuerda, sino solo en lo que pregunta. 🌱 ¿Dónde vive, entonces, un punto de vista?