1ЗадачаHard
Single-head self-attention с Q/K/V projections
Условие
Implement single-head self-attention using NumPy-style matrix operations.
Given token matrix X and projection matrices W_q, W_k, W_v, compute:
Q = X @ W_q
K = X @ W_k
V = X @ W_v
Attention = softmax(Q @ K.T / sqrt(d_k)) @ VApply softmax row-wise and return the attention output.
Решение прямо на странице
Напишите код, запустите проверки и только потом открывайте разбор.
Проверка решения
Нажмите «Запустить проверки» или Ctrl+Enter.
Показать разбор
Подсказки
- Shapes
Если X имеет shape (n, d_model), то Q и K имеют shape (n, d_k), а scores - (n, n).
- Stable softmax
Перед exp вычитайте max строки: exp(x - max_row).
Идея решения
Реализация идет в четыре шага:
- Получить Q, K, V через matrix multiplication.
- Посчитать scores = Q @ K.T / sqrt(d_k).
- Применить stable softmax по каждой строке scores.
- Умножить weights @ V и округлить результат.
Важно не забыть K.T: scores должны иметь shape (n_tokens, n_tokens).
Эталонный код
def single_head_self_attention(
X: list[list[float]],
W_q: list[list[float]],
W_k: list[list[float]],
W_v: list[list[float]],
) -> list[list[float]]:
import math
def transpose(matrix: list[list[float]]) -> list[list[float]]:
return [list(col) for col in zip(*matrix)]
def matmul(a: list[list[float]], b: list[list[float]]) -> list[list[float]]:
b_t = transpose(b)
return [
[sum(x * y for x, y in zip(row, col)) for col in b_t]
for row in a
]
Q = matmul(X, W_q)
K = matmul(X, W_k)
V = matmul(X, W_v)
d_k = len(Q[0])
raw_scores = matmul(Q, transpose(K))
scaled_scores = [[value / math.sqrt(d_k) for value in row] for row in raw_scores]
weights: list[list[float]] = []
for row in scaled_scores:
row_max = max(row)
exps = [math.exp(value - row_max) for value in row]
denom = sum(exps)
weights.append([value / denom for value in exps])
output = matmul(weights, V)
return [[round(value, 6) for value in row] for row in output]Сложность
Время: O(n * d_model * (d_k + d_k + d_v) + n^2 * (d_k + d_v)). Память: O(n^2).
Сначала считаются Q/K/V projections, затем матрица scores размера n x n и умножение attention weights на V.