split_heads — reshape y luego transpose
Con H cabezas, divide las D = H × D_k características en grupos y luego da a las cabezas su propio eje:
x.reshape(B, S, H, D // H).transpose(0, 2, 1, 3) # (B, S, D) -> (B, H, S, D_k)
| Etapa | Forma | Significado |
|---|---|---|
| Q, K o V proyectada | (2, 6, 12) |
lote, token, características |
| Reshape | (2, 6, 3, 4) |
lote, token, cabeza, características de cabeza |
| Transpose | (2, 3, 6, 4) |
lote, cabeza, token, características de cabeza |
| Q @ Kᵀ | (2, 3, 6, 6) |
lote, cabeza, token consulta, token clave |
La permutación es (0, 2, 1, 3): los ejes 1 y 2 intercambian su lugar, así que la cabeza queda antes que el token. La función rechaza un número de cabezas que no divida a D, y los booleanos.