問題
マルチヘッドアテンションにより、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 を計算し、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 なので、スコアはすべて 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}
};解説
トークンが1つだけなので各ヘッドは自分自身だけに注目し、連結した出力は V = [2, 4] です。w_o が2つの列を入れ替えるので [4, 2] になります。
解き方のヒント
まずは自分で考えてから、ヒントを 1 つずつ開いてみましょう。
ヒント 1
まず入力を queries、keys、values に投影し、それから特徴次元を等しい head サイズのチャンクに分割します。
ヒント 2
各 head 内で同じ attention 手順を独立に実行し、小さい head 次元でスケールします。
ヒント 3
最終的な出力 projection の前に、head の出力を元の列順で連結し、行列の次元を慎重に揃えます。
目安の計算量
時間 O(n^2 * d + n * d^2)、追加メモリ O(n * d)
参考資料
この課題はCodeUtility用に調整されています。上記の説明と制約に従ってください。