注意力机制与 Q/K/V 计算

02 / 学习目标与先修自测

学完本课你能做什么，以及开始前需要会什么

先确认自己能否完成打分、掩码、归一化和加权计算，再用下面的先修自测检查基础。

先修自测：下面 5 题请先自己写答案，再展开对照。任何一题答不出，按括号里的位置补看 10 分钟再回来——这些能力在本课每一步都会用到。

补看入口（本站内，可直接点开）：矩阵乘法与维度核对（矩阵乘法与转置）、浮点与稳定 softmax、标准输入输出与首次独立提交（读入实数）。

03 / 概念与符号

注意力解决什么问题，Q、K、V 各是什么

模块 6 · 第 3 课 的 k 近邻用「距离」决定谁对预测有影响；注意力换成一种可训练的写法：每个位置的输出，是所有位置的值向量按「匹配程度」加权平均的结果，权重由内容本身算出来。

先看要解决的问题。一个长度为 n 的序列（比如一句话里的 n 个词元）中，第 i 个位置想从其它位置「取信息」，但不同位置的重要程度不同：与当前位置匹配的应当权重大，不相关的权重小，且所有权重加起来等于 1，这样输出仍与 V 同一量级。注意力就是「算一组和为 1 的权重，再用它们对 V 加权平均」这件事的标准做法。

维度变化是自查第 3 条的内容，请在纸上写一遍：Q、K、V 各是 (n,d)；S = Q·Kᵀ/√d 是 (n,n)；W 与 S 同形状 (n,n)；O = W·V 是 (n,d)。检查方法：矩阵乘法要求左矩阵列数等于右矩阵行数——Q (n,d) 乘 Kᵀ (d,n) 合法，W (n,n) 乘 V (n,d) 合法。

04 / 从直觉到公式

为什么用点积、为什么除以 √d、为什么先掩码后 softmax

四步公式里每一步都有原因。下面用能手算的小数字逐条说明，再给出规则。

第一步，用点积衡量匹配程度。q=(2,1) 与 k=(1,3) 的点积是 2×1 + 1×3 = 5；与 k′=(1,0) 的点积是 2。点积越大，两个向量方向越接近（在长度相近时），所以把它当作「i 想找的」与「j 能提供的」之间的匹配分数。整张分数表就是 Q·Kᵀ：第 i 行第 j 列是 Q[i]·K[j]。

第二步，除以 √d。点积是 d 个乘积的和，d 越大，分数的量级通常越大，softmax 之后权重会过度集中到最大分数上。除以 √d 是缓和这种增长的约定：当 q、k 的各分量相互独立、均值为 0 且各维方差相同（通常按单位方差建模）时，Σ q_c·k_c 的方差正比于 d，除以 √d 后方差就不再随 d 变化；它不是对任意输入的保证——q=k=全 1 向量时，d=1、4、16 的缩放后分数分别是 1、2、4，仍随 d 增长。用第 05 节的 3×2 小例（d=2）对比：第 3 行不缩放时权重为 [0.2119, 0.2119, 0.5761]，除以 √2 后为 [0.2483, 0.2483, 0.5035]——缩放让分布更平缓。题目明确写了缩放因子，按题目写：AI006 用 √d，AI031 每个头用 √hd。

第三步，softmax 把一行分数变成和为 1 的权重：每个分数取指数，再除以本行指数之和。指数保证权重为正，且分数大的权重大。实现时先减本行最大值：Python 的指数函数（math.exp）算到 math.exp(710) 直接溢出抛出 OverflowError（math.exp(709) 已约 8.2×10³⁰⁷）。会不会碰到这条线，看的是传给指数函数的分数——也就是除以 √d（或 √hd）之后的值：AI006 缩放后的分数可达 4×10⁶，必然溢出；AI031 缩放后最多 256，不会溢出（第 08 节的表给了两题的上界）。负的大分数不会上溢，只会下溢成 0。减最大值后每个指数都在 (0, 1] 之间，结果与不减时在数学上相同（分子分母同乘一个常数）。浮点上，比最大值小很多的分数（AI006 里差值可达 4×10⁶）取指数会下溢成 0，对应权重就是 0——与数学上「接近 0」的结果一致，不影响答案。

