Problema
L'attenzione multi-head permette a un Transformer di prestare attenzione a più cose contemporaneamente. Con input X (n token × d colonne) e matrici dei pesi d × d:
- Q = X × w_q, K = X × w_k, V = X × w_v
- Dividi le d colonne in
num_headshead di d_k = d / num_heads colonne consecutive (la head h usa le colonne da h×d_k a (h+1)×d_k - 1). - Per ogni head, esegui l'attenzione a prodotto scalare scalato sulle sue colonne: per il token i, score_j = (Q_i · K_j) / √d_k sulle colonne di quella head, softmax su tutti gli n token j, e output = Σ weight_j × V_j (colonne di quella head).
- Rimetti gli output delle head nelle loro colonne (concatenazione) e moltiplica per w_o.
Implementa multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads).
Input
x: Matrice di input n × d.
w_q, w_k, w_v, w_o: Matrici dei pesi d × d.
num_heads: Il numero di head (d è divisibile per esso).
Output
Restituisci la matrice di output n × d.
Vincoli
1 ≤ n ≤ 30, 1 ≤ d ≤ 32, 1 ≤ num_heads ≤ d e d è divisibile per num_heads
-5 ≤ ogni valore ≤ 5
Esempi
Input:
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
Output atteso:
expected = [
[0.5, 0.5],
[0.5, 0.5]
]Input:
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;
Output atteso:
std::vector<std::vector<double>> expected = {
{0.5, 0.5},
{0.5, 0.5}
};Input:
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;
Output atteso:
double[][] expected = new double[][] {
new double[] {0.5, 0.5},
new double[] {0.5, 0.5}
};Spiegazione
w_q è tutta a zero, quindi tutti i punteggi sono 0 e ogni token fa la media delle righe dei value: V = X, quindi entrambe le righe diventano [0.5, 0.5].
Input:
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
Output atteso:
expected = [
[4.0, 2.0]
]Input:
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;
Output atteso:
std::vector<std::vector<double>> expected = {
{4.0, 2.0}
};Input:
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;
Output atteso:
double[][] expected = new double[][] {
new double[] {4.0, 2.0}
};Spiegazione
Con un solo token, ogni head presta attenzione solo a se stessa, quindi l'output concatenato è V = [2, 4]. w_o scambia le due colonne: [4, 2].
Come affrontarlo
Prova prima da solo, poi apri un suggerimento alla volta.
Suggerimento 1
Proietta prima l'input in queries, keys e values, poi dividi la dimensione delle feature in chunk uguali grandi quanto un head.
Suggerimento 2
Esegui la stessa procedura di attention in ogni head in modo indipendente, usando la dimensione ridotta dell'head per la scala.
Suggerimento 3
Concatena gli output degli head nel loro ordine originale delle colonne prima della proiezione finale, e mantieni allineate con cura le dimensioni delle matrici.
Complessità attesa
tempo O(n^2 * d + n * d^2), spazio aggiuntivo O(n * d)
Riferimenti
Questa sfida è adattata per CodeUtility. Segui la traccia e i vincoli sopra.