在线注意力：免于存储权重矩阵的重缩放

摘要：归一化求和可以分块计算

计算一行注意力输出，无需保存全部权重。只需维护最大值、归一化项和累积向量，但更大评分到来时必须共同重缩放。本文推导恒等式，用会导致朴素指数溢出的 logits 验证，并展示具体错误：只重缩放分母，即使 softmax 看似数值稳定，输出仍然错误。贡献是附已执行代码的推导说明，并非新算法或 GPU 基准。

1. 要计算的对象

查询 q 与 N 个 d_k 维键 k_j 产生无量纲评分 s_j=qᵀk_j/√d_k，偏置可并入评分。每个键对应 d_v 维值向量 v_j。输出 o 是值向量的凸组合：系数为正且和为一。本文处理一行、无 dropout，并要求至少有一个未被掩码排除的键。被排除的键不贡献。问题是计算顺序，而非改变数学函数。

p_j = exp(s_j) / Σ_i exp(s_i) o = Σ_j p_j v_j m = max_j s_j l = Σ_j exp(s_j−m) u = Σ_j exp(s_j−m) v_j o = u/l

从全部评分减去同一最大值，等于将分子分母同乘 exp(−m)，比值不变。评分有限时，指数参数非正，各项不超过一，至少一项等于一。因此在实数算术下 1≤l≤N。这避免指数溢出，却不保证无限精度：极小项仍可能下溢，求和仍会舍入。

2. 不变量与尺度变换

处理 t 个键后，维护前缀最大值 m_t、相对此最大值的指数和 l_t，以及用同样权重累积的值向量 u_t。新对 (s,v) 到来，令 m′=max(m_t,s)。每个旧贡献 exp(s_i−m_t) 都要乘 exp(m_t−m′)，变成 exp(s_i−m′)。u_t 也必须乘同一因子，因为它使用相同权重累积。然后再加入新键贡献。

m′ = max(m, s) a = exp(m−m′), b = exp(s−m′) l′ = a·l + b u′ = a·u + b·v o′ = u′/l′

证明来自循环不变量：将 l 与 u 的定义代入更新式，正好得到扩展前缀上的和。单键初始化为 m=s_1、l=1、u=v_1，可避免 −∞−(−∞)。空状态需显式处理。必须区分未归一化分子 u 与输出 o；若把 o 当 u 重缩放，就遗漏了之前的 l 因子。

3. 让错误显现的例子

取 s=(1000,1001,999)，v_1=(1,0)、v_2=(0,2)、v_3=(3,−1)。Python 对 exp(1000) 报溢出，而稳定权重正比于 (e^−1,1,e^−2)。首键后 l=1、u=(1,0)。第二键把最大值从 1000 提高到 1001，旧分子必须变成 (e^−1,0)，而非仍为 (1,0)。于是 l=1+e^−1、u=(e^−1,2)。第三键不改变最大值，再加 e^−2·(3,−1)。

步骤	m	l	o₁	o₂
1	1000.0	1.000000000	1.000000000	0.000000000
2	1001.0	1.367879441	0.268941421	1.462117157
3	1001.0	1.503214724	0.514820191	1.240451338

最终 o=(0.514820191, 1.240451338)。第二步若只缩放 l 而不缩放 u，第一分量变成 0.935332675。第二分量恰巧仍正确，因为旧贡献为零；只检查一个分量可能漏掉缺陷。两个输出都有限。没有 NaN 和溢出，并不等于算法正确。

上图：每个键处理后的输出。下图：单个 N×N FP32 矩阵的理论大小，不含其他张量，也不是模型总内存。

4. 不保存权重的块合并

同一代数可合并两个互不重叠集合 A 与 B，各用 (m,l,u) 表示。取 m=max(m_A,m_B)，把双方归一化项与分子转换到新尺度，再相加。结果表示并集，因此在实数算术中满足结合律与交换律。浮点归约树变化可能改变末位比特。这支持并行分块，却不声称与顺序计算逐比特一致。

m = max(m_A, m_B) l = exp(m_A−m)l_A + exp(m_B−m)l_B u = exp(m_A−m)u_A + exp(m_B−m)u_B

附件在 257 个键、四维值向量上，对比块大小 1、7、32、257 与稳定完整计算。Python 随机种子为 20260924；评分均匀分布在 [−a,a]，a=1、10、1000，值在 [−2,2]。十二次比较中最大绝对误差约 6.66×10^−16。这只是 Python float 合成数据检查，不是普适误差界或 FP16/BF16 测量。代码、种子与结果均已保存。

5. 内存、复杂度及待测项目

每行的持久分子状态需要 d_v+2 个标量，此外还有查询、键值块和临时缓冲。d_v=64、FP32 时为 264 字节，但不是核函数总内存。完整 N×N FP32 评分矩阵每个头、每个批样本占 4N² 字节：N=4096 时 64 MiB，8192 时 256 MiB。朴素实现还可能存储归一化权重。避免这些矩阵不同于压缩 K、V，也不会消除解码器所需的 KV cache。

全部查询键对的稠密计算仍为二次复杂度 O(N²(d_k+d_v))。递推不会把精确注意力变成线性注意力。Python 标量循环可能比向量化矩阵乘法更慢；硬件收益需要适当分块、融合和访存。测量需声明 GPU、dtype、尺寸、掩码、batch、版本、预热和同步。本文未执行该基准，也不报告加速比。

全部被掩码的行需要明确约定：0/0 不定义分布。代码拒绝无键状态，并要求有限评分，掩码键须预先过滤。异号值可能在分子中抵消，低精度还需要选择累加器。梯度与 dropout 也需要额外逻辑；验证单行前向计算不等于认证完整训练实现。

6. 来源与结论

Milakov 与 Gimelshein 描述了在线归一化（2018，arXiv v2）；实验使用 Tesla V100、CUDA 9.1 和不同 batch，时间不能移植到本代码。Dao 等的 FlashAttention（2022，arXiv v2）将分块计算与减少 HBM 流量、反向重计算联系起来。我们阅读了算法、分析与实验。本文讲解其基础之一，不复现基准，也不综述后续全部发展。

可验证结论是分子分母共享的不变量。最大值不仅是防溢出技巧，也定义累积贡献的数值尺度。它改变时，必须转换全部旧贡献。正是这一步，使我们在实数算术下保持注意力结果，同时无需存储权重矩阵。

Milakov M., Gimelshein N. (2018), Online normalizer calculation for softmax, arXiv:1805.02867v2.

Dao T., Fu D. Y., Ermon S., Rudra A., Ré C. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, arXiv:2205.14135v2.

from math import exp
scores = [1000., 1001., 999.]
values = [[1., 0.], [0., 2.], [3., -1.]]
m, total, numerator = scores[0], 1., values[0][:]
for s, v in zip(scores[1:], values[1:]):
    new_m = max(m, s)
    old_scale, new_scale = exp(m-new_m), exp(s-new_m)
    total = total*old_scale + new_scale
    numerator = [u*old_scale + x*new_scale for u, x in zip(numerator, v)]
    m = new_m
print([u/total for u in numerator])

代码、数据与说明 · JSON. 教学计算使用 Python 3.14.0，图使用 Matplotlib 3.11.2。分析由 AI 辅助，不声称经过同行评审或人工审核。原创 ImageGen 封面仅作示意，不记录 EL-AI 人员、场所或实际安装。来源查阅于 2026 年 9 月 24 日。