大模型 Attention 模块开发

已知大模型常用的 Attention 模块 定义如下：

此处考虑二维情况，其中：

这道题表面上是矩阵运算模拟，但因为有「全 1 输入 + 上三角权重」的特殊初始化，实际蕴含一条极其优雅的化简路径。

推导过程：

第一步：计算 $\mathbf{q}$ 向量

$W$ 是 $m \times h$ 的上三角全 1 矩阵，$W_{ij} = 1$ 当且仅当 $i \leq j$。因此第 $j$ 列的 1 的个数等于满足 $0 \leq i \leq \min(j, m-1)$ 的行数：

以 $m=4,\; h=5$ 为例，$W$ 是 $4 \times 5$ 的上三角全 1 矩阵：

$Q = XW$，而 $X$ 全 1，所以 $q_j$ 就是 $W$ 第 $j$ 列的列和。逐列数 1 的个数：

列索引越大，从上往下覆盖的行越多（$j+1$），但不会超过总行数 $m$，所以 $q_j = \min(j+1,\; m)$。

第二步：分析 $QK^T$

由于 $Q$ 的 $n$ 行完全相同（都是 $\mathbf{q}$），$K$ 也是，所以 $QK^T$ 的每个元素都是 $\|\mathbf{q}\|^2 = \sum_{k=0}^{h-1} q_k^2$。

第三步：简化 softmax

$QK^T / \sqrt{h}$ 的所有元素相同，设为 $c$。softmax 按行归一化：每行有 $h$ 个 $c$，所以每个元素变为 $c / (h \cdot c) = 1/h$。

第四步：计算 $Y$

因此 $Y$ 每行就是 $\mathbf{q}$，总和为 $n \times \sum_{j=0}^{h-1} q_j$。

基于上述分析，最终实现只需一行核心计算：

n, m, h = map(int, input().split())
ans = n * sum(min(j + 1, m) for j in range(h))
print(round(ans))
  

import math

n, m, h = map(int, input().split())

# 构造全 1 矩阵 X
X = [[1] * m for _ in range(n)]
# 构造上三角全 1 矩阵 W
W = [[1 if i <= j else 0 for j in range(h)] for i in range(m)]

# 矩阵乘法
def matmul(A, B):
    r1, c1 = len(A), len(A[0])
    c2 = len(B[0])
    C = [[0.0] * c2 for _ in range(r1)]
    for i in range(r1):
        for k in range(c1):
            for j in range(c2):
                C[i][j] += A[i][k] * B[k][j]
    return C

Q = matmul(X, W)
K = matmul(X, W)
V = matmul(X, W)

# QK^T
S = matmul(Q, [[K[j][i] for j in range(n)] for i in range(h)])
# 除以 sqrt(h)
for i in range(n):
    for j in range(n):
        S[i][j] /= math.sqrt(h)

# softmax
A = [[0.0] * n for _ in range(n)]
for i in range(n):
    row_sum = sum(S[i])
    for j in range(n):
        A[i][j] = S[i][j] / row_sum

# Y = A @ V
Y = matmul(A, V)
print(round(sum(sum(row) for row in Y)))
      

在本题的数据范围（$n, m, h < 100$）下，暴力模拟完全可以 AC。但数学化简的思路揭示了一个有趣的事实：当输入具有特殊结构时，Attention 的矩阵运算可以退化到近乎平凡的计算。