P5337. 滑动窗口注意力机制

场次：2026年9月2日 · 来源：CodeFun2000 P5337

题目描述

标准的 $Transformer$ 自注意力计算为：

$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

对于序列长度为 $N$、维度为 $d_k$ 的输入，其时间复杂度和空间复杂度分别为 $O(N^2 \cdot d_k)$ 和 $O(N^2)$。

近期研究表明，大多数任务中的依赖关系具有局部性，远距离 token 通常可以通过多层堆叠间接完成交互。因此，可以通过限制每个 token 的注意力范围来显著降低计算复杂度。

滑动窗口注意力机制是一种典型方法，其原理如下：

1. 为每个 $token$ 设定固定大小为 $w$ 的对称窗口。
2. 位置 $i$ 的 $query$ 只与满足 $|i-j| \le w$ 的位置 $j$ 的 $key$ 计算注意力，其中 $w \ll N$。
3. 为保证数值稳定性，在计算 $Softmax$ 前减去当前窗口内的最大值：

$$\text{softmax}(x_i) = \frac{e^{x_i-\max(x)}}{\sum_j e^{x_j-\max(x)}}$$

以 $(N=8,w=1)$ 为例，滑动窗口注意力可视化如下：

位置：$0$ $1$ $2$ $3$ $4$ $5$ $6$ $7$

$Q0: [x] [x] [] [] [] [] [] [] \rightarrow$ 窗口 $[0,1]$，左边界限制为 $0$

$Q1: [x] [x] [x] [] [] [] [] [] \rightarrow$ 窗口 $[0,1,2]$

$Q2: [] [x] [x] [x] [] [] [] [] \rightarrow$ 窗口 $[1,2,3]$

$Q3: [] [] [x] [x] [x] [] [] [] \rightarrow$ 窗口 $[2,3,4]$

$Q4: [] [] [] [x] [x] [x] [] [] \rightarrow$ 窗口 $[3,4,5]$

$Q5: [] [] [] [] [x] [x] [x] [] \rightarrow$ 窗口 $[4,5,6]$

$Q6: [] [] [] [] [] [x] [x] [x] \rightarrow$ 窗口 $[5,6,7]$

$Q7: [] [] [] [] [] [] [x] [x] \rightarrow$ 窗口 $[6,7]$，右边界限制为 $7$

其中 $[x] =$ 计算注意力，$[] =$ 掩码为 $0$（不计算）。

输入格式

第一行输入四个整数：

$batch_size$, $seq_len$, $d_k$, $window_size$

后续包含 $batch_size$ 组数据，每组依次包含：

$seq_len$ 行，每行 $d_k$ 个浮点数，表示 Query 矩阵。

$seq_len$ 行，每行 $d_k$ 个浮点数，表示 Key 矩阵。

$seq_len$ 行，每行 $d_k$ 个浮点数，表示 Value 矩阵。

数据范围

$batch_size \ge 1$

$seq_len \ge 1$

$d_k \ge 1$

$1\le window_size \le seq_len$

输出格式

输出 $batch_size$ 组结果，每组包含 $seq_len$ 行，每行 $d_k$ 个浮点数，表示该位置的输出向量。

组间用一个空行分隔，最后一组后无空行。

最终结果保留 $2$ 位小数，四舍五入。

注意力打分及 $Softmax$ 权重在题目说明中保留 $6$ 位小数展示，实际计算过程中无需提前进行舍入。

如果输入违反约束，输出 $0$。

示例

示例 1

输入：

1 4 4 1
0.1 0.2 0.3 0.4
0.2 0.3 0.4 0.5
0.3 0.4 0.5 0.6
0.4 0.5 0.6 0.7
0.1 0.0 0.1 0.0
0.0 0.1 0.0 0.1
0.1 0.1 0.0 0.0
0.0 0.0 0.1 0.1
1.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 1.0 0.0
0.0 0.0 0.0 1.0

输出：

0.50 0.50 0.00 0.00
0.33 0.34 0.33 0.00
0.00 0.33 0.33 0.34
0.00 0.00 0.50 0.50

说明：

输入 $batch_size=1$，$seq_len=4$，$d_k=4$，$window_size=1$。

缩放因子 $scale=\sqrt{d_k}=2$。

$$Q=\begin{bmatrix} 0.1 & 0.2 & 0.3 & 0.4\ 0.2 & 0.3 & 0.4 & 0.5\ 0.3 & 0.4 & 0.5 & 0.6\ 0.4 & 0.5 & 0.6 & 0.7 \end{bmatrix}$$