1. 可见集合 vis = { j : 第 i 行能看到 j }        AI006：j ≤ i　AI031：mask[i][j] = 1
2. 若 vis 为空 → 本行输出全 0，结束（AI031 规定）
3. 分数 s_j = Q[i]·K[j] / 缩放因子，j ∈ vis        AI006：√d　AI031：每头 √hd
4. m = vis 中最大分数；w_j = e^(s_j − m) / Σ e^(s_k − m)，k ∈ vis
5. 输出 O[i] = Σ w_j · V[j]，j ∈ vis

05 / 完整手算例

同一组输入走完四步：全量注意力，然后加因果掩码

n=3、d=2，Q=K=[[1,0],[0,1],[1,1]]，V=[[1,2],[3,4],[5,6]]，缩放因子 √2≈1.4142。所有中间值保留四位小数显示，但计算时保留完整精度（用第 06 节的检查脚本可以逐项复现）。

现在加因果掩码（AI006 的规则：第 i 行只能看 j ≤ i）。公式一个字不变，变的只是每行参与 softmax 的集合。

06 / 从步骤到程序

把四步规则写成代码：核心片段与完整参考程序

先对照步骤表看核心循环里每个变量对应第 04 节规则的哪一行，再自己写完整程序；两份完整参考程序在展开区里，写完再展开对照。

scale = math.sqrt(d)
out = []
for i in range(n):
    vis = visible(i)                     # 掩码：哪些 j 可见，以题目要求为准（因果 = 只看 j <= i）
    row = [0.0] * d
    if not vis:                          # 整行被屏蔽：按题目要求输出全 0（AI031），不执行 softmax
        out.append(row)
        continue
    scores = [dot(Q[i], K[j]) / scale for j in vis]
    mx = max(scores)                     # 稳定 softmax：先减本行（可见集合内）最大值
    ws = [math.exp(s - mx) for s in scores]
    total = sum(ws)
    for j, w in zip(vis, ws):
        w /= total                       # 归一成权重
        for p in range(d):
            row[p] += w * V[j][p]        # 按权重加权求和 V
    out.append(row)

import math

Q = [[1, 0], [0, 1], [1, 1]]
K = [[1, 0], [0, 1], [1, 1]]
V = [[1, 2], [3, 4], [5, 6]]
d = 2
# 固定例子的预期输出（与第 05 节表格一致）。改了 Q、K、V 之后：把它换成你手算的结果，或设为 None 只看打印不核对
EXPECTED = [[3.0, 4.0], [3.4067, 4.4067], [3.5105, 4.5105]]

def dot(a, b):
    return sum(x * y for x, y in zip(a, b))

for i in range(3):
    scores = [dot(Q[i], K[j]) / math.sqrt(d) for j in range(3)]
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    total = sum(exps)
    weights = [e / total for e in exps]
    out = [sum(w * V[j][c] for w, j in zip(weights, range(3))) for c in range(2)]
    print(i + 1, "分数", [round(s, 4) for s in scores], "最大值", round(m, 4), "指数", [round(e, 4) for e in exps], "和", round(total, 4))
    print("  权重", [round(w, 4) for w in weights], "输出", [round(o, 4) for o in out])
    assert abs(sum(weights) - 1.0) < 1e-9          # 通用检查：权重和为 1，任何输入都成立
    if EXPECTED is not None:
        assert [round(o, 4) for o in out] == EXPECTED[i]   # 固定例子核对；改了数据就更新 EXPECTED
# 期望输出（固定例子）：
# 1 分数 [0.7071, 0.0, 0.7071] 最大值 0.7071 指数 [1.0, 0.4931, 1.0] 和 2.4931
#   权重 [0.4011, 0.1978, 0.4011] 输出 [3.0, 4.0]
# 2 分数 [0.0, 0.7071, 0.7071] 最大值 0.7071 指数 [0.4931, 1.0, 1.0] 和 2.4931
#   权重 [0.1978, 0.4011, 0.4011] 输出 [3.4067, 4.4067]
# 3 分数 [0.7071, 0.7071, 1.4142] 最大值 1.4142 指数 [0.4931, 0.4931, 1.0] 和 1.9861
#   权重 [0.2483, 0.2483, 0.5035] 输出 [3.5105, 4.5105]

import math
import sys


def read_matrix(tokens, start, n, d):
    """从 tokens[start] 起连续读 n 行、每行 d 个实数，返回矩阵和新的读取位置。"""
    rows = []
    pos = start
    for _ in range(n):
        rows.append([float(x) for x in tokens[pos:pos + d]])
        pos += d
    return rows, pos


