Problema
La atención multicabeza permite que un Transformer atienda a varias cosas a la vez. Con la entrada X (n tokens × d columnas) y matrices de pesos d × d:
- Q = X × w_q, K = X × w_k, V = X × w_v
- Divide las d columnas en
num_headscabezas de d_k = d / num_heads columnas consecutivas (la cabeza h usa las columnas h×d_k a (h+1)×d_k - 1). - Para cada cabeza, ejecuta la atención por producto escalar escalado sobre sus columnas: para el token i, score_j = (Q_i · K_j) / √d_k sobre las columnas de esa cabeza, softmax sobre los n tokens j, y salida = Σ weight_j × V_j (columnas de esa cabeza).
- Coloca las salidas de las cabezas de nuevo en sus columnas (concatenación) y multiplica por w_o.
Implementa multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads).
Entrada
x: Matriz de entrada n × d.
w_q, w_k, w_v, w_o: Matrices de pesos d × d.
num_heads: El número de cabezas (d es divisible entre él).
Salida
Devuelve la matriz de salida n × d.
Restricciones
1 ≤ n ≤ 30, 1 ≤ d ≤ 32, 1 ≤ num_heads ≤ d y d es divisible entre num_heads
-5 ≤ cada valor ≤ 5
Ejemplos
Entrada:
x = [
[1.0, 0.0],
[0.0, 1.0]
]
w_k = [
[1.0, 0.0],
[0.0, 1.0]
]
w_o = [
[1.0, 0.0],
[0.0, 1.0]
]
w_q = [
[0.0, 0.0],
[0.0, 0.0]
]
w_v = [
[1.0, 0.0],
[0.0, 1.0]
]
num_heads = 1
Salida esperada:
expected = [
[0.5, 0.5],
[0.5, 0.5]
]Entrada:
std::vector<std::vector<double>> x = {
{1.0, 0.0},
{0.0, 1.0}
};
std::vector<std::vector<double>> w_k = {
{1.0, 0.0},
{0.0, 1.0}
};
std::vector<std::vector<double>> w_o = {
{1.0, 0.0},
{0.0, 1.0}
};
std::vector<std::vector<double>> w_q = {
{0.0, 0.0},
{0.0, 0.0}
};
std::vector<std::vector<double>> w_v = {
{1.0, 0.0},
{0.0, 1.0}
};
long long num_heads = 1;
Salida esperada:
std::vector<std::vector<double>> expected = {
{0.5, 0.5},
{0.5, 0.5}
};Entrada:
double[][] x = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
double[][] w_k = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
double[][] w_o = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
double[][] w_q = new double[][] {
new double[] {0.0, 0.0},
new double[] {0.0, 0.0}
};
double[][] w_v = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
long num_heads = 1;
Salida esperada:
double[][] expected = new double[][] {
new double[] {0.5, 0.5},
new double[] {0.5, 0.5}
};Explicación
w_q es todo ceros, así que todas las puntuaciones son 0 y cada token promedia las filas de valores: V = X, así que ambas filas pasan a ser [0.5, 0.5].
Entrada:
x = [
[2.0, 4.0]
]
w_k = [
[1.0, 0.0],
[0.0, 1.0]
]
w_o = [
[0.0, 1.0],
[1.0, 0.0]
]
w_q = [
[1.0, 0.0],
[0.0, 1.0]
]
w_v = [
[1.0, 0.0],
[0.0, 1.0]
]
num_heads = 2
Salida esperada:
expected = [
[4.0, 2.0]
]Entrada:
std::vector<std::vector<double>> x = {
{2.0, 4.0}
};
std::vector<std::vector<double>> w_k = {
{1.0, 0.0},
{0.0, 1.0}
};
std::vector<std::vector<double>> w_o = {
{0.0, 1.0},
{1.0, 0.0}
};
std::vector<std::vector<double>> w_q = {
{1.0, 0.0},
{0.0, 1.0}
};
std::vector<std::vector<double>> w_v = {
{1.0, 0.0},
{0.0, 1.0}
};
long long num_heads = 2;
Salida esperada:
std::vector<std::vector<double>> expected = {
{4.0, 2.0}
};Entrada:
double[][] x = new double[][] {
new double[] {2.0, 4.0}
};
double[][] w_k = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
double[][] w_o = new double[][] {
new double[] {0.0, 1.0},
new double[] {1.0, 0.0}
};
double[][] w_q = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
double[][] w_v = new double[][] {
new double[] {1.0, 0.0},
new double[] {0.0, 1.0}
};
long num_heads = 2;
Salida esperada:
double[][] expected = new double[][] {
new double[] {4.0, 2.0}
};Explicación
Con un solo token, cada cabeza solo se atiende a sí misma, así que la salida concatenada es V = [2, 4]. w_o intercambia las dos columnas: [4, 2].
Cómo abordarlo
Inténtalo primero por tu cuenta y luego abre las pistas una a una.
Pista 1
Primero proyecta la entrada en queries, keys y values, y luego divide la dimensión de características en bloques iguales por cabeza.
Pista 2
Ejecuta el mismo procedimiento de atención dentro de cada cabeza de forma independiente, usando la dimensión más pequeña de la cabeza para escalar.
Pista 3
Concatena las salidas de las cabezas en su orden original de columnas antes de la proyección final, y mantén alineadas con cuidado las dimensiones de las matrices.
Complejidad esperada
tiempo O(n^2 * d + n * d^2), espacio extra O(n * d)
Referencias
Este reto está adaptado para CodeUtility. Sigue el enunciado y las restricciones indicadas arriba.