$$K=\begin{bmatrix} 0.1 & 0.0 & 0.1 & 0.0\ 0.0 & 0.1 & 0.0 & 0.1\ 0.1 & 0.1 & 0.0 & 0.0\ 0.0 & 0.0 & 0.1 & 0.1 \end{bmatrix}$$

$$V=\begin{bmatrix} 1.0 & 0.0 & 0.0 & 0.0\ 0.0 & 1.0 & 0.0 & 0.0\ 0.0 & 0.0 & 1.0 & 0.0\ 0.0 & 0.0 & 0.0 & 1.0 \end{bmatrix}$$

根据窗口大小 $window_size=1$，注意力打分为：

位置 $0$ 的权重：

$score[0]=[Q[0]K[0],Q[0]K[1]]/scale=[0.020000,0.030000]$

归一化后为：

$[0.497500,0.502500]$

位置 $1$ 的权重：

$score[1]=[Q[1]K[0],Q[1]K[1],Q[1]K[2]]/scale=[0.030000,0.040000,0.025000]$

归一化后为：

$[0.332772,0.336116,0.331112]$

位置 $2$ 的权重：

$score[2]=[Q[2]K[1],Q[2]K[2],Q[2]K[3]]/scale=[0.050000,0.035000,0.055000]$

归一化后为：

$[0.334434,0.329455,0.336111]$

位置 $3$ 的权重：

$score[3]=[Q[3]K[2],Q[3]K[3]]/scale=[0.045000,0.060000]$

归一化后为：

$[0.496250,0.503750]$

上述归一化权重保留 $6$ 位小数精度。

最后使用归一化后的权重对对应位置的 $V$ 进行加权求和，并将最终结果保留 $2$ 位小数：

$output[0]=0.497500\times V[0]+0.502500\times V[1]$

$=[0.50,0.50,0.00,0.00]$

$output[1]=0.332772\times V[0]+0.336116\times V[1]+0.331112\times V[2]$

$=[0.33,0.34,0.33,0.00]$

$output[2]=0.334434\times V[1]+0.329455\times V[2]+0.336111\times V[3]$

$=[0.00,0.33,0.33,0.34]$

$output[3]=0.496250\times V[2]+0.503750\times V[3]$

$=[0.00,0.00,0.50,0.50]$

示例 2

输入：

1 2 2 0
0.3 0.4
0.4 0.5
0.5 0.6
0.6 0.7
0.1 0.0
0.0 0.1

输出：

0

说明：

输入 $batch_size=1$，$seq_len=2$，$d_k=2$，$window_size=0$。

由于 $window_size=0$ 违反输入范围要求，因此输出 $0$。

提示

仅能使用编程语言内置的函数。

---

解题思路

第一步：理解问题本质

标准自注意力会让每个 query 与序列中的所有 key 计算分数：

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

当序列长度为 $N$ 时，需要计算 $N\times N$ 个分数，时间复杂度和空间复杂度分别是 $O(N^2d_k)$ 和 $O(N^2)$。

滑动窗口注意力利用“远距离依赖可以通过多层堆叠间接完成”这一观察，把每个位置的注意力范围限制在左右各 $w$ 个位置内。这样每个 query 只需要和至多 $2w+1$ 个 key 计算分数，复杂度降为 $O(Nwd_k)$。

第二步：暴力解法——完整注意力

最直接的做法是不利用窗口，仍然计算完整的 $N\times N$ 注意力矩阵：

# 每个 query 都要和全部 N 个 key 打分，于是产生 N×N 的分数矩阵
for i in range(N):
    scores = []
    # 除以 sqrt(d) 是缩放：维度越大点积越大，不除会让 softmax 被最大项独占
    for j in range(N):
        scores.append(dot(Q[i], K[j]) / sqrt(d))
    # softmax 在这一整行上归一化，权重之和为 1
    weights = softmax(scores)
    # 用这组权重对全部 V 向量加权求和，得到位置 i 的输出向量
    output[i] = weighted_sum(weights, V)

这个做法能得到正确答案，但每个位置都要扫描全部 $N$ 个 key。当 $N$ 很大时，$N^2$ 的分数计算和存储都会成为瓶颈。

第三步：优化解法——只计算窗口内的分数

对于位置 $i$，需要参与计算的 key 下标范围是：

$$left=\max(0,i-w),\qquad right=\min(N-1,i+w)$$

注意窗口在序列两端会被截断，而不是补零：

• $i=0$ 时，左边没有位置，窗口是 $[0,w]$；
• $i=N-1$ 时，右边没有位置，窗口是 $[N-1-w,N-1]$。

