Numerical-Stability

副标题 / 摘要 Self-Attention 的公式很短，但工程细节很长：从 Q/K/V 计算到 softmax 数值稳定、mask 与缩放，每一步都影响效果与性能。本文用 ACERS 结构给出推导、实践步骤与可运行示例。 预计阅读时长：12~16 分钟 标签：attention、transformer、softmax SEO 关键词：Self-Attention, Softmax, Scaled Dot-Product, 数值稳定 元描述：Self-Attention 的计算公式与 softmax 稳定实现方法，含工程实践与示例代码。 目标读者 想真正理解 Self-Attention 公式含义的学习者 需要处理训练不稳定/溢出的工程实践者 关注注意力数值稳定与实现细节的开发者 背景 / 动机 在 Transformer 中，Self-Attention 是计算量最大、数值最敏感的模块之一。 很多训练不稳定、输出 NaN 的问题，都来自 softmax 的溢出/下溢或 mask 的错误处理。 理解公式与稳定实现，可以显著减少工程“踩坑”。 核心概念 Q/K/V：查询、键和值，来自输入线性投影 缩放点积注意力：$\text{softmax}(QK^\top/\sqrt{d_k})V$ 数值稳定：通过减去行最大值避免 softmax 溢出 思路推导（从朴素到稳定实现） 朴素做法 先算所有相似度 $S = QK^\top$，再做 softmax 得到权重 $P$，最后 $O = PV$。 这个实现最直观，但当 $S$ 很大时会出现 exp 溢出。 关键观察 softmax 对每行同时加上或减去一个常数不改变输出： $\text{softmax}(x) = \text{softmax}(x - \max(x))$。 ...

副标题 / 摘要 注意力中的缩放项 \u221a(d_k) 不是装饰，而是数值稳定的关键：它控制 QK^T 的方差，避免 softmax 饱和和梯度消失。本文用公式与实验解释其必要性，并给出工程场景建议。 预计阅读时长：12~16 分钟 标签：attention、transformer、scaled-dot-product SEO 关键词：Attention, Scaled Dot-Product, \u221a(d_k) 元描述：解释注意力缩放项的数学动机与工程收益。 目标读者 想理解 Transformer 注意力细节的入门读者 需要排查训练不稳定问题的工程实践者 关注数值稳定性与性能优化的开发者 背景 / 动机 在点积注意力中，维度越大，QK^T 的数值越大，softmax 越容易饱和。 一旦饱和，梯度接近 0，训练会变慢甚至不稳定。 \u221a(d_k) 的缩放项就是为了解决这个问题。 核心概念 点积注意力：$QK^\top$ 衡量相似度。 缩放项 \u221a(d_k)：控制相似度的尺度。 softmax 饱和：输入过大导致概率趋近 0/1，梯度变小。 A — Algorithm（题目与算法） 用通俗语言说明主题内容 维度大时，QK^T 变大，softmax 过于“自信”。 缩放 \u221a(d_k) 后，数值回到合理范围，梯度更健康。 基础示例（1） d_k=64 时，如果不缩放，softmax 输出会非常尖锐。 基础示例（2） d_k=512 时，缩放与否会直接影响训练是否稳定。 实践指南 / 步骤 使用标准缩放：$QK^\top / \sqrt{d_k}$。 如果做自定义注意力，先验证 softmax 分布是否过尖锐。 在混合精度训练下，缩放更重要。 可运行示例（缩放与不缩放的对比） import torch import torch.nn.functional as F def attn_scores(d, scale=True): q = torch.randn(1, 1, d) k = torch.randn(1, 8, d) scores = q @ k.transpose(-2, -1) if scale: scores = scores / (d ** 0.5) probs = F.softmax(scores, dim=-1) return probs.max().item(), probs.min().item() for d in [32, 128, 512]: mx_s, mn_s = attn_scores(d, scale=True) mx_u, mn_u = attn_scores(d, scale=False) print(f"d={d} scaled max={mx_s:.3f} min={mn_s:.3f} | unscaled max={mx_u:.3f} min={mn_u:.3f}") 解释与原理 如果 $q_i, k_i \sim \mathcal{N}(0, 1)$， ...