def dot(a, b):
    return sum(x * y for x, y in zip(a, b))


def softmax_row(scores):
    """数值稳定的 softmax：先减本行最大值，再取指数并归一化。"""
    mx = max(scores)
    exps = [math.exp(s - mx) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


def causal_attention(Q, K, V, n, d):
    scale = math.sqrt(d)
    out = []
    for i in range(n):
        visible = range(0, i + 1)                 # 因果：第 i 行只看 j <= i
        scores = [dot(Q[i], K[j]) / scale for j in visible]
        weights = softmax_row(scores)
        row = [0.0] * d
        for w, j in zip(weights, visible):
            for p in range(d):
                row[p] += w * V[j][p]
        out.append(row)
    return out


def fmt4(x):
    """保留四位小数；绝对值小于 0.00005 时输出 0.0000（避免 -0.0000）。"""
    if abs(x) < 0.00005:
        return "0.0000"
    return f"{x:.4f}"


def main():
    tokens = sys.stdin.read().split()
    n, d = int(tokens[0]), int(tokens[1])
    Q, pos = read_matrix(tokens, 2, n, d)
    K, pos = read_matrix(tokens, pos, n, d)
    V, pos = read_matrix(tokens, pos, n, d)
    for row in causal_attention(Q, K, V, n, d):
        print(" ".join(fmt4(x) for x in row))


main()

import math
import sys


def read_matrix(tokens, start, n, d, cast=float):
    rows = []
    pos = start
    for _ in range(n):
        rows.append([cast(x) for x in tokens[pos:pos + d]])
        pos += d
    return rows, pos


def softmax_row(scores):
    mx = max(scores)
    exps = [math.exp(s - mx) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


def multi_head_attention(Q, K, V, mask, n, d, h):
    hd = d // h
    scale = math.sqrt(hd)                       # 每个头按自己的维度 hd 缩放
    out = [[0.0] * d for _ in range(n)]
    for i in range(n):
        visible = [j for j in range(n) if mask[i][j] == 1]
        if not visible:                          # 整行被屏蔽：输出全 0，不做 softmax
            continue
        for g in range(h):
            lo, hi = g * hd, (g + 1) * hd        # 第 g 头的列区间 [lo, hi)
            scores = []
            for j in visible:
                s = sum(Q[i][c] * K[j][c] for c in range(lo, hi)) / scale
                scores.append(s)
            weights = softmax_row(scores)
            for w, j in zip(weights, visible):
                for c in range(lo, hi):
                    out[i][c] += w * V[j][c]
    return out


def fmt2(x):
    s = f"{x:.2f}"
    return "0.00" if s == "-0.00" else s


def main():
    tokens = sys.stdin.read().split()
    n, d, h = int(tokens[0]), int(tokens[1]), int(tokens[2])
    Q, pos = read_matrix(tokens, 3, n, d)
    K, pos = read_matrix(tokens, pos, n, d)
    V, pos = read_matrix(tokens, pos, n, d)
    mask, pos = read_matrix(tokens, pos, n, n, cast=int)
    for row in multi_head_attention(Q, K, V, mask, n, d, h):
        print(" ".join(fmt2(x) for x in row))


main()

07 / 多头拆分与拼接

第 g 头取哪些列、按什么缩放、怎么拼回去

多头注意力把 d 维按列切成 h 段，每段 hd=d/h 维，各头独立完成四步，再按头的顺序拼回 d 维。AI031 直接给出 Q、K、V，考查的就是拆分与拼接的下标。

列区间的规则：第 g 头（g 从 0 数到 h−1）取列 [g·hd, (g+1)·hd)，左闭右开。d=4、h=2 时 hd=2：头 0 取列 0、1，头 1 取列 2、3。缩放因子是本头的维度 √hd，不是总维度 √d——因为每个头的点积只累加 hd 个乘积。同一行的掩码对所有头相同。

行 0 = 头 0 的 [1.3395, 0.6605] ‖ 头 1 的 [2.0000, 2.0000] → 保留两位 → 1.34 0.66 2.00 2.00
行 1 = 头 0 的 [0.0000, 2.0000] ‖ 头 1 的 [4.0000, 0.0000] → 0.00 2.00 4.00 0.00
与题目样例输出逐位一致。

08 / 边界、反例与复杂度

四种错误做法各输出什么：三种用 AI031 样例，一种用样例的掩码变式

下面每一行都是一个具体的错误实现。前三行的输入是 AI031 的样例（第 07 节那组数）；第四行必须把样例第 0 行的掩码改成 0 0 才能暴露错误——原样例没有整行被屏蔽的行。你可以把参考程序改成对应写法自己验证。

复杂度：每行对每个可见位置做一次 d 维点积，再做一次 d 维加权，总量是 O(n²·d)；多头版本是 O(h·n²·hd) = O(n²·d)。按完整矩阵估算，AI006 打分阶段不超过 80×80×16 ≈ 1.0×10⁵ 次乘加，AI031 不超过 40×40×16 ≈ 2.6×10⁴ 次；加权阶段另有同量级运算。空间 O(n·d) 存三个矩阵和输出，加 O(n²) 存掩码。避免在内层重复构造矩阵或计算缩放系数；耗时还取决于实现与评测环境，提交前需在目标语言核对。

09 / 渐进练习与参考答案

跟做 → 改一个条件 → 独立实现 → 迁移

每题先在纸上或文件里做完，再展开答案。答案给出计算过程或代码，以及做错时最可能错在哪一步。

练习 1（跟做）：仍用第 05 节的 Q=K=[[1,0],[0,1],[1,1]]，把 V 换成 [[1,0],[0,1],[2,2]]，计算全量注意力（不带掩码）第 0 行的输出。提示：权重不变，只有第四步的加权和要重算。

练习 2（改一个条件）：AI031 样例的掩码改成全 1（两行都能看到 j=0 和 j=1），其它不变。计算第 1 行的输出，保留两位小数。

练习 3（改一个条件）：AI031 样例改成 h=1（单头，hd=4，缩放 √4=2），掩码不变。计算第 0 行的输出，直接保留两位小数。不要先保留四位再取两位：两次舍入可能与一次舍入不同，例如 1.23496 直接取两位是 1.23，先取四位得 1.2350 再取两位就成了 1.24。

练习 4（独立实现）：写出 AI031 需要的两个辅助函数并通过断言——visible_from_mask(row) 由一行 0/1 掩码返回可见位置列表；fmt2(x) 保留两位小数并把 -0.00 归一为 0.00。

def visible_from_mask(row):
    # 待完成：row 是一行 0/1 掩码，返回可见的 j 列表（升序）；1 表示可见
    ...

def fmt2(x):
    # 待完成：保留两位小数，"-0.00" 归一为 "0.00"
    ...

assert visible_from_mask([1, 1]) == [0, 1]
assert visible_from_mask([0, 1]) == [1]
assert visible_from_mask([0, 0]) == []
assert fmt2(1.339) == "1.34" and fmt2(-0.004) == "0.00" and fmt2(-0.006) == "-0.01"

def visible_from_mask(row):
    return [j for j, m in enumerate(row) if m == 1]

def fmt2(x):
    s = f"{x:.2f}"
    return "0.00" if s == "-0.00" else s

练习 5（迁移）：AI006 的第二组样例是 n=3、d=1，Q=[1,1,1]，K=[0,0,10]，V=[1,2,100]，期望输出 1.0000、1.5000、99.9911。不运行程序，解释第 2 行（下标 2）为什么是 99.9911 而不是 100.0000，并算出它的三个权重。

10 / 读题要求、题解入口与复习自评

两道题的输入输出对照，以及完成本课之后怎么复习

提交前逐项核对输入布局、掩码规则与输出精度；具体要求以题目页为准。

参考程序与题解：需要对照解法时，展开第 06 节的两份完整参考程序（AI006、AI031）。AI006 的题目页另有思路与 Python / Java / C++ 参考代码。

复习与自评：「学习完成检查」六条是自评，勾选表示你认为自己能做到，它不会改变题目的通过（AC）状态。本课记为完成的条件：必做题 AI006、AI031 都通过，并勾选全部六条学习完成检查；复习题不计入完成状态。复习时用下面三个问题自测：① 不看正文，写出第 04 节的五行规则；② 不看表格，重算第 05 节因果掩码下第 1 行的权重 [0.3302, 0.6698]；③ 说出 √hd 与 √d 在 AI031 样例上分别给出什么输出（第 08 节第一行）。答不出哪一条，就回到对应的节重读，再做第 09 节对应的练习。