Multi-Head-Self-Attention

Problem

Multi-Head-Attention lässt einen Transformer auf mehrere Dinge gleichzeitig achten. Mit Eingabe X (n Tokens × d Spalten) und Gewichtsmatrizen der Größe d × d:

Q = X × w_q, K = X × w_k, V = X × w_v
Teile die d Spalten in num_heads Heads mit je d_k = d / num_heads aufeinanderfolgenden Spalten (Head h verwendet die Spalten h×d_k bis (h+1)×d_k - 1).
Führe für jeden Head skalierte Skalarprodukt-Attention auf seinen Spalten aus: Für Token i ist score_j = (Q_i · K_j) / √d_k über die Spalten dieses Heads, Softmax über alle n Tokens j und Ausgabe = Σ weight_j × V_j (Spalten dieses Heads).
Setze die Ausgaben der Heads in ihre Spalten zurück (Konkatenation) und multipliziere mit w_o.

Implementiere multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads).

Eingabe

x: n × d-Eingabematrix.

w_q, w_k, w_v, w_o: d × d-Gewichtsmatrizen.

num_heads: Anzahl der Heads (d ist durch sie teilbar).

Ausgabe

Gib die n × d-Ausgabematrix zurück.

Einschränkungen

1 ≤ n ≤ 30, 1 ≤ d ≤ 32, 1 ≤ num_heads ≤ d und d ist durch num_heads teilbar

-5 ≤ jeder Wert ≤ 5

Beispiele

Beispiel 1

Eingabe:
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

Erwartete Ausgabe:
expected = [
    [0.5, 0.5],
    [0.5, 0.5]
]

Eingabe:
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;

Erwartete Ausgabe:
std::vector<std::vector<double>> expected = {
    {0.5, 0.5},
    {0.5, 0.5}
};

Eingabe:
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;

Erwartete Ausgabe:
double[][] expected = new double[][] {
    new double[] {0.5, 0.5},
    new double[] {0.5, 0.5}
};

Erläuterung

w_q besteht nur aus Nullen, also sind alle Scores 0 und jedes Token mittelt die Value-Zeilen: V = X, also werden beide Zeilen zu [0.5, 0.5].

Beispiel 2

Eingabe:
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

Erwartete Ausgabe:
expected = [
    [4.0, 2.0]
]

Eingabe:
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;

Erwartete Ausgabe:
std::vector<std::vector<double>> expected = {
    {4.0, 2.0}
};

Eingabe:
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;

Erwartete Ausgabe:
double[][] expected = new double[][] {
    new double[] {4.0, 2.0}
};

Erläuterung

Mit nur einem Token achtet jeder Head nur auf sich selbst, also ist die konkatenierte Ausgabe V = [2, 4]. w_o vertauscht die beiden Spalten: [4, 2].

So gehst du vor

Versuch es zuerst allein und öffne dann einen Hinweis nach dem anderen.

Hinweis 1

Projiziere die Eingabe zuerst in Queries, Keys und Values und teile dann die Feature-Dimension in gleich große head-Abschnitte auf.

Hinweis 2

Führe dieselbe Attention-Prozedur in jedem Head unabhängig aus und verwende für das Skalieren die kleinere Head-Dimension.

Hinweis 3

Verkette die Head-Outputs in ihrer ursprünglichen Spaltenreihenfolge vor der finalen Ausgabeprojektion und halte die Matrixdimensionen sorgfältig im Gleichgewicht.

Erwartete Komplexität

Zeit O(n^2 * d + n * d^2), zusätzlicher Speicher O(n * d)

Referenzen

Diese Aufgabe wurde für CodeUtility angepasst. Beachte die oben genannten Vorgaben und Einschränkungen.

Das lernst du

Problem

Eingabe

Ausgabe

Einschränkungen

Beispiele

Erläuterung

Erläuterung

So gehst du vor

Referenzen

Community-Lösungen