Problema
A atenção multi-cabeça permite que um Transformer atenda a várias coisas ao mesmo tempo. Com a entrada X (n tokens × d colunas) e matrizes de pesos d × d:
- Q = X × w_q, K = X × w_k, V = X × w_v
- Divida as d colunas em
num_headscabeças de d_k = d / num_heads colunas consecutivas (a cabeça h usa as colunas h×d_k até (h+1)×d_k - 1). - Para cada cabeça, execute a atenção por produto escalar escalado nas suas colunas: para o token i, score_j = (Q_i · K_j) / √d_k nas colunas dessa cabeça, softmax sobre todos os n tokens j, e saída = Σ weight_j × V_j (colunas dessa cabeça).
- Coloque as saídas das cabeças de volta em suas colunas (concatenação) e multiplique por w_o.
Implemente 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: Matrizes de pesos d × d.
num_heads: O número de cabeças (d é divisível por ele).
Saída
Retorne a matriz de saída n × d.
Restrições
1 ≤ n ≤ 30, 1 ≤ d ≤ 32, 1 ≤ num_heads ≤ d e d é divisível por num_heads
-5 ≤ cada valor ≤ 5
Exemplos
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
Saída 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;
Saída 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;
Saída esperada:
double[][] expected = new double[][] {
new double[] {0.5, 0.5},
new double[] {0.5, 0.5}
};Explicação
w_q é toda zero, então todas as pontuações são 0 e cada token faz a média das linhas de valores: V = X, então as duas linhas se tornam [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
Saída 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;
Saída 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;
Saída esperada:
double[][] expected = new double[][] {
new double[] {4.0, 2.0}
};Explicação
Com um único token, cada cabeça atende só a si mesma, então a saída concatenada é V = [2, 4]. w_o troca as duas colunas: [4, 2].
Como abordar
Tente sozinho primeiro e depois abra uma dica de cada vez.
Dica 1
Projete a entrada em queries, keys e values primeiro, depois divida a dimensão de features em blocos iguais do tamanho de cada head.
Dica 2
Execute o mesmo procedimento de attention dentro de cada head de forma independente, usando a dimensão menor da head para o scale.
Dica 3
Conecte os outputs das heads na ordem original das colunas antes da projeção final de saída e mantenha as dimensões das matrizes alinhadas com cuidado.
Complexidade esperada
tempo O(n^2 * d + n * d^2), espaço extra O(n * d)
Referências
Este desafio foi adaptado para o CodeUtility. Siga o enunciado e as restrições acima.