每个位置只计算 $right-left+1$ 个分数，然后在这个窗口内做 Softmax。总计算量从 $O(N^2d_k)$ 降到 $O(Nwd_k)$。

第四步：数值稳定性与输出细节

Softmax 的公式是：

$$\text{softmax}(x_i)=\frac{e^{x_i-\max(x)}}{\sum_j e^{x_j-\max(x)}}$$

减去窗口内最大值不会改变结果，但可以避免 $e^{x_i}$ 在数值较大时溢出。即使题目数据范围不大，也建议保留这一步，保证实现的通用性。

另外需要注意：

1. 每个 batch 单独计算，batch 之间输出一个空行；
2. 最终结果保留两位小数；
3. 结果接近 $0$ 时不要输出 -0.00；
4. 如果 window_size 不在 $[1,seq\_len]$ 内，直接输出 0。

---

完整代码实现

import math
import sys


def format_number(value: float) -> str:
    """保留两位小数，并避免输出 -0.00。"""
    if abs(value) < 0.005:
        value = 0.0
    return f"{value:.2f}"


def solve():
    # 输入既有整数（头部 4 个参数）也有浮点数（Q/K/V），
    # 所以先按空白切成字符串，再按需要分别转换。
    tokens = sys.stdin.buffer.read().split()
    if len(tokens) < 4:
        print(0)
        return

    # 头部：batch 数、序列长度、每个 token 的维度、窗口半径
    batch_size, seq_len, d_k, window_size = map(int, tokens[:4])

    # 题目约束：batch_size >= 1、seq_len >= 1、d_k >= 1、
    # 1 <= window_size <= seq_len。任何一条不满足都输出 0。
    if batch_size < 1 or seq_len < 1 or d_k < 1:
        print(0)
        return
    if window_size < 1 or window_size > seq_len:
        print(0)
        return

    # 每个 batch 有 Q、K、V 三个 seq_len x d_k 矩阵。
    needed = 3 * batch_size * seq_len * d_k
    if len(tokens) < 4 + needed:
        print(0)
        return

    pos = 4
    # 缩放因子 1 / sqrt(d_k)，在算完点积后统一除。
    scale = math.sqrt(d_k)
    groups = []

    for _ in range(batch_size):
        # 读取当前 batch 的 Query 矩阵，形状 seq_len x d_k。
        query = []
        for _ in range(seq_len):
            query.append([float(x) for x in tokens[pos:pos + d_k]])
            pos += d_k

        # 读取 Key 矩阵，形状 seq_len x d_k。
        key = []
        for _ in range(seq_len):
            key.append([float(x) for x in tokens[pos:pos + d_k]])
            pos += d_k

        # 读取 Value 矩阵，形状 seq_len x d_k。
        value = []
        for _ in range(seq_len):
            value.append([float(x) for x in tokens[pos:pos + d_k]])
            pos += d_k

        lines = []
        for i in range(seq_len):
            # 位置 i 的窗口是 [i-w, i+w]。
            # 序列两端没有位置，所以窗口要截断，而不是补零。
            left = max(0, i - window_size)
            right = min(seq_len - 1, i + window_size)

            # 只计算窗口内的注意力分数：
            # score = Q[i] · K[j] / sqrt(d_k)
            scores = []
            qi = query[i]
            for j in range(left, right + 1):
                dot = 0.0
                kj = key[j]
                for t in range(d_k):
                    dot += qi[t] * kj[t]
                scores.append(dot / scale)

            # 数值稳定的 Softmax：
            # 先减去窗口内最大值，再取 exp，最后归一化。
            max_score = max(scores)
            exp_scores = [math.exp(s - max_score) for s in scores]
            total = sum(exp_scores)
            weights = [e / total for e in exp_scores]

            # 用窗口内的权重对 V 加权求和：
            # output[t] = sum_j weight[j] * V[j][t]
            output = [0.0] * d_k
            for offset, j in enumerate(range(left, right + 1)):
                weight = weights[offset]
                vj = value[j]
                for t in range(d_k):
                    output[t] += weight * vj[t]

            lines.append(" ".join(format_number(x) for x in output))

        groups.append("\n".join(lines))

    # batch 之间用空行分隔，最后一组后没有空行
    print("\n\n".join(groups))


if __name__ == "__main__":
    solve()

NumPy 版本

以下版本使用 NumPy。只有在评测环境允许第三方库时才能提交；CodeFun2000 的普通 Python 环境不一定安装 NumPy。

