Đề bài
Multi-head attention cho phép Transformer chú ý đến nhiều thứ cùng lúc. Với đầu vào X (n token × d cột) và các ma trận trọng số kích thước d × d:
- Q = X × w_q, K = X × w_k, V = X × w_v
- Chia d cột thành
num_headsđầu, mỗi đầu gồm d_k = d / num_heads cột liên tiếp (đầu h dùng các cột từ h×d_k đến (h+1)×d_k - 1). - Với mỗi đầu, chạy scaled dot-product attention trên các cột của nó: với token i, score_j = (Q_i · K_j) / √d_k trên các cột của đầu đó, softmax trên cả n token j, và đầu ra = Σ weight_j × V_j (các cột của đầu đó).
- Đặt đầu ra của các đầu vào đúng cột của chúng (ghép lại) rồi nhân với w_o.
Hãy cài đặt multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads).
Đầu vào
x: Ma trận đầu vào n × d.
w_q, w_k, w_v, w_o: Các ma trận trọng số d × d.
num_heads: Số đầu (d chia hết cho số này).
Đầu ra
Trả về ma trận đầu ra n × d.
Ràng buộc
1 ≤ n ≤ 30, 1 ≤ d ≤ 32, 1 ≤ num_heads ≤ d và d chia hết cho num_heads
-5 ≤ mỗi giá trị ≤ 5
Ví dụ
Đầu vào:
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
Kết quả mong đợi:
expected = [
[0.5, 0.5],
[0.5, 0.5]
]Đầu vào:
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;
Kết quả mong đợi:
std::vector<std::vector<double>> expected = {
{0.5, 0.5},
{0.5, 0.5}
};Đầu vào:
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;
Kết quả mong đợi:
double[][] expected = new double[][] {
new double[] {0.5, 0.5},
new double[] {0.5, 0.5}
};Giải thích
w_q toàn số 0, nên mọi điểm bằng 0 và mỗi token lấy trung bình các hàng value: V = X, nên cả hai hàng đều thành [0.5, 0.5].
Đầu vào:
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
Kết quả mong đợi:
expected = [
[4.0, 2.0]
]Đầu vào:
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;
Kết quả mong đợi:
std::vector<std::vector<double>> expected = {
{4.0, 2.0}
};Đầu vào:
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;
Kết quả mong đợi:
double[][] expected = new double[][] {
new double[] {4.0, 2.0}
};Giải thích
Chỉ có một token nên mỗi đầu chỉ chú ý đến chính nó, đầu ra ghép lại là V = [2, 4]. w_o hoán đổi hai cột: [4, 2].
Hướng tiếp cận
Hãy tự thử trước, sau đó mở từng gợi ý một.
Gợi ý 1
Trước hết chiếu input thành queries, keys, và values, rồi chia chiều feature thành các chunk bằng nhau theo head.
Gợi ý 2
Chạy cùng quy trình attention độc lập trong mỗi head, dùng head dimension nhỏ hơn để scale.
Gợi ý 3
Nối output của các head theo đúng thứ tự cột ban đầu trước phép chiếu output cuối cùng, và giữ kích thước matrix khớp cẩn thận.
Độ phức tạp mong đợi
thời gian O(n^2 * d + n * d^2), bộ nhớ phụ O(n * d)
Tham khảo
Thử thách này đã được điều chỉnh cho CodeUtility. Hãy làm theo đề bài và các ràng buộc ở trên.