Многоголовое внимание позволяет Transformer обращать внимание на несколько вещей одновременно. Для входа X (n токенов × d столбцов) и матриц весов d × d:
- Q = X × w_q, K = X × w_k, V = X × w_v
- Разбейте d столбцов на
num_headsголов по d_k = d / num_heads последовательных столбцов (голова h использует столбцы с h×d_k по (h+1)×d_k - 1). - Для каждой головы выполните масштабированное скалярное внимание на её столбцах: для токена i score_j = (Q_i · K_j) / √d_k по столбцам этой головы, softmax по всем n токенам j, и выход = Σ weight_j × V_j (столбцы этой головы).
- Верните выходы голов в их столбцы (конкатенация) и умножьте на w_o.
Реализуйте multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads).
Входные данные
x: Входная матрица n × d.
w_q, w_k, w_v, w_o: Матрицы весов d × d.
num_heads: Число голов (d делится на него).
Выходные данные
Верните выходную матрицу n × d.
Ограничения
1 ≤ n ≤ 30, 1 ≤ d ≤ 32, 1 ≤ num_heads ≤ d, и d делится на num_heads
-5 ≤ каждое значение ≤ 5
Примеры
Ввод:
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
Ожидаемый вывод:
expected = [
[0.5, 0.5],
[0.5, 0.5]
]Ввод:
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;
Ожидаемый вывод:
std::vector<std::vector<double>> expected = {
{0.5, 0.5},
{0.5, 0.5}
};Ввод:
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;
Ожидаемый вывод:
double[][] expected = new double[][] {
new double[] {0.5, 0.5},
new double[] {0.5, 0.5}
};Пояснение
w_q состоит из нулей, поэтому все оценки равны 0, и каждый токен усредняет строки значений: V = X, поэтому обе строки становятся [0.5, 0.5].
Ввод:
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
Ожидаемый вывод:
expected = [
[4.0, 2.0]
]Ввод:
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;
Ожидаемый вывод:
std::vector<std::vector<double>> expected = {
{4.0, 2.0}
};Ввод:
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;
Ожидаемый вывод:
double[][] expected = new double[][] {
new double[] {4.0, 2.0}
};Пояснение
При одном токене каждая голова обращает внимание только на себя, поэтому объединённый выход — V = [2, 4]. w_o меняет два столбца местами: [4, 2].
Ссылки
Это задание адаптировано для CodeUtility. Следуйте условию и ограничениям выше.