思路与上面的纯 Python 版完全一致，区别在于把「逐位置、逐 key」的循环换成张量运算：一个 batch 的 $Q/K/V$ 读成三维数组，$K$、$V$ 两侧各补 $w$ 行 $0$ 之后，用 sliding_window_view 得到形状 $(N,d_k,2w+1)$ 的窗口视图，于是所有位置的窗口分数 $(N,2w+1)$、Softmax 权重、加权输出都能一次算完。两端越界的槽位用 $-\infty$ 掩掉，Softmax 的归一化范围仍然只限于真实窗口。

import math
import sys

import numpy as np


def format_number(value: float) -> str:
    """保留两位小数，并避免输出 -0.00。"""
    if abs(value) < 0.005:
        value = 0.0
    return f"{value:.2f}"


def solve():
    # 输入既有整数（头部 4 个参数）也有浮点数（Q/K/V），
    # 所以先按空白切成字符串，再按需要分别转换。
    tokens = sys.stdin.buffer.read().split()
    if len(tokens) < 4:
        print(0)
        return

    # 头部：batch 数、序列长度、每个 token 的维度、窗口半径
    batch_size, seq_len, d_k, window_size = (int(x) for x in tokens[:4])

    # 题目约束：batch_size >= 1、seq_len >= 1、d_k >= 1、
    # 1 <= window_size <= seq_len。任何一条不满足都输出 0。
    if batch_size < 1 or seq_len < 1 or d_k < 1:
        print(0)
        return
    if window_size < 1 or window_size > seq_len:
        print(0)
        return

    # 每个 batch 有 Q、K、V 三个 seq_len x d_k 矩阵。
    needed = 3 * batch_size * seq_len * d_k
    if len(tokens) < 4 + needed:
        print(0)
        return

    # 一次把所有浮点数读进来，组织成 (batch, 3, N, d) 的四维张量：
    # 第 1 维区分 Q / K / V，后两维是每个位置的 d_k 维向量。
    data = np.array(tokens[4:4 + needed], dtype=np.float64)
    data = data.reshape(batch_size, 3, seq_len, d_k)

    # 缩放因子 1 / sqrt(d_k)，在算完点积后统一除。
    scale = math.sqrt(d_k)
    # 每个位置的窗口宽度固定是 2w+1，两端不足的部分在下面用掩码屏蔽。
    width = 2 * window_size + 1
    # 窗口内第 k 个槽位对应的相对下标：k - w，取值 -w .. w。
    offsets = np.arange(width) - window_size

    groups = []
    for b in range(batch_size):
        query, key, value = data[b]  # 每个都是 N x d

        # 把 K、V 两侧各补 w 行 0，这样位置 i 的窗口在内存里就是
        # 从第 i 行开始的连续 2w+1 行，可以直接用滑动窗口视图取出来。
        pad = np.zeros((window_size, d_k))
        key_pad = np.concatenate([pad, key, pad], axis=0)
        value_pad = np.concatenate([pad, value, pad], axis=0)

        # sliding_window_view 得到 (N, d, 2w+1) 的视图，不复制数据：
        # key_win[i, :, k] 就是窗口内第 k 个位置的 key 向量。
        key_win = np.lib.stride_tricks.sliding_window_view(key_pad, width, axis=0)
        value_win = np.lib.stride_tricks.sliding_window_view(value_pad, width, axis=0)

        # 所有位置、所有窗口槽位的分数一次算完：(N, 2w+1)。
        scores = np.einsum("id,idw->iw", query, key_win) / scale

        # 两端越界的槽位（i-w < 0 或 i+w > N-1）标记出来，
        # 窗口是被截断而不是补零参与 Softmax，所以这里置 -inf 再屏蔽。
        pos_idx = np.arange(seq_len)[:, None] + offsets[None, :]
        valid = (pos_idx >= 0) & (pos_idx < seq_len)
        scores = np.where(valid, scores, -np.inf)

        # 数值稳定的 Softmax：先减去窗口内最大值，再取 exp、归一化。
        # 被屏蔽的槽位是 -inf，exp 后正好是 0，权重自然为 0。
        row_max = scores.max(axis=1, keepdims=True)
        exp_scores = np.exp(scores - row_max)
        weights = exp_scores / exp_scores.sum(axis=1, keepdims=True)

        # 用窗口内的权重对 V 加权求和：(N, 2w+1) x (N, d, 2w+1) -> (N, d)。
        output = np.einsum("iw,idw->id", weights, value_win)

        lines = [
            " ".join(format_number(float(x)) for x in row)
            for row in output
        ]
        groups.append("\n".join(lines))

    # batch 之间用空行分隔，最后一组后没有空行
    print("\n\n".join(groups))


