Tres Matrices Por Cabeza, O Una Gran Matrix

Vista Por Cabeza

Cada cabeza necesita su propia proyección de consulta, proyección de clave y proyección de valor. Para la cabeza h:

Q_h = X · W_Q^h    donde W_Q^h tiene forma [d_model, head_dim]
K_h = X · W_K^h    donde W_K^h tiene forma [d_model, head_dim]
V_h = X · W_V^h    donde W_V^h tiene forma [d_model, head_dim]

X lleva la forma de entrada [batch, seq_len, d_model]. Después de la proyección, Q_h, K_h, V_h cada uno lleva la forma [batch, seq_len, head_dim].

Vista Fusionada

Las matrices por cabeza se sitúan una al lado de la otra en memoria. Una sola matriz fusionada W_Q de forma [d_model, d_model] produce todas las cabezas a la vez:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

La multiplicación de matrices fusionada realiza una sola llamada BLAS en lugar de 12. Los tensor cores de CUDA alcanzan el rendimiento máximo en multiplicaciones de matrices de este tamaño; las multiplicaciones por cabeza subutilizarían el hardware.

Conteo de Parámetros

Tres matrices fusionadas W_Q, W_K, W_V, cada una d_model × d_model. Más la proyección de salida W_O, también d_model × d_model. Para ANDREA-120M:

paráms por atención de la capa = 4 × 768² = 2,359,296 ≈ 2.36M
paráms en 12 capas      = 12 × 2.36M ≈ 28.3M

Aproximadamente un cuarto de los parámetros totales de ANDREA-120M se encuentran en las proyecciones de atención. Los tres cuartos restantes se encuentran en la subcapa MLP y en los embeddings.

Nombrando las Proyecciones

Doce Vectores se Convierten en Uno

Después de que Cada Cabeza Compute

Cada cabeza produce un tensor de salida con forma [batch, seq_len, head_dim]. Doce cabezas producen doce tensores de este tipo. La concatenación a lo largo de la dimensión de características los apila de nuevo juntos:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # para ANDREA-120M

Concat invierte la división. La dimensión total de características regresa a d_model. No hay pérdida de información en las dimensiones; la diferencia radica en lo que contiene cada chunk: el chunk de la cabeza 1 refleja el patrón de atención aprendido de la cabeza 1.

La Proyección de Salida W_O

La concatenación por sí sola deja las cabezas aisladas: la salida de la cabeza 4 está junto a la salida de la cabeza 7, sin que ninguna sea consciente de la otra. La proyección de salida W_O de forma [d_model, d_model] las mezcla:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

Después de W_O, cada dimensión de salida lleva una combinación lineal aprendida de todas las doce cabezas. La información fluye libremente entre cabezas a través de esta única multiplicación de matrices.

Por qué las cabezas se especializan

Nada en la arquitectura obliga a que la cabeza 4 aprenda el tiempo verbal o la cabeza 9 aprenda puntuación pareada. La especialización surge de la presión del gradiente: durante el entrenamiento, las cabezas que contribuyen de manera redundante reciben gradientes más pequeños que las cabezas que contribuyen de manera única. A lo largo de miles de pasos, cada cabeza se asienta en un nicho que reduce la pérdida total de la manera más efectiva.

Empíricamente, los transformadores entrenados muestran cabezas que manejan: patrones posicionales (la cabeza mira el token anterior), patrones sintácticos (la cabeza mira el corchete de cierre coincidente), patrones semánticos (la cabeza mira la entidad nombrada más reciente). Ninguna etiqueta entrena esta especialización. Solo la señal de entrenamiento, propagada a través de W_O, ordena las cabezas.

Por qué doce cabezas, no una cabeza más ancha

Cómo CUDA Almacena Las Cabezas

Un Solo Tensor, Reorganizado

El motor de entrenamiento de ANDREA microgpt_cuda.cu no asigna doce búferes separados para doce cabezas. Asigna un tensor fusionado & trata la dimensión de la cabeza como un patrón de stride:

// después de Q = X · W_Q (una multiplicación de matrices, fusionada entre cabezas)
// Q tiene forma [batch, seq_len, d_model]
// cambiar forma a    [batch, seq_len, n_head, head_dim]
// transponer a  [batch, n_head, seq_len, head_dim]
// ahora cada cabeza es contigua en las dos dimensiones internas

El traspuesto mueve n_head delante de seq_len. ¿Por qué? Porque la siguiente operación (Q_h · K_h^T) necesita que la porción seq_len × head_dim de cada cabeza sea contigua en memoria. Las multiplicaciones de matrices CUDA se ejecutan más rápido en tensores contiguos.

Un Kernel, Muchas Cabezas

Un solo kernel CUDA de atención se ejecuta en todas las cabezas en paralelo. Cada bloque de hilos maneja un par (batch, head); los hilos dentro del bloque cooperan en la tesela seq_len × head_dim. El kernel nunca sabe que procesa múltiples cabezas; la cuadrícula de lanzamiento maneja el paralelismo.

La Configuración Refleja el Hardware

La elección de ANDREA-120M de n_head=12, head_dim=64 se alinea con los tensor cores de RTX 4090, que prefieren teselas de matmul en múltiplos de 16. head_dim=64 = 4 × 16 coincide exactamente con la forma de la tesela. head_dim=32 (ANDREA-12M) también coincide pero subutiliza la tesela. head_dim=72 no coincidiría y forzaría kernels de respaldo.

Imagen Final

Cinco pasos. Tres multiplicaciones de matrices tocan la entrada (proyecciones Q, K, V). Una multiplicación de matrices toca las cabezas concatenadas (W_O). Un kernel de atención maneja todas las cabezas en paralelo. ANDREA-120M ejecuta los cinco pasos una vez por capa, doce capas de profundidad, en cada pasada hacia adelante.

Leyendo la Pipeline