多头掩码注意力（AI031）- 算法机考高频题型 · 注意力 / 多头题解 | AlgoMooc

AI031. 多头掩码注意力

中等

通过率 100% · 提交 3 · 通过 3

注意力多头

这道 OJ 的配套动画

多头注意力 · 对应本题的思路、过程和代码

算法机考高频题型改编 · 聚焦「注意力 / 多头」 · 难度：中等

实现多头掩码注意力：Q、K、V 各为 n×d，按列分成 h 段（每段 hd=d/h 维）。可见性掩码 mask 为 n×n 的 0/1 矩阵，mask[i][j]=1 表示位置 i 可以看到位置 j，0 表示屏蔽；第 i 行的掩码作用于第 i 个查询，所有头共用同一份掩码。每个头对可见的 j 打分 s=(Q^g_i·K^g_j)/√hd，做数值稳定的 softmax（归一化指数函数）后对 V^g 加权求和；整行被屏蔽时该行输出全 0；h 个头的结果按原顺序拼接回 d 维。

这类题属于算法机考高频题型中「注意力 / 多头」方向的高频题型，通常考察对「注意力 / 多头」的建模能力与边界条件处理。掌握本题的解题思路后，可举一反三应对同类真题方向，稳步提升机考通过率。

输入描述

第一行三个整数 n d h（h 整除 d）。随后 n 行为 Q、n 行为 K、n 行为 V（每行 d 个实数，最多两位小数）；最后 n 行、每行 n 个 0 或 1，为掩码 mask。

输出描述

n 行、每行 d 个数（四舍五入保留两位小数，单个空格分隔）；-0.00 输出 0.00。舍入规则：用语言的标准两位小数格式化（f"{x:.2f}"、printf("%.2f")、String.format("%.2f")）即可；测试数据保证每个待输出值的精确值到 0.005 舍入边界的距离不小于 1e-8，三种语言结果一致，无需自行实现舍入。

数据范围

1 <= n <= 40；1 <= d <= 16；h 整除 d；|分量| <= 8。

示例

示例 1

输入示例

2 4 2
1.00 0.00 0.50 0.50
0.00 1.00 0.50 0.50
1.00 0.00 1.00 0.00
0.00 1.00 0.00 1.00
2.00 0.00 0.00 4.00
0.00 2.00 4.00 0.00
1 1
0 1

输出示例

1.34 0.66 2.00 2.00
0.00 2.00 4.00 0.00

说明：

两个头；索引为 1 的第二行只可见 j=1

时间限制 4000 ms · 内存限制 256 MB

本平台为独立第三方培训机构，与华为技术有限公司无任何关联；课程的服务内容与权益以购买协议为准，学习效果因个人情况而异。「华为 OD」「华为可信」等仅为对岗位与考试方向的客观描述，相关商标归各自权利人所有。

问小欧 · 图解算法年卡¥99 开通 →

题库上一题下一题

未登录也能直接「运行」调试;「提交」判题需微信登录,你写的代码会自动保留、登录后回到本题继续。

字号13Ctrl+/ 注释 · Ctrl+↵ 提交

1
2

# 在此编写你的代码

写完代码点「提交」，将对全部测试用例判题。

解题思路

令每头维度 hd=d/h。按列把 Q、K、V 拆成 h 组；每组独立计算缩放点积，先应用掩码，再做减最大值的稳定 Softmax，并用概率加权 V。最后按头编号把结果拼回 d 维。

正确性说明

每个头只读取自己的连续 hd 个维度，因此各头计算互不交叉；头内 Softmax 对且仅对可见键归一化，加权和正是该头定义的输出。按原头序拼接覆盖全部 d 个维度一次，所以得到完整多头输出。

复杂度分析

时间复杂度 O(n²d)；输出矩阵占 O(nd)，当前实现逐头逐行计算，除输出外每行临时空间为 O(n)（输入 Q、K、V 共占 O(nd)）。

边界与易错点

整行掩码全为 0 时直接输出零向量；缩放因子使用 √hd 而不是 √d，并把 -0.00 规范为 0.00。

参考代码

import math
import sys


def fmt(v: float) -> str:
    s = f"{v:.2f}"
    return "0.00" if s == "-0.00" else s


def main() -> None:
    data = sys.stdin.read().split()
    it = iter(data)
    n = int(next(it)); d = int(next(it)); h = int(next(it))
    read_mat = lambda r, c: [[float(next(it)) for _ in range(c)] for _ in range(r)]
    q = read_mat(n, d)
    k = read_mat(n, d)
    v = read_mat(n, d)
    mask = [[int(next(it)) for _ in range(n)] for _ in range(n)]
    hd = d // h
    inv = 1.0 / math.sqrt(hd)
    out = [[0.0] * d for _ in range(n)]
    for g in range(h):
        lo = g * hd
        for i in range(n):
            visible = [j for j in range(n) if mask[i][j] == 1]
            if not visible:
                continue  # 全屏蔽行：该头输出保持全 0
            scores = []
            for j in visible:
                dot = sum(q[i][lo + t] * k[j][lo + t] for t in range(hd))
                scores.append(dot * inv)
            m = max(scores)
            exps = [math.exp(s - m) for s in scores]
            total = sum(exps)
            for idx, j in enumerate(visible):
                p = exps[idx] / total
                for t in range(hd):
                    out[i][lo + t] += p * v[j][lo + t]
    sys.stdout.write("\n".join(" ".join(fmt(x) for x in row) for row in out) + "\n")


if __name__ == "__main__":
    main()