if __name__ == "__main__":
    solve()

---

示例推演

以样例 1 为例，$N=4$、$d_k=4$、$w=1$，缩放因子 $\sqrt{4}=2$。

位置 0

窗口是 $[0,1]$，只和 $K_0,K_1$ 计算分数：

$$\text{score}=[0.020000,\ 0.030000]$$

Softmax 后：

$$\text{weights}=[0.497500,\ 0.502500]$$

输出：

$$0.4975\times V_0+0.5025\times V_1=[0.50,\ 0.50,\ 0.00,\ 0.00]$$

位置 1

窗口是 $[0,2]$，分数为：

$$[0.030000,\ 0.040000,\ 0.025000]$$

Softmax 后：

$$[0.332772,\ 0.336116,\ 0.331112]$$

输出：

$$[0.33,\ 0.34,\ 0.33,\ 0.00]$$

位置 2

窗口是 $[1,3]$，分数为：

$$[0.050000,\ 0.035000,\ 0.055000]$$

Softmax 后：

$$[0.334434,\ 0.329455,\ 0.336111]$$

输出：

$$[0.00,\ 0.33,\ 0.33,\ 0.34]$$

位置 3

窗口是 $[2,3]$，分数为：

$$[0.045000,\ 0.060000]$$

Softmax 后：

$$[0.496250,\ 0.503750]$$

输出：

$$[0.00,\ 0.00,\ 0.50,\ 0.50]$$

四个位置的结果合起来就是样例输出。

对于样例 2，window_size=0 不在 $[1,seq\_len]$ 内，因此直接输出 0。

---

复杂度分析

解法	时间复杂度	空间复杂度	说明
完整注意力	$O(B\cdot N^2\cdot d_k)$	$O(B\cdot N^2)$	每个位置与所有 key 计算
滑动窗口注意力	$O(B\cdot N\cdot(2w+1)\cdot d_k)$	$O(B\cdot N\cdot d_k)$	本题采用；辅助空间为 $O(N\cdot d_k)$

其中 $B$ 是 batch_size。当 $w\ll N$ 时，滑动窗口版本比完整注意力快一个数量级。

---

易错点总结

1. 窗口边界要截断

$i-w$ 可能小于 $0$，$i+w$ 可能大于 $N-1$。必须使用：

# 两端要截断而不是补零：i 靠近开头时左边没有位置，靠近结尾时右边没有位置
left = max(0, i - window_size)
# 上界取 seq_len - 1，因为下标最大只到最后一个位置
right = min(seq_len - 1, i + window_size)

不能直接使用 i-w 到 i+w，否则会访问越界。

2. 窗口内做 Softmax，而不是全序列

每个位置的最大值和指数和都只来自当前窗口。混入窗口外的分数会改变权重，结果必然错误。

3. 不要提前舍入

题目说明中的分数和权重保留 6 位小数只是展示用，实际计算必须使用完整精度。提前四舍五入会累积误差，导致最后两位小数不一致。

4. 多 batch 的输出格式

batch 之间用一个空行分隔，最后一组后不能多输出空行。代码中使用 "\n\n".join(groups)，正好满足这个要求。

5. 避免 -0.00

浮点运算可能产生一个绝对值很小但符号为负的数，直接格式化会得到 -0.00。代码在格式化前把绝对值小于 $0.005$ 的数归零。

6. 非法输入要输出 0

window_size 必须满足 $1\le window\_size\le seq\_len$。样例 2 就是 window_size=0 的非法情况。除此之外，batch_size、seq_len、d_k 也必须为正整数。

---

扩展思考

• 窗口更大时：当 $w\ge N$ 时，滑动窗口退化为完整注意力，此时两种写法结果相同，只是滑动窗口版本仍然只扫描窗口。
• 因果滑动窗口：自回归模型通常使用因果掩码，位置 $i$ 只能看到 $[i-w,i]$，而不是左右对称窗口。只需要把 right 改成 $i$ 即可。
• 稀疏注意力：除了滑动窗口，还有空洞窗口、全局 token + 局部窗口等变体，目标都是在保留局部依赖的同时降低 $O(N^2)$ 的开销。
• 批量矩阵运算：如果使用 NumPy 等库，可以把一个 batch 的 Q、K、V 堆成三维张量，利用矩阵乘法一次性算出所有窗口分数；本题要求只用内置函数，所以按位置逐个计算。

---

相关题目

• 模域双向循环注意力网络前向计算
• 滑动窗口