首页
应用
关于
Search
1
Pytorch DDP
2,489 阅读
2
Pytorch 常见问题
1,520 阅读
3
视频时序切分
1,358 阅读
4
Semi-Supervised + Noisy Label
1,055 阅读
5
中文场景下的CLIP图文预训练
1,054 阅读
AIGC
AIGC Daily Papers
AIGC Fundamentals
其他
职场经验复盘
多模态理解
购房/投资
阅读
广告
Segmentation
LeetCode
Pytorch
Python
Shell
C++
常用链接
Search
标签搜索
AIGC
论文速读
人工智能
ai
视频生成
DiT
对齐
蒸馏
图像生成
扩散模型
attention
transformer
视频编辑
3D生成
音频生成
稀疏注意力
diffusion
基础知识
多模态
文生图
Jefxiong
累计撰写
224
篇文章
累计收到
8
条评论
首页
应用
栏目
AIGC
AIGC Daily Papers
AIGC Fundamentals
其他
职场经验复盘
多模态理解
购房/投资
阅读
广告
Segmentation
LeetCode
Pytorch
Python
Shell
C++
常用链接
页面
关于
搜索到
224
篇与
人工智能炼丹君
的结果
2026-09-25
AIGC 基本功|变分下界与重参数化-ELBO
AIGC 基本功|变分下界与重参数化-ELBO 所属方向:数学基础 | 难度:入门 | 前置知识:无知识树先修节点;需了解条件概率、概率密度、期望、矩阵乘法与链式求导 关键词:变分推断、ELBO、重参数化、KL 散度、后验坍缩、β-VAE、IWAE、得分函数估计量 01. 为什么需要它 先摆三组可复现的数字,全部出自文末四份完整脚本。本轮在 Windows CPU、Python 3.11.9、NumPy 2.4.6 实跑;固定 float64 与随机种子。前三份脚本做数学实验,第四份还依赖 matplotlib 3.10.7 生成三图。表中的“最优”来自第 4.3 节的线性高斯闭式,不是训练真实图像模型得到的性能。 第一组:KL 项掉到 0,潜变量整条死掉。 固定解码器噪声 σ²=0.25,让四个独立数据方向分别具有方差 λ=4.0、1.0、0.25、0.0625,只改 KL 项的权重 β,看看最优解长什么样: 数据方向方差 λ β=0.5 β=1.0 β=2.0 β=4.0 死亡阈值 β* 4.0 MSE 0.125 存活 MSE 0.250 存活 MSE 0.500 存活 MSE 1.000 存活 16.00 1.0 MSE 0.125 存活 MSE 0.250 存活 MSE 0.500 存活 MSE 1.000 死亡 4.00 0.25 MSE 0.125 存活 MSE 0.250 死亡 死亡 死亡 1.00 0.0625 死亡 死亡 死亡 死亡 0.25 死亡那一格的具体解是 a=0、s=1、KL=0:编码器把均值直接输出 0、方差输出 1,等于彻底放弃这个维度。此时重建误差精确等于 λ——也就是把 x 全猜成 0 的水平,这个维度一点信息都没传。问题在于看 loss 曲线你只会看到「KL 顺利降到 0,总损失还在降」,很难意识到这是失败而不是收敛。 第二组:把 ELBO 当成似然来监控,会被方差骗。 在一个 D=6、K=3 的高斯线性模型上,真实 log p(x) = -5.434331,解析 ELBO = -7.419767,两者差 1.985436。但用训练时真正用的那个单样本蒙特卡洛估计量去估 ELBO,重复 4000 次,标准差是 4.308——比它和 log p(x) 之间的差距还大一倍多,结果有 40.7% 的采样值直接越过了 -5.434 这条「上界」。如果把单样本曲线直接当成精确似然或模型排序,可能被波动误导;重复评估与报告不确定性才有助于比较。 第三组:不用重参数化,梯度方差能大到没法训练。 同一个模型上估 ∇ E_q[log p(x|z)],重参数化的梯度方差和是 377.55,得分函数估计量是 3043.11,差 8.1 倍。听起来还能忍?把数据维度 D 从 8 拉到 2048: D 重参数化方差 得分函数方差 比值 8 2.08 2 469 1 188× 64 1.98 104 328 52 630× 512 1.84 7 579 291 4 121 907× 2048 1.84 114 230 222 61 916 893× 重参数化的方差几乎不随 D 动,得分函数估计量按 D² 往上冲,到 D=2048(约 45×45 个标量)比值是 6191 万倍。这是本教学扫描中一个很大的方差差距,足以说明原始得分函数估计量可能难用;它不证明所有高维模型都无法用得分函数训练。这里固定 K=3、归一化 W 的列范数、s=0.6、每个 D 用 4000 次采样,并未训练 2048 维图像模型。2048 只是观测维度,约等于 45×45 个标量的数量。 这三组数字背后是同一件事:我们想最大化的是 log p(x),但它算不动,只能换成它的一个下界 ELBO 来优化;而换成下界之后,「换掉了什么」「这个下界的估计量有多吵」「梯度怎么穿过采样」这三笔账必须自己算清楚。不搞清楚 KL 项的价,你就不知道 β 该往哪调;不搞清楚下界的松紧,你就不知道扩散模型的训练目标从哪来。 02. 最小可用理解 三句话讲完核心: 算不动的边际似然换成一个能算的下界:$p_\theta(x)$ 要对所有 z 积分,算不动,于是引入一个我们自己挑的分布 $q_\phi(z \mid x)$,用 Jensen 不等式把优化目标换成 ELBO,它只需要在 q 下求期望,采样就能算。 下界和真值之间差的那一坨,正好是 q 和真实后验的 KL:固定生成模型时,对数边际密度等于 ELBO 加上后验 KL(完整公式见 3.3)。所以「下界有多松」等价于「q 离真后验有多远」,而不是「q 离先验有多远」——后者的 KL 在 ELBO 中被减去,是另一回事。 要让这个下界能被反向传播,就把随机性从梯度路径上挪走:写成 $z = m_\phi + s_\phi \odot \epsilon$,$\epsilon$ 与参数无关,固定一份噪声后,z 是参数的确定性函数,梯度沿它传回 $m_\phi$ 和 $s_\phi$;在可交换微分与期望等正则条件下,这给出有效的路径梯度估计。 03. 数学推导 先统一符号:$x\in\mathbb R^D$ 是一个观测,$z\in\mathbb R^K$ 是隐变量;$\theta$ 控制生成模型,$\phi$ 控制编码器;$m_\phi(x),s_\phi(x)\in\mathbb R^K$ 分别是均值和逐维标准差,$s_j>0$。$\odot$ 表示逐元素乘,$\ell=\log s$,工业代码的 logvar 则是 $2\ell$。$I_D,I_K$ 是相应单位矩阵,$\mathrm{tr}$ 是迹。本文对数为自然对数,KL 和对数密度的单位是 nats;连续密度可以大于 1,对数密度也可以为正。 3.1 边际似然为什么算不动 生成模型的设定很简单:先从一个固定先验里采隐变量 $z \sim p(z)$(通常是标准正态),再用解码器生成观测 $x \sim p_\theta(x \mid z)$。要拟合数据,目标是对数边际似然: $$\log p_\theta(x) = \log \int p_\theta(x \mid z) \, p(z) \, dz$$ 麻烦全在这个积分上。$p_\theta(x \mid z)$ 是由神经网络给出参数的条件概率分布,一般非线性解码器的边际积分没有可用闭式解;本文第 4 节的线性高斯模型是特意选出的可解特例。用数值积分的话,z 是 K 维,网格点数随 K 指数爆炸,K=32 就已经不可能。而 K 小了模型表达能力又不够——这正是我们不愿意接受的取舍。 3.2 换个能算的目标 既然积分算不动,就绕开它。引入在目标分布支撑上为正、满足相关可积条件的以 x 为条件的分布 $q_\phi(z \mid x)$(后面简称 $q$),把被积函数乘一个 $q/q$: $$\log p_\theta(x) = \log \int q_\phi(z \mid x) \, \frac{p_\theta(x, z)}{q_\phi(z \mid x)} \, dz = \log \mathbb{E}_{q_\phi}\!\left[\frac{p_\theta(x, z)}{q_\phi(z \mid x)}\right]$$ log 是凹函数,Jensen 不等式给出 $\log \mathbb{E}[Y] \ge \mathbb{E}[\log Y]$,于是 $$\log p_\theta(x) \ge \mathbb{E}_{q_\phi}\!\left[\log \frac{p_\theta(x, z)}{q_\phi(z \mid x)}\right] \equiv \mathcal{L}(\theta, \phi)$$ 右边就是 ELBO(Evidence Lower Bound)。它好算:$\log p_\theta(x, z) = \log p_\theta(x \mid z) + \log p(z)$ 两项都能直接求值,期望用采样估计就行。代价是我们不再直接优化 log p(x),而是优化它的一个下界。 3.3 差的那一项到底是什么 这是全文最关键的一步,也是最容易被跳过的一步。把 $\log p_\theta(x)$ 写成对 q 的期望(它和 z 无关,所以这么写是恒等的),再硬塞一个 $\log \frac{q}{q}$ 进去: $$\log p_\theta(x) = \mathbb{E}_{q}\!\left[\log p_\theta(x)\right] = \mathbb{E}_{q}\!\left[\log \frac{p_\theta(x, z)}{p_\theta(z \mid x)}\right] = \mathbb{E}_{q}\!\left[\log \frac{p_\theta(x, z)}{q_\phi(z \mid x)}\right] + \mathbb{E}_{q}\!\left[\log \frac{q_\phi(z \mid x)}{p_\theta(z \mid x)}\right]$$ 第一项正是 ELBO,第二项按定义就是 KL 散度。合起来: $$\log p_\theta(x) = \mathcal{L}(\theta, \phi) + \mathrm{KL}\big(q_\phi(z \mid x) \,\|\, p_\theta(z \mid x)\big)$$ KL 恒非负,所以 $\mathcal{L} \le \log p_\theta(x)$,和 Jensen 的结论一致——但这一版多给了一个信息:等号成立当且仅当 q 等于真实后验。这里固定 x 和生成模型 θ,缝隙才完全等于 q 到该模型真实后验的 KL;联合训练改变 θ 时,后验与 log pθ(x) 本身也会改变。“q 等于后验”指概率分布几乎处处相同,并要求式中的量有定义。 3.4 把 ELBO 拆成两项看 把 $\log p_\theta(x, z)$ 展开,ELBO 可以写成更有物理含义的形式: $$\mathcal{L}(\theta, \phi) = \mathbb{E}_{q_\phi}\!\left[\log p_\theta(x \mid z)\right] - \mathrm{KL}\big(q_\phi(z \mid x) \,\|\, p(z)\big)$$ 第一项是重建项:从 q 里采 z,解码器能不能还原出 x。第二项是先验对齐项:q 别离先验太远——因为生成时我们是从先验采 z 的,标准高斯先验具有全空间支撑,“覆盖不到”应理解为采到的概率很小,而非严格为零。这个惩罚逐个 x 约束 q;仅让汇总后验 q(z)=E数据[q(z|x)] 接近先验,不能保证每个条件后验都接近先验。 这里有个必须记住的区分,后面第 08 节还会回到它: $\mathrm{KL}(q \,\|\, p(z))$ —— 出现在 ELBO 里,是要被最小化的一项,物理含义是「编码分布别跑太偏」。 $\mathrm{KL}(q \,\|\, p(z \mid x))$ —— 不出现在 ELBO 里,是下界的缝隙,物理含义是「q 离真后验还差多少」。 在我们那个 D=6 的例子里,这两个数是 3.364364 和 1.985436,不是一回事,也不是简单的包含关系。β-VAE 显式改的是先验 KL 的权重,但重新训练后 q 与生成模型都会变化,所以真实后验差距也会变化,且不保证变小。 顺带一个很容易漏掉的推论:把 KL 乘上 β 之后得到的目标 $$\mathcal{L}_\beta = \mathbb{E}_{q_\phi}\!\left[\log p_\theta(x \mid z)\right] - \beta \cdot \mathrm{KL}\big(q_\phi(z \mid x) \,\|\, p(z)\big)$$ 当 $\beta\ge1$,仍有 $\mathcal L_\beta=\mathcal L-(\beta-1)\mathrm{KL}(q\Vert p)\le\mathcal L\le\log p_\theta(x)$,所以它仍是下界,只是更松;$0<\beta<1$ 时则不再保证下界。不同 β 的目标包含不同惩罚,不能直接当成同一种似然估计比较。 3.5 梯度怎么穿过采样 ELBO 对解码器参数 $\theta$ 的梯度没问题,重建项直接可导。麻烦在对 $\phi$ 的梯度:期望的分布本身依赖于 $\phi$,而采样操作不可导。设 $f(z) = \log p_\theta(x \mid z)$,要估的是 $\nabla_\phi \mathbb{E}_{q_\phi}[f(z)]$。 方法一:得分函数估计量(REINFORCE)。 直接把导数挪进期望: $$\nabla_\phi \mathbb{E}_{q_\phi}[f(z)] = \mathbb{E}_{q_\phi}\!\left[f(z) \, \nabla_\phi \log q_\phi(z \mid x)\right]$$ 在可交换微分与积分、分布支持集不随参数变化等常见正则条件下,这个式子成立;它不要求 f 对 z 可导,也可用于离散 z;方差可能较大,因为整个 f 都被乘进梯度里。此处 f 不显含 φ;若 f 也依赖 φ,另需加上 E_q[∇φ f]。下面两种估计量只比较重建项梯度,先验 KL 的解析梯度另算。对高斯先验,KL 对 m 的梯度是 m,对 ℓ 的梯度是 s²−1;最大化 ELBO 时从重建梯度里减去它们。 方法二:重参数化。 把 z 写成参数的确定性函数外加一个与参数无关的噪声: $$z = m_\phi(x) + s_\phi(x) \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$ 于是期望可以改写成对 $\epsilon$ 的期望,导数直接进去: $$\nabla_\phi \mathbb{E}_{q_\phi}[f(z)] = \mathbb{E}_{\epsilon}\!\left[\nabla_\phi f\big(m_\phi(x) + s_\phi(x) \odot \epsilon\big)\right] = \mathbb{E}_{\epsilon}\!\left[\nabla_z f(z) \cdot \nabla_\phi z\right]$$ 关键在于 $\nabla_z f$ 用到了 f 对 z 的局部形状,相当于「知道往哪个方向挪 z 会让 f 变大」,这是方法一完全没有利用的信息。 方差为什么差这么多——把与 z 无关的部分拆开。 高斯解码器 $p(x|z)=\mathcal N(Wz+b,\sigma^2 I_D)$ 中,$W\in\mathbb R^{D\times K}$、$b\in\mathbb R^D$、$\sigma^2>0$ 固定。令 $d=x-b$: $$f(z)=-\frac D2\log(2\pi\sigma^2)-\frac{\|d-Wz\|^2}{2\sigma^2}=c_0+\frac{d^\top Wz}{\sigma^2}-\frac{z^\top W^\top Wz}{2\sigma^2},\quad c_0=-\frac D2\log(2\pi\sigma^2)-\frac{d^\top d}{2\sigma^2}$$ 路径梯度对 $c_0$ 求导为零;得分函数把它乘上 $S(z)=\nabla_\phi\log q(z|x)$,虽然 E_q[S]=0,不改变梯度期望,却可能造成很大的方差。归一化常数和固定观测残差都在 c₀ 中,仅凭 D 扫描不能把全部差距归因于前者。本实验同时把 W 列范数固定,避免梯度尺度随 D 自然增长,故 D² 趋势是这组设置的观测,不是普适定律。 用与当前采样 z 无关的 baseline c,E_q[(f-c)S]=E_q[fS]。脚本用解析 E_q[f] 作预言机 baseline,D=2048 的方差比从 61 916 893× 降至约 5.3×;它不是通用最优 baseline。若按梯度各分量的总方差选一个标量,最优常数为 E[f‖S‖²]/E[‖S‖²],通常并非 E[f]。实际可用历史批次或独立样本估计 c,并在这个梯度估计里停止对 c 求导;直接减去同批样本均值会引入相关性,在独立同分布 n 个样本的简单情形把期望缩成 (1−1/n) 倍。 一个反例说明边界。 固定 q=N(m,1),在 m=0 估计 f(z)=sin(10z) 的均值梯度,pathwise 是 10cos(10z),score 是 z sin(10z)。200000 个共同噪声样本实跑,方差分别 50.062 与 0.499,后者更低。两者理论均值均为 10exp(−50)。因此重参数化在本文高斯重建中明显有利,并不保证任何可微 f 上都方差更小。反例与推导核验见第 09 节。 3.6 下界能有多紧:IWAE 单样本 ELBO 的缝是 $\mathrm{KL}(q \,\|\, p(z \mid x))$。要缩缝有两条路:把 q 变强(换更灵活的后验族),或者换一个更紧的界。后者的经典做法是 IWAE 的 k 样本界: $$L_k = \mathbb{E}_{z_1 \dots z_k \sim q_\phi}\!\left[\log \frac{1}{k} \sum_{l=1}^{k} \frac{p_\theta(x, z_l)}{q_\phi(z_l \mid x)}\right]$$ 注意顺序:先对 k 个重要性权重取平均,再取 log。单样本 ELBO 是「log 再取平均」(每个样本的 log 比值取期望),IWAE 是「平均再 log」,Jensen 保证后者更紧。可以证明 $\log p_\theta(x) \ge L_{k+1} \ge L_k \ge L_1 = \mathcal{L}$,收敛还需要更强的尾部与可积条件,具体充分条件见下文。 这里 $\widehat L_k=\log(k^{-1}\sum_l w_l)$ 是随机结果,而 $L_k=\mathbb E[\widehat L_k]$ 才是下界,$w_l=p(x,z_l)/q(z_l|x)>0$。支持覆盖使 E_q[w]=p(x),Jensen 给出 Lk≤log p(x)。单调性也能一步看清:k+1 个权重的均值等于所有“删去一个权重”的 k 样本均值的平均;先应用 log 的凹性,再取期望,就得到 L{k+1}≥L_k。这个论证针对同一固定模型与 q 下的独立同分布采样,并要求相应期望存在。 IWAE 原文定理 1 给出重要性权重有界时趋近真值的充分条件。一般讨论极限还须保证期望可随极限交换,例如相关 log 估计量的均匀可积性;只有支持覆盖和 E[w] 有限不足以直接交换 log 与期望。即便界更紧,单次结果以及有限重复的均值仍可越过 log p(x),也不保证逐次单调。 明确反例:离散真实后验为 (0.5,0.5),q=(0.9,0.1),设 p(x)=1,则权重为 (5/9,5)。一条先采到 5、再采到两次 5/9 的路径,其 log 平均权重依次为 1.6094、1.0217、0.7115,全高于真实 log p(x)=0 且连续下降;对全部路径取期望得到的 L₁、L₂、L₃ 则是 −0.3681、−0.2761、−0.2200,保持下界与单调性。多采样平均 log 权重只降低同一 ELBO 估计的方差;改成 log 平均权重才改变期望目标及其偏差。 代价是:梯度不再是对单个样本的简单求和,k 个权重互相耦合(每个权重的梯度都带上了其他权重的归一化因子),总计算量随 k 线性上涨;编码器梯度的信噪比还可能随 k 增大而下降(见 Rainforth 等,2018),这反而不利于训练编码器——工程上这是个明确的取舍,不是免费的午餐。 这张图要看什么:左图四根柱子——ELBO(-7.420)加 KL(q‖p(z|x))(1.985)正好等于 log p(x)(-5.434),而 KL(q‖p(z))(3.364)是另一根完全不同的柱子,别混。中图是单样本蒙特卡洛 ELBO 的 4000 次采样分布,标准差 4.31,虚线是真实上界 -5.434,40.7% 的样本落在它右边。右图是 IWAE 的 k 扫描,k 从 1 涨到 500,6000 次重复的平均缝隙从 2.049 缩到 0.00262,标准误分别为 0.04643 和 0.00080;这是有限重复的估计,不能当成解析界。k=1、5、50、500 时单次越过 log p(x) 的比例分别为 34.7%、43.5%、47.7%、48.2%。标准误描述均值的不确定性,不是严格置信界。 04. 代码实现 4.1 恒等式核对:先造一个能算到机器精度的模型 要验证「差的那一项到底是什么」,必须三样东西都有解析解:边际似然、真实后验、ELBO。高斯线性模型满足这一点。将联合 log 密度中的 z 项配方: $$\log p(x,z)=\text{const}-\frac12z^\top Pz+h^\top z=\text{const}-\frac12(z-\mu)^\top P(z-\mu),\quad P=I_K+\frac{W^\top W}{\sigma^2},\quad h=\frac{W^\top(x-b)}{\sigma^2},\quad\mu=P^{-1}h$$ 所以后验 p(z|x)=N(μ,P⁻¹),边际 p(x)=N(b,WWᵀ+σ²I_D)。P 正定,实际计算可用 solve/Cholesky;这里 K=3,为便于展示保留显式逆。我们的模型为:$z \sim \mathcal{N}(0, I_K)$,$x \mid z \sim \mathcal{N}(Wz + b, \sigma^2 I_D)$。边际仍是高斯,后验是共轭高斯,ELBO 也能闭式算。 import numpy as np D, K, SIGMA2 = 6, 3, 0.35 ** 2 rng_model = np.random.default_rng(20260925) W = rng_model.normal(0, 1, size=(D, K)) * 0.9 b = rng_model.normal(0, 0.3, size=(D,)) def true_posterior(x): """p(z|x) 的闭式解:precision = I + W^T W / sigma^2。""" prec = np.eye(K) + W.T @ W / SIGMA2 cov = np.linalg.inv(prec) return cov @ (W.T @ (x - b) / SIGMA2), cov def elbo_analytic(x, m, s): """ELBO 的解析值。注意 ||x - Wz - b||^2 对 q 求期望会多出一项 trace。""" var = s ** 2 resid = x - (W @ m + b) recon = -0.5 * (D * np.log(2 * np.pi * SIGMA2) + (resid @ resid + np.trace(W.T @ W @ np.diag(var))) / SIGMA2) kl = 0.5 * np.sum(var + m ** 2 - 1.0 - np.log(var)) return float(recon - kl), float(recon), float(kl) def kl_gauss_gauss(m, s, mean2, cov2): """KL(N(m, diag(s^2)) || N(mean2, cov2)),两个高斯的闭式。""" var, diff = s ** 2, mean2 - m cov2_inv = np.linalg.inv(cov2) _, logdet2 = np.linalg.slogdet(cov2) return 0.5 * (np.trace(cov2_inv @ np.diag(var)) + diff @ cov2_inv @ diff - K + logdet2 - np.sum(np.log(var))) 重建项里那个 np.trace(W.T @ W @ np.diag(var)) 是最容易漏掉的一项:对 $z$ 求期望时,$\|x - Wz - b\|^2$ 里的 $Wz$ 项会因为 z 的随机性多出一份方差贡献。漏了它,ELBO 就不是下界了。 关键一步是故意把 q 指定错:取真实后验的均值再加扰动,协方差只留对角线,再乘 0.8 并加 0.05。在固定状态下,真实边际方差为 [0.21735,0.04503,0.04927],q 方差为 [0.22388,0.08602,0.08942],三个都更大,不能称为“更自信”。这份 q 同时有均值偏移、方差失配和缺失相关性。这样 KL(q‖p(z|x)) 严格大于 0,缝才看得见。跑起来: [1] 解析三项 E_q[log p(x|z)] = -4.055403 KL(q || p(z)) = 3.364364 ELBO = -7.419767 KL(q || p(z|x)) = 1.985436 ELBO + KL(q||p(z|x)) = -5.434331 log p(x) = -5.434331 |误差| = 1.776e-15 恒等式对到 2e-15 以内。同时注意:KL(q‖p(z))=3.364 比 KL(q‖p(z|x))=1.985 还大——这两个量的大小关系没有必然规律,别用其中一个去猜另一个。 均值正确时,“把协方差的非对角项丢掉”仍不是 reverse KL 意义下的最优对角近似。设 q=N(m,diag(v)),真实后验精度为 P,KL 中依赖 v_j 的项为 (P_jj v_j−log v_j)/2,对它求导得到 v_j=1/P_jj,均值最优为 m=μ。于是: $$\min_{m,v}\mathrm{KL}(q\Vert p(z|x))=\frac12\log\frac{\prod_j P_{jj}}{\det P}$$ 本例直接使用真实边际方差 diag(P⁻¹) 时,残留缝隙为 0.79398;使用 1/diag(P) 时为 0.40038。这才是在固定生成模型下该对角族的近似误差;原来 1.98544 还包含人为指定 q 的额外误差。满协方差 q 等于后验时,缝才为零。 4.2 两种梯度估计量,方差实测 同一个模型上,$\nabla_m$ 和 $\nabla_\ell$($\ell = \log s$)的解析梯度都能写出来,所以「谁对谁错」有标准答案,剩下的差别纯粹是方差: def sample_reparam(x, m, s, W, b, sigma2, rng, n): """重参数化:z = m + s ⊙ ε,梯度经 z 反传。""" eps = rng.normal(size=(n, m.size)) z = m + s * eps score_z = (x - (z @ W.T + b)) @ W / sigma2 # ∇_z log p(x|z) return score_z, score_z * s * eps # ∂z/∂m = 1, ∂z/∂ℓ = s ⊙ ε def sample_score(x, m, s, W, b, D, sigma2, rng, n, baseline=None): """得分函数:∇ E[f] = E[f · ∇ log q],f = log p(x|z)。""" eps = rng.normal(size=(n, m.size)) z = m + s * eps r = x - (z @ W.T + b) f = -0.5 * (D * np.log(2 * np.pi * sigma2) + np.einsum("nd,nd->n", r, r) / sigma2) if baseline is not None: f = f - baseline return f[:, None] * (eps / s), f[:, None] * (eps ** 2 - 1.0) 两行代码的差别就是全部:重参数化用的是 $\nabla_z f$,得分函数用的是 $f$ 本身。跑 40000 次独立采样: 真实梯度(解析) ∇_m = [ 0.4228 4.9273 -0.6448] ∇_ℓ = [-1.9319 -2.4426 -3.8296] 估计量对比(40000 次独立采样,方差 = 6 个参数分量方差之和) 重参数化 方差和= 377.55 均值相对误差= 0.41% 5%相对RMS需 n≈3068 得分函数(无 baseline) 方差和= 3043.11 均值相对误差= 3.95% 5%相对RMS需 n≈24722 得分函数(预言机 baseline) 方差和= 1586.45 均值相对误差= 3.09% 5%相对RMS需 n≈12889 方差比:得分函数 / 重参数化 = 8.1× 加 baseline 之后 = 4.2× 在上述正则条件与独立 baseline 下,三个估计量理论上无偏;有限样本均值仍有误差。表中的样本数是按相对均方根误差计算的估计:设单样本梯度协方差为 V、真实梯度为 g,则 n 个独立样本均值满足 E‖ĝ−g‖²=tr(V)/n,所以 n≈tr(V)/(0.05²‖g‖²)。3068 与 24722 不是“95% 概率落在 5% 内”的保证。方差和是 6 个参数分量方差之和,正文表取总体式 ddof=0 的样本估计。 这张图要看什么:左图是 $\nabla_m$ 第 0 个分量(真值 0.4228)在三种估计量下的分布——均值都压在真值附近,标准差分别是 5.7、16.9、12.0,尾巴长度差一个量级。右图的双对数坐标里,重参数化那条线几乎是平的(方差 2.08 → 1.84),得分函数曲线在本设置中近似呈 D² 增长;而把蓝色点(加了 baseline)和橙色点对比,比值从 1188× 一路到 6191 万×,加完 baseline 却稳定在 5~7×——说明减去完整均值能显著降低本组 score 方差;归一化常数与固定残差都参与,不能由此隔离各自的贡献。 4.3 β 加权:把维度死掉的过程解出来 这一节的目标是要一个能解到全局最优的玩具模型,这样「维度死了」是解析结论而不是训练运气。构造:每个坐标独立,$x_j \sim \mathcal{N}(0, \lambda_j)$,$q(z_j \mid x_j) = \mathcal{N}(a_j x_j, s_j^2)$,$p(x_j \mid z_j) = \mathcal{N}(b_j z_j, \sigma^2)$。每个维度就是一份独立副本,逐个维度单独求最优即可: SIGMA2 = 0.25 # 解码器观测噪声方差 def kl_of(a, lam, s): return 0.5 * (a ** 2 * lam + s ** 2 - 1.0 - np.log(s ** 2)) def recon_mse(a, b, lam, s): """E_{x,z}[(x - b z)^2],重建误差的期望。""" return lam * (1.0 - b * a) ** 2 + b ** 2 * s ** 2 def objective(p, lam, beta): a, b, ell = p s = np.exp(ell) return -recon_mse(a, b, lam, s) / (2 * SIGMA2) - beta * kl_of(a, lam, s) 目标对 $(a, b, \ell)$ 的梯度全部解析,脚本里再用中心差分核对一遍(实测最大偏差 5.5e-10),避免推错。数值梯度上升保留多起点作对照:a=b=0 是一个驻点,单起点可能停住;它不构成全局最优证明。下面直接把三变量问题降成一变量,给出全局解。 先固定 a,s 对 b 求导,得到 b*=aλ/(a²λ+s²)。令 r=a²λ+s²>0、t=s²/r∈(0,1],代回重建 MSE=λt,数据平均 KL=(r−1−log r−log t)/2。r−1−log r≥0,在 r=1 取到最小;其余目标为: $$J(t)=-\frac{\lambda t}{2\sigma^2}+\frac\beta2\log t,\quad J'(t)=-\frac\lambda{2\sigma^2}+\frac\beta{2t},\quad J''(t)=-\frac\beta{2t^2}<0$$ 因此 β>0、λ>0 时唯一最优 t=min(βσ²/λ,1)。选正号的一组解为 s²=t、a=√((1−t)/λ)、b=√(λ(1−t));a,b 同时取负也是同一重建解。t=1 时 a=b=0、s=1,MSE=λ;否则 MSE=βσ²、平均 KL=−log(t)/2。阈值 β*=λ/σ² 与连续的 MSE=min(βσ²,λ) 就这样推出来,等号处也关闭。λ 是数据方差,不是以 nats 衡量的信息量。 本轮扫描和配图采用这份闭式。独立多起点优化在 (λ,β)=(4,0.25) 仍有目标差 1.51×10⁻⁵,MSE 0.0627248 对闭式 0.0625,说明有限迭代会留下误差;原文 5.53×10⁻³ 的偏差也不应当作精确最优证明。闭式内部一致性误差 8.88×10⁻¹⁶;另以原三变量目标的随机扰动与一变量网格核对,均未发现更优值。 跑完 β 扫描,把最优解和闭式预测放在一起对: λ β s²闭式 βσ²/λ a²λ闭式 1-s² MSE闭式 βσ² 1.0000 1.00 0.25000 0.25000 0.75000 0.75000 0.25000 0.25000 1.0000 2.00 0.50000 0.50000 0.50000 0.50000 0.50000 0.50000 1.0000 3.00 0.75000 0.75000 0.25000 0.25000 0.75000 0.75000 0.2500 0.50 0.50000 0.50000 0.50000 0.50000 0.12500 0.12500 最大内部一致性偏差 = 8.88e-16;独立数值优化误差另列 λ=4.0000 理论死亡阈值 β* = λ/σ² = 16.00 λ=1.0000 理论死亡阈值 β* = λ/σ² = 4.00 λ=0.2500 理论死亡阈值 β* = λ/σ² = 1.00 λ=0.0625 理论死亡阈值 β* = λ/σ² = 0.25 存活时最优解落在闭式 $s^2 = \beta\sigma^2/\lambda$、$a^2\lambda = 1 - s^2$、$\mathrm{MSE} = \beta\sigma^2$ 上。这个闭式有个反直觉的推论:存活时的重建误差只由 β 和解码器噪声决定,跟这个维度携带多少信息 λ 无关——λ 只决定「这个维度值不值得用」。 05. 工业级实现对照 最小实现里 KL 是两个高斯的闭式,官方示例也是闭式,但分布、归约和 dtype 必须对齐。本节固定 diffusers v0.35.1 与 pytorch/examples 提交 acc295dc7b90,以 2026-10-12 的核对为准;主分支后续可变化。已从原始源码 AST 原样提取相关方法,在 PyTorch 2.10.0+cpu 上真实执行;没有下载权重、MNIST 或运行整个生成管线。 1. 官方最小实现:pytorch/examples。 pytorch/examples 的 vae/main.py 里,loss_function 就是标准的 BCE 重建加闭式 KL,没有任何 β(以本轮核对的固定源码为准): def loss_function(recon_x, x, mu, logvar): BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum') # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2) KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) return BCE + KLD 和本文 3.4 节那一项完全对得上:$\mathrm{KL} = \frac{1}{2}\sum(\mu^2 + s^2 - 1 - \log s^2)$,代码里写成 -0.5 * sum(1 + logvar - mu^2 - exp(logvar)),logvar 就是 $\log s^2$。要加 β 得自己动手,官方示例没有。 2. 为什么参数化 logvar 而不是 s。 方差必须为正,直接学 s 要加约束;学 $\log s^2$ 则值域是全体实数,网络怎么输出都合法。diffusers 的 DiagonalGaussianDistribution 还额外做了截断(以本轮核对的固定源码为准): self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) self.logvar = torch.clamp(self.logvar, -30.0, 20.0) self.std = torch.exp(0.5 * self.logvar) self.var = torch.exp(self.logvar) clamp 到 [-30, 20] 限制指数输入,但不等于所有精度安全:$\exp(20)$ 约为 4.85 亿,fp32 可表示但 fp16 已无法表示;下界 -30 避免过小方差与过大负 logvar。还需核对计算 dtype,截断本身不是所有精度下的安全保证。 3. 采样就是重参数化,一行代码。 同一个类里: def sample(self, generator=None): sample = randn_tensor(self.mean.shape, generator=generator, device=self.parameters.device, dtype=self.parameters.dtype) x = self.mean + self.std * sample return x 就是本文 3.5 的 $z = m + s \odot \epsilon$。mode() 直接返回均值——一些确定性编码场景使用 mode(),但 diffusers img2img 默认使用后验 sample(generator=...),训练和推理均可采样。 4. KL 的求和维度与归一化。 diffusers 里 kl() 写成(以本轮核对的固定源码为准): return 0.5 * torch.sum(torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=[1, 2, 3]) 对通道、高、宽三个维度求和,得到的是每个样本一个标量,得到 [B] 后,训练侧可按其目标再取 batch 平均。这里源码固定求和 dim=[1,2,3],该结论只适用于 [B,C,H,W];若直接输入五维 [B,C,T,H,W],返回 [B,W],本轮实跑形状为 [2,4],不能据类名假定它已归约全部视频维度。应核对具体视频 VAE 类与外部 loss,不能直接复用此四维示例。这个区别很实际:若重建归约保持不变,却仅把每样本的 KL 再除以 $C \times H \times W$,其相对权重就会相应变小;若两个 loss 一同同比缩放,其相对权重则不变。单独把 KL 改成逐元素平均时,潜变量维度越多,其相对权重越小,可能使先验对齐变弱、尺度漂移;过强 KL 才更直接推动后验坍缩。 5. 训练时 KL 根本不在模型里。 值得注意的是 AutoencoderKL 的模型文件里没有 kl_loss 这个东西,它只通过 encode() 返回 DiagonalGaussianDistribution 后验,KL 由外部训练脚本算。这是工程上的职责划分:模型只负责给出分布,损失函数怎么组合(KL 权重、perceptual loss、对抗损失)交给训练配置。 6. 潜空间的缩放。 SD1.x 常见 scaling_factor=0.18215 是乘到原始 latent 上的系数,按缩放后目标标准差为 1 的约定,对应原始标准差约 $1/0.18215\approx5.49$,缩放后才接近 1,不能把 0.18 当成原始潜空间尺度。它来自特定 VAE 与扩散训练的尺度约定,不由 ELBO 唯一决定。 7. 最小实现和工业实现的差异。 工业代码增加 logvar 截断、dtype 控制、分布采样接口及 latent 尺度约定。KL 逐维求和来自概率定义,batch 平均来自训练目标;它们不是任意工程细节。sample() 与 mode() 都有推理用途,要依具体管线选择。 06. 代价与边界 省了什么。 把一个对 K 维积分的不可解问题,变成了「采样 + 两个可导项」,梯度能用标准反向传播算。这是 VAE、扩散模型、以及一大票潜变量模型能训练起来的全部前提。 赔了什么,四条。 一是下界不是似然。 优化 ELBO 不等于优化 log p(x),中间隔着 KL(q‖p(z|x))。在固定生成模型下,这个缝受 q 族表达能力与实际优化结果共同影响:均值场对角高斯拟合不了有相关性的真实后验,缝就永远在。本例里缝是 1.985 nats,在一个 log p(x) 只有 -5.43 的小模型上,数值上差约 36%,但连续对数密度受坐标单位影响,不宜把这个比例当作通用误差尺度。 二是 KL 的方向有偏好。 $\mathrm{KL}(q\Vert p(z\mid x))$ 是 reverse KL:当 $q$ 把概率放到真实后验很低的地方,代价很大,因此在受限的单峰近似族下常表现为 mode-seeking,可能只覆盖一个峰;反向的 $\mathrm{KL}(p\Vert q)$ 则倾向覆盖目标的质量,即 mass-covering。VAE 的模糊不能单归因于这个方向,逐像素重建目标、解码器分布与压缩瓶颈都有关。 三是推断被参数化摊平(amortized)之后又多了一层误差。 经典变分推断对每个 x 单独优化一组变分参数;VAE 用一个共享的编码器网络去输出所有 x 的 $m_\phi(x)$ 和 $s_\phi(x)$。这一步是为了快——测试时一次前向就得到近似后验参数,不用重新迭代——但共享映射还增加了约束。即便每个 x 单独看,最优的对角高斯后验就在那儿,共享编码器也可能一辈子到不了。Cremer 等的分析 区分近似差距与摊销差距。若再把实际优化单独列账,固定 θ 与 x,可写成 log p(x)−L实际=(log p(x)−L族最优)+(L族最优−L共享最优)+(L共享最优−L实际),分别对应近似、共享约束、优化差距。在 L实际属于该共享模型族等前提下各项非负。当前实验直接指定一份 x 的 q,没有训练共享编码器,不能测出摊销差距。这也解释了为什么给编码器加容量在有些任务上能明显提 ELBO——在固定生成模型时,这改善的是近似推断。 四是加权 KL 会直接杀死维度。 这是第 4.3 节那张表的完整结论: λ=4.0000 在 β≤12.0 内都存活 λ=1.0000 从 β=4.0 起死亡(β=1 时 KL=0.6931) λ=0.2500 从 β=1.0 起死亡 λ=0.0625 从 β=0.25 起死亡 死亡阈值是 $\beta^\ast = \lambda / \sigma^2$。读出这个式子的含义:一个维度要活下来,数据方差 λ 必须盖过「KL 的价」β 乘上「解码器噪声」σ²。β 翻倍,能活下来的维度门槛就翻倍;解码器越准(σ² 越小),越多的维度能活。这仅解释线性高斯模型的阈值。强大的自回归解码器可通过其他条件预测数据而忽略 z,不能简单等同于把固定高斯观测方差 σ² 调小。 这张图要看什么:左图四条曲线是不同 λ 下重建 MSE 随 β 的变化,实心点表示维度存活、空心方块表示已死——注意存活段的 MSE 就是 $\beta\sigma^2$ 这条直线,跟 λ 无关;而死掉的段 MSE 平在 λ 上不再变化。右图是 $(\lambda, \beta)$ 平面上的相图,斜线是 $\beta = \lambda/\sigma^2$,横轴是 λ/σ²、纵轴是 β,线的上方关闭、下方存活。调 β 之前先看一眼这个平面:你真正要判断的是「在当前 β 下,哪些数据方向位于存活区域」。 缓解手段:free bits。 归约顺序对应 IAF 原文附录 C.8、公式 15 的分组平均后截断;本文另加 β 作教学对照。 在本脚本中,先对数据取每维平均 KL,再换成 $\max(\mathbb E_{\mathrm{data}}[\mathrm{KL}],C)$,是在 KL 小于 C 时去掉进一步压低它的梯度,给重建项使用这部分容量的机会;这不是保证至少传 C nats 的硬约束。下面的线性实验确实在边界附近找到更好的重建,但一般神经网络不保证被救活。这个“先平均再截断”的目标与“每个样本先截断再平均”通常不同,数值和全局解也不同。本模型 r=1 时记 R=−log(t)/2≥0,MSE=λexp(−2R),目标为 −λexp(−2R)/(2σ²)−βmax(R,C),所以 R*=max(C,0.5log(λ/(βσ²)),0)。这个特例确实用到 C 附近,但普遍神经网络无此保证。闭式结果(β=4): λ C=0(纯 β) C=0.05 C=0.2 1.0000 a=0.000 KL=0.000 MSE=1.000 a=0.308 KL=0.050 MSE=0.905 a=0.574 KL=0.200 MSE=0.670 0.2500 a=0.000 KL=0.000 MSE=0.250 a=0.617 KL=0.050 MSE=0.226 a=1.148 KL=0.200 MSE=0.168 λ=1.0 的维度被救回来了(MSE 1.000 → 0.670),λ=0.25 的也从 0.250 降到 0.168。代价是放松了对先验 KL 的惩罚;是否改善或损害解耦需要独立评估——这是一笔明码标价的交易。 什么时候不该用。 归一化流和自回归模型在相应建模假设下可计算精确似然,不一定需要变分下界;VQ 模型可精确计算离散 token 序列的自回归概率,但一般仍不能精确边缘化得到像素似然。ELBO 也不是训练高维生成模型的唯一途径,score matching、流匹配和对抗训练是其他路线。同一数据、同一似然约定下标准 ELBO 可比较为下界,但 q 的质量影响松紧,不能直接据此断言真实似然或感知质量的排序。 07. 经典论文脉络 Kingma & Welling, 1312.6114(2013)——VAE 原文。贡献是把「变分推断 + 重参数化 + 神经网络编码器」拼成一个能用 SGD 训的东西,并给出 SGVB 估计量。它同时确立了沿用至今的 loss 形式:重建项减 KL。本文 3.2~3.5 节基本是这篇的复述。 Rezende, Mohamed & Wierstra, 1401.4082(2014)——几乎同期独立提出的重参数化,论文里叫 stochastic backpropagation。贡献是把这个方法从「VAE 的一个技巧」推广到论文所讨论的可微潜变量模型与分布族,仍需可用的噪声变换和正则条件,并系统讨论了高斯之外的分布族怎么处理。想理解重参数化的适用边界(哪些分布能做、哪些只能退回得分函数),这篇比 VAE 原文讲得更清楚。 Burda, Grosse & Salakhutdinov, 1509.00519(2015)——IWAE。贡献就是本文 3.6 那个 $L_k$:把「log 再平均」换成「平均再 log」,得到一个随 k 单调变紧的界,并证明了收敛性。它澄清了一个当时普遍的误解——「多采几个样本只是降方差」,实际上换的是界本身。 Higgins et al., ICLR 2017——β-VAE。给 KL 项加权重 β > 1,旨在增强因子化与解耦表征,效果须独立评测;目标与重建取舍可对照同组作者的 2018 原论文。本文第 4.3 节给出线性高斯教学特例的代价面:只在存活区域,MSE 随 β 成正比;到阈值后饱和。它不是该论文非线性模型的通用定理,也未在本实验测量解耦指标。 Bowman et al., 1511.06349(2015)——讨论句子 VAE 后验坍缩的代表性工作。用 VAE 做句子生成,强大的自回归解码器会直接忽略潜变量,KL 掉到 0。这篇提出的 KL annealing(训练初期把 β 从 0 慢慢涨到 1)是可尝试的缓解手段之一,效果依赖具体模型和优化过程,并非保证解决坍缩。 补充两篇:Kingma et al. 的 IAF(1606.04934)走的是另一条路——不改界,改 q,用可逆变换把后验族变强来直接缩小那个缝;Kingma & Welling 的综述(1906.02691)适合把上述脉络串起来通读。 08. 常见误解 误解一:「ELBO 里的 KL 就是 q 和真实后验的差距」。 ELBO 里减去的是 $\mathrm{KL}(q\Vert p(z))$;下界的差距是 $\mathrm{KL}(q\Vert p(z\mid x))$。调整 β 会通过训练改变 q 和生成模型,后者也可能变化,只是没有保证随前者一起下降。 误解二:「ELBO 是下界,所以蒙特卡洛估计值不会超过 log p(x)」。 下界性质是对期望成立的,不是对每个样本成立。本例单样本估计的标准差是 4.308,而缝只有 1.985,结果 40.7% 的样本越过了真实上界。若有限样本估计越过 log p(x),采样波动是一种解释;仍应核对归约、似然常数与实现,不能仅凭越界认定程序正确。 误解三:「重参数化是为了让采样可微」。 准确说法是把参数相关的随机变量表示为固定噪声上的可导变换,从而得到路径梯度。采样「不可微」这个表述本身就有问题:得分函数估计量里 $\nabla_\phi \log q_\phi$ 是对参数求导,完全可微,它对 f 连可导性都不要求。路径估计用上了 ∇z f,在本文高斯重建实验中方差较低;第 3.5 节的正弦反例说明排序可反转。其首要机制是把参数相关采样表示成固定噪声的可导映射,方差降低是有条件的好处。 误解四:「KL 掉到 0 说明 KL 项优化到位了」。 这是后验坍缩的典型症状。本例 λ=1.0、β=4 时最优解就是 a=0、s=1、KL=0、MSE=1.000(正好等于 λ),编码器彻底放弃了这个维度。看数据平均 KL 为零可推出几乎所有 x 的 q(z|x)=p(z),说明潜变量不携带 x 信息;非零 KL 也未必等于有用信息。辅助统计量 A=Var数据(E_q[z|x])=a²λ 只检查均值活动度。active_unit() 使用 IWAE 第 5.2 节 的经验阈值 A>0.01;它不等价于理论 a≠0,方差编码的信息也可能被它漏掉。临近阈值时仍存活但 A≤0.01 的方向会被统计判为 inactive。 误解五:「β 越大解耦越好,重建变差只是小代价」。 在线性模型中,最优 MSE 为 $\min(\beta\sigma^2,\lambda)$,连续增加后饱和;每维 KL 也连续趋零。离散的存活计数会出现台阶,但不意味着误差发生跳变。更强的 KL 约束还可能让所有维度关闭,不能只用 β 大小判断解耦质量。 09. 动手验证 前三个实验只依赖 numpy,配图脚本还需 matplotlib;完整四脚本均在附录。把它们保存到同一个目录运行,reparam_gradients.py 和 make_figures.py 会导入相邻脚本: python -B elbo_identity.py python -B reparam_gradients.py python -B beta_kl_weight.py python -B make_figures.py 本轮 CPU 实跑约 1.7、0.7、1.0、6.7 秒,仅供环境对照,不承诺其他机器的耗时。模型生成与状态种子 20260925,MC 种子 7,IWAE 种子 13,梯度小模型种子 11,维度扫描种子 23。matplotlib 字体选 Microsoft YaHei/SimHei/DejaVu Sans;其他平台请换成已安装的中文字体。计算线程在本轮进程中设为 1。 实验一:比较三份 q。 把 make_q 的均值设成真实后验均值,分别用 np.diag(cov) 与 1 / np.diag(np.linalg.solve(cov,np.eye(K))) 作方差。预期 KL 缝隙为 0.79398 与 0.40038;原指定 q 为 1.98544。要到零还需完整协方差,不只是换均值或对角方差。 实验二:换独立 baseline 并查无偏条件。 用独立 pilot 样本或上一批 f 的均值作当前批次 c,比较方差与均值误差。不要直接对当前批次减自身均值后声称严格无偏;可用 leave-one-out 排除当前样本。方差改善依 f、score 与 baseline 的关系决定,不承诺一定到 4.2×。 实验三:比较 free bits 的两种归约顺序。 运行: from beta_kl_weight import closed_form_optimal print(closed_form_optimal(lam=0.25, beta=4.0, free_bits=0.5)) 这里对应数据平均 KL 外取 max,低方差方向在这个解析模型中会使用 latent;若改为逐样本 max 后求期望,须重新计算目标与最优解,不能沿用这里的闭式。 实验四:重现两个反例。 离散 IWAE 权重路径与正弦梯度的最小实验如下;第二个种子固定 73、20 万共同噪声,独立运行的方差应约为 50 与 0.5,具体有限样本数值允许变化。 import numpy as np weights = np.array([5.0, 5.0/9.0, 5.0/9.0]) print(np.log(np.cumsum(weights)/np.arange(1,4))) z = np.random.default_rng(73).normal(size=200000) print(np.var(10*np.cos(10*z)), np.var(z*np.sin(10*z))) 第一行输出 [1.60943791, 1.02165125, 0.71149632];这条可发生的路径同时反驳“单次一定是下界”和“单次随 k 单调”。它不反驳对全部路径取期望的定理。 10. 延伸阅读 按知识树的依赖关系,从这篇出发有三个方向: 往下(直接依赖本文):扩散过程的前向与反向推导——离散潜变量链的变分界与连续 SDE/score 视角怎样联系;SDE 训练并非只有 ELBO 一条路线。另一个直接后继是 VAE 结构与训练目标。 往旁(同一个估计量,另一副马甲):策略梯度与 PPO 基础——本文 3.5 的得分函数估计量在 RL 里叫 REINFORCE,方差问题一模一样,baseline 的作用也一模一样。读过这篇再去看 PPO,会省掉一次重新理解。 已发布的相关篇目:视频 VAE 的常见 loss 组合——KL 之外还有哪些 loss 会加进来,以及它们各自的量级;FlashAttention 为什么不需要存下注意力矩阵 与 性能建模与 Profiling——本文没碰的工程侧。 完整的知识树见博客目录页 AIGC 基本功知识树,按依赖顺序排好了先修课。 附录:完整代码 09 节用到的脚本全文如下(reparam_gradients.py、make_figures.py、elbo_identity.py、beta_kl_weight.py)。复制到本地存成同名文件,按各脚本开头的依赖说明准备环境后即可运行。 reparam_gradients.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """重参数化 vs 得分函数(REINFORCE)梯度估计量,方差实测。 要估的量:∇_φ E_{q_φ}[log p(x|z)],φ = (m, ℓ),ℓ = log s。 这一项没法解析求的时候只能采样,两种采样方式的方差差好几个数量级。 真实梯度在这个高斯线性模型里能解析算出来,所以「谁对谁错」有标准答案, 剩下的差别就纯粹是方差。 两部分实验: A. 固定一个小模型(D=6),看三种估计量的分布 B. 扫数据维度 D,看方差比怎么长——这里才是重参数化真正救命的地方 只依赖 numpy,直接 `python reparam_gradients.py` 即可运行。 """ import numpy as np import elbo_identity as EI from elbo_identity import ( D, K, SIGMA2, W, b, elbo_analytic, make_one_x, make_q, ) N_TRIAL = 40000 # ══════════════════════════════════════════════════════════════ # A. 小模型:三种估计量的分布 # ══════════════════════════════════════════════════════════════ def true_grad(x, m, s, Wm, bm, sigma2): """∇_m 与 ∇_ℓ 的解析梯度,ℓ = log s。""" g_m = Wm.T @ (x - (Wm @ m + bm)) / sigma2 g_ell = -(s ** 2) * np.diag(Wm.T @ Wm) / sigma2 return g_m, g_ell def log_lik_rows(x, z, Wm, bm, Dm, sigma2): """批量算 log p(x|z),z: [n, K]。""" r = x - (z @ Wm.T + bm) return -0.5 * (Dm * np.log(2 * np.pi * sigma2) + np.einsum("nd,nd->n", r, r) / sigma2) def sample_reparam(x, m, s, Wm, bm, sigma2, rng, n): """重参数化:z = m + s ⊙ ε,梯度经 z 反传。""" eps = rng.normal(size=(n, m.size)) z = m + s * eps score_z = (x - (z @ Wm.T + bm)) @ Wm / sigma2 # ∇_z log p(x|z) return score_z, score_z * s * eps # ∂z/∂m=1, ∂z/∂ℓ=s⊙ε def sample_score(x, m, s, Wm, bm, Dm, sigma2, rng, n, baseline=None): """得分函数:∇ E[f] = E[f · ∇ log q],f = log p(x|z)。""" eps = rng.normal(size=(n, m.size)) z = m + s * eps f = log_lik_rows(x, z, Wm, bm, Dm, sigma2) if baseline is not None: f = f - baseline return f[:, None] * (eps / s), f[:, None] * (eps ** 2 - 1.0) def summarize(name, gm_hat, ge_hat, g_m, g_ell): var = float(np.var(gm_hat, axis=0).sum() + np.var(ge_hat, axis=0).sum()) gnorm = float(np.linalg.norm(np.concatenate([g_m, g_ell]))) err = float(np.linalg.norm(np.concatenate([gm_hat.mean(0) - g_m, ge_hat.mean(0) - g_ell]))) / gnorm n_needed = int(np.ceil((var ** 0.5 / (0.05 * gnorm)) ** 2)) print(f" {name:<24s} 方差和={var:11.2f} 均值相对误差={err:6.2%} " f"5%相对RMS需 n≈{n_needed}") return var def part_a(): print("=" * 68) print("A. 小模型(D=6, σ=0.35):三种估计量的分布") print("=" * 68) # 与 elbo_identity.state() 同一颗种子 => 同一个 x、同一个 q rng_state = np.random.default_rng(EI.STATE_SEED) x, _ = make_one_x(rng_state) m, s = make_q(x, rng_state) g_m, g_ell = true_grad(x, m, s, W, b, SIGMA2) _, recon, kl = elbo_analytic(x, m, s) print(f"\n真实梯度(解析)") print(f" ∇_m = {np.array2string(g_m, precision=4)}") print(f" ∇_ℓ = {np.array2string(g_ell, precision=4)}") print(f" E_q[log p(x|z)] = {recon:.4f},KL(q||p(z)) = {kl:.4f}") rng = np.random.default_rng(11) n = N_TRIAL print(f"\n估计量对比({n} 次独立采样,方差 = 6 个参数分量方差之和)") gm_r, ge_r = sample_reparam(x, m, s, W, b, SIGMA2, rng, n) gm_s, ge_s = sample_score(x, m, s, W, b, D, SIGMA2, rng, n) gm_sb, ge_sb = sample_score(x, m, s, W, b, D, SIGMA2, rng, n, baseline=recon) v_r = summarize("重参数化", gm_r, ge_r, g_m, g_ell) v_s = summarize("得分函数(无 baseline)", gm_s, ge_s, g_m, g_ell) v_sb = summarize("得分函数(预言机 baseline)", gm_sb, ge_sb, g_m, g_ell) print(f"\n 方差比:得分函数 / 重参数化 = {v_s / v_r:.1f}×") print(f" 加 baseline 之后 = {v_sb / v_r:.1f}×") print(f"\n∇_m 第 0 个分量(真值 {g_m[0]:.4f})的估计分布") for name, g in (("重参数化", gm_r[:, 0]), ("得分函数", gm_s[:, 0]), ("得分函数+baseline", gm_sb[:, 0])): print(f" {name:<20s} mean={g.mean():9.4f} std={g.std():8.3f} " f"min={g.min():9.2f} max={g.max():9.2f}") print(f"\n均值误差随样本数的收敛(∇_m 第 0 个分量)") print(f" {'n':>7s} {'重参数化':>12s} {'得分函数':>12s}") for n2 in (10, 100, 1000, 10000, N_TRIAL): print(f" {n2:7d} {abs(gm_r[:n2, 0].mean() - g_m[0]):12.5f} " f"{abs(gm_s[:n2, 0].mean() - g_m[0]):12.5f}") return dict(g_m0=g_m[0], gm_r=gm_r[:, 0], gm_s=gm_s[:, 0], gm_sb=gm_sb[:, 0], var_r=v_r, var_s=v_s, var_sb=v_sb) # ══════════════════════════════════════════════════════════════ # B. 方差比随数据维度 D 怎么长 # ══════════════════════════════════════════════════════════════ def build_model(Dm, Km, seed): """构造一个 decoder 列范数归一化的线性高斯模型,让梯度尺度不随 D 漂。""" rng = np.random.default_rng(seed) Wm = rng.normal(size=(Dm, Km)) Wm /= np.linalg.norm(Wm, axis=0, keepdims=True) # 每列范数 = 1 bm = rng.normal(0.0, 0.3, size=(Dm,)) return Wm, bm def part_b(): print("\n" + "=" * 68) print("B. 方差比随数据维度 D 的变化(σ²=1,decoder 列范数归一化)") print("=" * 68) rng = np.random.default_rng(23) rows = [] print(f"\n {'D':>6s} {'重参数化':>12s} {'得分函数':>14s} {'比值':>10s} " f"{'+baseline 方差':>14s} {'比值':>10s}") for Dm in (8, 64, 512, 2048): Wm, bm = build_model(Dm, K, seed=1000 + Dm) z_true = rng.normal(size=(K,)) x = Wm @ z_true + bm + rng.normal(size=(Dm,)) # σ² = 1 m = Wm.T @ (x - bm) # 一个合理的变分均值 s = np.full((K,), 0.6) sigma2 = 1.0 n = 4000 gm_r, ge_r = sample_reparam(x, m, s, Wm, bm, sigma2, rng, n) gm_s, ge_s = sample_score(x, m, s, Wm, bm, Dm, sigma2, rng, n) # 预言机 baseline:E_q[log p(x|z)] 的解析值 resid = x - (Wm @ m + bm) recon = -0.5 * (Dm * np.log(2 * np.pi * sigma2) + (resid @ resid + np.trace(Wm.T @ Wm @ np.diag(s ** 2))) / sigma2) gm_sb, ge_sb = sample_score(x, m, s, Wm, bm, Dm, sigma2, rng, n, baseline=recon) v_r = float(np.var(gm_r, axis=0).sum() + np.var(ge_r, axis=0).sum()) v_s = float(np.var(gm_s, axis=0).sum() + np.var(ge_s, axis=0).sum()) v_sb = float(np.var(gm_sb, axis=0).sum() + np.var(ge_sb, axis=0).sum()) rows.append((Dm, v_r, v_s, v_sb)) print(f" {Dm:6d} {v_r:12.2f} {v_s:14.2f} {v_s / v_r:9.1f}× " f"{v_sb:14.2f} {v_sb / v_r:9.1f}×") print("\n 为什么长这么快:f = log p(x|z) 的均值里有一大坨与参数无关的") print(" 与 z 无关的归一化常数和残差均值。它们乘到 score 中会增加方差;") print(" 本扫描固定 K=3、归一化 W 列范数、s=0.6,非一般维度定律。") print(" baseline 减去完整 E[f];pathwise 不一定在所有问题上都更低方差。") return rows if __name__ == "__main__": a = part_a() rows = part_b() make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「变分下界与重参数化」的三张解释图。 数值全部来自同目录的三个脚本(elbo_identity / reparam_gradients / beta_kl_weight),这里只负责把它们画出来——改了那几个脚本这里要重跑, 避免图与正文数字不一致。 三张图分别回答: 1. ELBO 离真实目标差多少,差的那一项是什么,多采样能不能补上 2. 重参数化到底省了多少方差,以及这个差距随数据维度怎么长 3. KL 项加权的 β 是怎么把潜变量维度一刀一刀切掉的 只依赖 numpy + matplotlib,直接 `python make_figures.py` 即可运行。 """ from pathlib import Path import matplotlib.pyplot as plt import numpy as np from matplotlib.patches import Rectangle import beta_kl_weight as BK import elbo_identity as EI import reparam_gradients as RG ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["Microsoft YaHei", "SimHei", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, }) INK = "#1f2937" C_ELBO = "#2f6fb0" # ELBO:蓝 C_GAP = "#e0a03c" # 缺口:橙 C_BAD = "#d1495b" # 越界 / 死亡:红 C_OK = "#2f9e6f" # 收紧 / 存活:绿 C_SCORE = "#8b5cf6" # 得分函数:紫 def style(ax, title, xlabel=None, ylabel=None): ax.set_title(title, fontsize=13.0, weight="bold", color=INK, pad=10) if xlabel: ax.set_xlabel(xlabel, fontsize=11, color="#475569") if ylabel: ax.set_ylabel(ylabel, fontsize=11, color="#475569") ax.tick_params(labelsize=10, colors="#64748b") for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("left", "bottom"): ax.spines[s].set_color("#cbd5e1") ax.grid(axis="y", color="#eef2f7", lw=1.0) ax.set_axisbelow(True) # ══════════════════════════════════════════════════════════════ # 图 1:ELBO 离 log p(x) 差多少 # ══════════════════════════════════════════════════════════════ def fig_elbo_gap(): st = EI.state() mc = EI.mc_table(st) iw = EI.iwae_table(st) log_px, elbo, gap, kl_prior = st["log_px"], st["elbo"], st["gap"], st["kl"] fig, axes = plt.subplots(1, 3, figsize=(16.2, 4.9)) fig.patch.set_facecolor("white") # ── (a) 分解:两条横向长条 ── ax = axes[0] style(ax, "(a) log p(x) 拆成两截", "nats(对数似然,0 在右边)") rows = [ ("ELBO(能算,优化它)", elbo, C_ELBO, "white"), ("log p(x)(本例闭式真值)", log_px, C_BAD, "white"), ("KL(q‖p(z)):loss 里那一项", -kl_prior, "#94a3b8", "white"), ("KL(q‖p(z|x)):上面两条的差", -gap, C_GAP, "white"), ] for i, (name, v, color, tc) in enumerate(rows): y = len(rows) - 1 - i ax.barh(y, -v, left=v, height=0.52, color=color) ax.text(v + 0.18, y, f"{v if i < 2 else abs(v):.3f}", ha="left", va="center", fontsize=10, color=INK, weight="bold") ax.axvline(log_px, color=C_BAD, ls="--", lw=1.3) ax.axvline(elbo, color=C_ELBO, ls="--", lw=1.3) ax.annotate("", xy=(elbo, 3.28), xytext=(log_px, 3.28), arrowprops=dict(arrowstyle="<->", color="#a16207", lw=1.8)) ax.text((elbo + log_px) / 2, 3.75, "差 = 1.985", ha="center", va="bottom", fontsize=10.5, color="#a16207", weight="bold", bbox=dict(fc="white", ec="none", alpha=0.9, pad=1.0)) ax.text(-9.35, 0.68, "上面两条之差 = 下面橙色那条,\n不是灰色那条——这两个 KL 常被混为一谈", ha="left", va="center", fontsize=9, color="#64748b", linespacing=1.5) ax.set_yticks(range(len(rows))) ax.set_yticklabels([r[0] for r in rows][::-1], fontsize=9.5) ax.set_xlim(-9.6, 1.5) ax.set_ylim(-0.6, 3.9) ax.grid(axis="x", color="#eef2f7", lw=1.0) ax.grid(axis="y", visible=False) # ── (b) 单样本 MC 会越过上界 ── ax = axes[1] style(ax, f"(b) L=1 的 MC 估计({EI.MC_REP} 次)", "单次估计值", "频数") vals = mc[1]["vals"] ax.hist(vals, bins=70, color="#cbd5e1", edgecolor="white", lw=0.4) over = vals[vals > log_px] ax.hist(over, bins=70, color=C_BAD, alpha=0.85, label=f"{mc[1]['frac']:.1%} 越过了 log p(x)") ax.axvline(elbo, color=C_ELBO, lw=2.0, label=f"ELBO = {elbo:.3f}") ax.axvline(log_px, color=C_BAD, lw=2.0, ls="--", label=f"log p(x) = {log_px:.3f}") ax.set_xlim(-26, 7) ax.legend(fontsize=9, frameon=False, loc="upper left") ax.set_yscale("log") ax.set_ylim(0.7, 900) ax.text(0.03, 0.05, "估计量是无偏的:它在 ELBO 周围晃,\n不是「恒小于 log p(x)」", transform=ax.transAxes, ha="left", va="bottom", fontsize=9.5, color="#64748b", linespacing=1.5) # ── (c) IWAE 收紧 ── ax = axes[2] style(ax, f"(c) 多采样收紧(每项 {EI.IW_REP} 次)", "k(每次采几个 z)", "nats") ks = np.array(EI.IW_KS, dtype=float) means = np.array([iw[k]["mean"] for k in EI.IW_KS]) stds = np.array([iw[k]["std"] for k in EI.IW_KS]) ax.errorbar(ks, means, yerr=stds, fmt="o-", color=C_OK, lw=2.0, capsize=4, markersize=6) ax.axhline(log_px, color=C_BAD, ls="--", lw=1.6, label=f"log p(x) = {log_px:.3f}") ax.axhline(elbo, color=C_ELBO, ls=":", lw=1.6, label=f"ELBO = {elbo:.3f}") ax.set_xscale("log") ax.set_ylim(-9.4, -4.7) ax.set_xlim(0.7, 900) for k, m in zip(ks, means): ax.annotate(f"还差 {log_px - m:.3f}", (k, m), textcoords="offset points", xytext=(8, 12 if k > 1 else -18), fontsize=9, color="#64748b", bbox=dict(fc="white", ec="none", alpha=0.85, pad=0.6)) ax.legend(fontsize=9, frameon=False, loc="lower right") ax.text(0.03, 0.06, "期望界更紧;误差棒是单次标准差", transform=ax.transAxes, fontsize=9.5, color="#64748b") fig.suptitle("这张图要看什么:ELBO 与真实目标之间那道缝," "固定 q 的 KL 缝隙与采样波动是两件事", fontsize=11.5, color="#64748b", y=1.02) fig.tight_layout() p = OUT / "elbo_gap.png" fig.savefig(p, bbox_inches="tight", facecolor="white") plt.close(fig) print(f" ✓ {p.name}") # ══════════════════════════════════════════════════════════════ # 图 2:重参数化省了多少方差 # ══════════════════════════════════════════════════════════════ def fig_grad_variance(): a = RG.part_a() rows = RG.part_b() fig, axes = plt.subplots(1, 2, figsize=(13.4, 5.0)) fig.patch.set_facecolor("white") # ── (a) 估计分布 ── ax = axes[0] style(ax, "(a) ∂L/∂m 第 0 个分量的估计分布", "估计值", "频数") bins = np.linspace(-70, 70, 90) for vals, color, name in ((a["gm_r"], C_ELBO, "重参数化"), (a["gm_s"], C_SCORE, "得分函数"), (a["gm_sb"], C_GAP, "得分函数+baseline")): ax.hist(vals, bins=bins, histtype="step", lw=1.8, color=color, label=f"{name} std={vals.std():.1f}") ax.axvline(a["g_m0"], color=C_OK, lw=2.2, label=f"真值 = {a['g_m0']:.3f}") ax.set_xlim(-70, 70) ax.set_yscale("log") ax.legend(fontsize=9.5, frameon=False, loc="upper left") ax.text(0.98, 0.05, "理论均值相同,有限样本均值有误差:\n" f"std = {a['gm_r'].std():.1f} / {a['gm_s'].std():.1f} / " f"{a['gm_sb'].std():.1f}", transform=ax.transAxes, ha="right", va="bottom", fontsize=9.5, color="#64748b", linespacing=1.5) # ── (b) 随维度怎么长 ── ax = axes[1] style(ax, "(b) 方差随数据维度 D 怎么长", "D(数据维度)", "梯度方差(6 个分量之和)") Ds = np.array([r[0] for r in rows], dtype=float) vr = np.array([r[1] for r in rows]) vs = np.array([r[2] for r in rows]) vsb = np.array([r[3] for r in rows]) ax.loglog(Ds, vs, "o-", color=C_SCORE, lw=2.0, markersize=7, label="得分函数(无 baseline)") ax.loglog(Ds, vsb, "s--", color=C_GAP, lw=2.0, markersize=6, label="得分函数 + baseline") ax.loglog(Ds, vr, "o-", color=C_ELBO, lw=2.0, markersize=7, label="重参数化") # 斜率 2 的参考线 ref = vs[0] * (Ds / Ds[0]) ** 2 ax.loglog(Ds, ref, ":", color="#cbd5e1", lw=2.0, label="∝ D² 参考线") for d, v, r in zip(Ds, vs, vr): ratio = v / r txt = f"{ratio:,.0f}×" if ratio < 1e4 else f"{ratio:.1e}×" ax.annotate(txt, (d, v), textcoords="offset points", xytext=(8, -14), fontsize=9.5, color=C_SCORE, weight="bold", bbox=dict(fc="white", ec="none", alpha=0.85, pad=0.6)) ax.set_ylim(2e-1, 1e10) ax.legend(fontsize=9.5, frameon=False, loc="upper left") ax.text(0.97, 0.03, "本归一化扫描中 pathwise 方差近似稳定", transform=ax.transAxes, fontsize=9.5, color=C_ELBO, ha="right", va="bottom", weight="bold") fig.suptitle("这张图要看什么:重参数化不是「更准」,是「腰围更小」——" "本实验分布更集中;方差排序不是普遍定理", fontsize=11.5, color="#64748b", y=1.02) fig.tight_layout() p = OUT / "grad_variance.png" fig.savefig(p, bbox_inches="tight", facecolor="white") plt.close(fig) print(f" ✓ {p.name}") # ══════════════════════════════════════════════════════════════ # 图 3:β 怎么把潜变量维度切掉 # ══════════════════════════════════════════════════════════════ def fig_beta_tradeoff(): table = BK.sweep(BK.LAMS, BK.BETAS, verbose=False) s2 = BK.SIGMA2 fig, axes = plt.subplots(1, 2, figsize=(13.4, 5.0)) fig.patch.set_facecolor("white") # ── (a) 重建误差 vs β ── ax = axes[0] style(ax, "(a) β 越大,重建越差——直到某一维被直接放弃", "β(KL 项的权重)", "该维的重建 MSE") colors = ["#2f6fb0", "#2f9e6f", "#e0a03c", "#d1495b"] bgrid = np.linspace(0.2, 14, 100) ax.plot(bgrid, bgrid * s2, ":", color="#94a3b8", lw=1.8, label="理论:MSE = βσ²(还活着)") for lam, c in zip(BK.LAMS, colors): betas, mses, alive = [], [], [] for beta in BK.BETAS: a, b, s, mse, kl, au = table[(lam, beta)] betas.append(beta) mses.append(mse) alive.append(BK.active_unit(a, lam)) betas = np.array(betas) mses = np.array(mses) alive = np.array(alive) ax.plot(betas[alive], mses[alive], "o-", color=c, lw=2.0, markersize=6, label=f"λ={lam}") if (~alive).any(): ax.plot(betas[~alive], mses[~alive], "s", color=c, markersize=7, mfc="white", mew=1.8) ax.axhline(mses[~alive][0], color=c, ls="--", lw=1.0, alpha=0.45) ax.set_xscale("log") ax.set_yscale("log") ax.set_ylim(0.028, 7.0) ax.legend(fontsize=9, frameon=False, loc="upper left", ncol=2) ax.text(0.5, 0.012, "实心圆 = 该维还活着(MSE 贴着 βσ² 往上走);" "空心方块 = 该维已死,MSE 停在 λ", transform=ax.transAxes, ha="center", va="bottom", fontsize=9.5, color="#64748b") # ── (b) 相图 ── ax = axes[1] style(ax, "(b) 存活 / 死亡的相图", "λ / σ²(该维的信噪比)", "β") snr = np.array(BK.LAMS) / s2 grid = np.logspace(-1, 1.6, 60) ax.loglog(grid, grid, "-", color="#94a3b8", lw=2.2, label="理论分界:β = λ/σ²") ax.fill_between(grid, grid, 1e3, color=C_BAD, alpha=0.06) ax.fill_between(grid, 1e-3, grid, color=C_OK, alpha=0.06) for lam in BK.LAMS: for beta in BK.BETAS: a, b, s, mse, kl, au = table[(lam, beta)] ok = BK.active_unit(a, lam) ax.plot(lam / s2, beta, "o" if ok else "s", color=C_OK if ok else C_BAD, markersize=7, mfc=C_OK if ok else "white", mew=1.6) ax.text(0.05, 0.93, "上方:KL 太贵,维度被关掉", transform=ax.transAxes, fontsize=10, color=C_BAD, weight="bold") ax.text(0.05, 0.06, "下方:数据方差信噪比盖过 KL 的价,维度存活", transform=ax.transAxes, fontsize=10, color=C_OK, weight="bold") ax.set_xlim(0.1, 40) ax.set_ylim(0.15, 20) ax.legend(fontsize=9.5, frameon=False, loc="center right") ax.grid(axis="y", color="#eef2f7", lw=1.0) fig.suptitle("这张图要看什么:线性高斯特例:β 给 KL 加权," "等于给每个潜变量维度标一个价——信噪比不够的维度直接归零", fontsize=11.5, color="#64748b", y=1.02) fig.tight_layout() p = OUT / "beta_tradeoff.png" fig.savefig(p, bbox_inches="tight", facecolor="white") plt.close(fig) print(f" ✓ {p.name}") if __name__ == "__main__": print("生成配图 →", OUT) fig_elbo_gap() fig_grad_variance() fig_beta_tradeoff() print("完成:3 张") elbo_identity.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """ELBO 恒等式核对:log p(x) = ELBO + KL(q || p(z|x))。 构造一个「后验有解析解」的模型,这样三样东西都能算到机器精度: * 边际似然 log p(x) —— 高斯线性模型,边际仍是高斯 * 真实后验 p(z|x) —— 共轭高斯,闭式解 * ELBO —— 重建项对 q 可解析求期望,KL 两项都是高斯闭式 有解析解才能验证「差的那一项到底是什么」,靠采样是验不出来的。 只依赖 numpy,直接 `python elbo_identity.py` 即可运行。 """ import numpy as np RNG = np.random.default_rng(20260925) # ── 模型:z ~ N(0, I_K),x|z ~ N(W z + b, sigma^2 I_D) ── D, K = 6, 3 SIGMA = 0.35 # 解码器观测噪声标准差 SIGMA2 = SIGMA ** 2 W = RNG.normal(0.0, 1.0, size=(D, K)) * 0.9 # 解码器权重 b = RNG.normal(0.0, 0.3, size=(D,)) # 解码器偏置 def make_one_x(rng=None): """从真实的边际分布里采一个 x,并返回它的解析 log p(x)。 rng 可显式传入:这样 reparam_gradients / make_figures 拿到的是同一个 x, 不会因为模块级 RNG 被别人消耗过而对不上数字。 """ if rng is None: rng = RNG z = rng.normal(size=(K,)) x = W @ z + b + SIGMA * rng.normal(size=(D,)) cov = W @ W.T + SIGMA2 * np.eye(D) sign, logdet = np.linalg.slogdet(cov) d = x - b quad = d @ np.linalg.solve(cov, d) log_px = -0.5 * (D * np.log(2 * np.pi) + logdet + quad) return x, log_px def true_posterior(x): """p(z|x) 的闭式解:precision = I + W^T W / sigma^2。""" prec = np.eye(K) + W.T @ W / SIGMA2 cov = np.linalg.inv(prec) mean = cov @ (W.T @ (x - b) / SIGMA2) return mean, cov # ── 变分后验 q(z|x) = N(m, diag(s^2)):故意用「对角」错误指定 ── def make_q(x, rng=None): """取真实后验的均值,协方差只留对角线,再乘 0.8 并加 0.05。 真实后验是有相关性的(cov 非对角),q 强行对角 => KL(q||p(z|x)) > 0, 这正是我们要留出来的那条缝。 """ if rng is None: rng = RNG mean, cov = true_posterior(x) m = mean + 0.15 * rng.normal(size=(K,)) # 均值也偏一点 var = np.diag(cov).copy() * 0.8 + 0.05 # 与原方差大小须逐项比较 return m, np.sqrt(var) def log_gauss_diag(z, mean, var, ): """log N(z; mean, diag(var))(省略常数也行,但这里算全)。""" return float(-0.5 * np.sum(np.log(2 * np.pi * var) + (z - mean) ** 2 / var)) def log_joint(x, z): """log p(x, z) = log p(x|z) + log p(z)。""" r = x - (W @ z + b) log_lik = -0.5 * (D * np.log(2 * np.pi * SIGMA2) + r @ r / SIGMA2) log_prior = -0.5 * (K * np.log(2 * np.pi) + z @ z) return log_lik + log_prior def elbo_analytic(x, m, s): """ELBO 的解析值。 E_q[log p(x|z)] 里 ||x - Wz - b||^2 对 q 求期望: ||x - Wm - b||^2 + tr(W^T W diag(s^2)) 第二项是采样噪声经过解码器放大出来的那部分,容易被漏掉。 """ var = s ** 2 resid = x - (W @ m + b) recon = -0.5 * (D * np.log(2 * np.pi * SIGMA2) + (resid @ resid + np.trace(W.T @ W @ np.diag(var))) / SIGMA2) kl = 0.5 * np.sum(var + m ** 2 - 1.0 - np.log(var)) return float(recon - kl), float(recon), float(kl) def kl_gauss_gauss(m, s, mean2, cov2): """KL(N(m, diag(s^2)) || N(mean2, cov2)),两个高斯的闭式。""" var = s ** 2 diff = mean2 - m cov2_inv = np.linalg.inv(cov2) trace = float(np.trace(cov2_inv @ np.diag(var))) quad = float(diff @ cov2_inv @ diff) _, logdet2 = np.linalg.slogdet(cov2) return 0.5 * (trace + quad - K + logdet2 - np.sum(np.log(var))) def mc_elbo(x, m, s, L, rng): """蒙特卡洛 ELBO:L 个样本取平均(L=1 就是训练时实际用的那个)。""" eps = rng.normal(size=(L, K)) z = m + s * eps residual = x - (z @ W.T + b) log_lik = -0.5 * (D * np.log(2 * np.pi * SIGMA2) + np.sum(residual ** 2, axis=1) / SIGMA2) var = s ** 2 kl = 0.5 * np.sum(var + m ** 2 - 1.0 - np.log(var)) return float(log_lik.mean() - kl) def iwae_log_bound(x, m, s, k, rng): """IWAE 的 k 样本界:log (1/k) sum_l p(x,z_l)/q(z_l),用 logsumexp 稳算。""" eps = rng.normal(size=(k, K)) z = m + s * eps var = s ** 2 residual = x - (z @ W.T + b) log_lik = -0.5 * (D * np.log(2 * np.pi * SIGMA2) + np.sum(residual ** 2, axis=1) / SIGMA2) log_prior = -0.5 * (K * np.log(2 * np.pi) + np.sum(z ** 2, axis=1)) log_q = -0.5 * np.sum(np.log(2 * np.pi * var) + (z - m) ** 2 / var, axis=1) logw = log_lik + log_prior - log_q mx = logw.max() return float(mx + np.log(np.mean(np.exp(logw - mx)))) # ── 下面两个表被 main 与 make_figures 共用,种子固定 => 数字永远一致 ── MC_LS = (1, 4, 16, 64) MC_REP = 4000 IW_KS = (1, 5, 50, 500) IW_REP = 6000 STATE_SEED = 20260925 def state(): """文章与配图共用的那一份状态:同一个 x、同一个 q。 用固定种子重放,保证单独跑本脚本与跑 make_figures 拿到同一组数字。 """ rng = np.random.default_rng(STATE_SEED) x, log_px = make_one_x(rng) m, s = make_q(x, rng) post_mean, post_cov = true_posterior(x) elbo, recon, kl = elbo_analytic(x, m, s) gap = kl_gauss_gauss(m, s, post_mean, post_cov) return dict(x=x, log_px=log_px, m=m, s=s, elbo=elbo, recon=recon, kl=kl, gap=gap) def mc_table(st): """蒙特卡洛 ELBO 的均值 / 标准差 / 越过 log p(x) 的比例。""" rng = np.random.default_rng(7) out = {} for L in MC_LS: vals = np.array([mc_elbo(st["x"], st["m"], st["s"], L, rng) for _ in range(MC_REP)]) out[L] = dict(mean=float(vals.mean()), std=float(vals.std()), frac=float(np.mean(vals > st["log_px"])), vals=vals) return out def iwae_table(st): """IWAE 的 k 样本界。""" rng = np.random.default_rng(13) out = {} for k in IW_KS: vals = np.array([iwae_log_bound(st["x"], st["m"], st["s"], k, rng) for _ in range(IW_REP)]) out[k] = dict(mean=float(vals.mean()), std=float(vals.std()), gap=float(st["log_px"] - vals.mean()), se=float(vals.std(ddof=1) / np.sqrt(IW_REP)), frac=float(np.mean(vals > st["log_px"])), vals=vals) return out def main(): print("=" * 68) print("ELBO 恒等式核对 log p(x) = ELBO + KL(q || p(z|x))") print("=" * 68) print(f"模型: D={D}, K={K}, sigma={SIGMA}") st = state() x, log_px, m, s = st["x"], st["log_px"], st["m"], st["s"] elbo, recon, kl, gap = st["elbo"], st["recon"], st["kl"], st["gap"] print(f"\n[1] 解析三项") print(f" E_q[log p(x|z)] = {recon: .6f}") print(f" KL(q || p(z)) = {kl: .6f}") print(f" ELBO = {elbo: .6f}") print(f" KL(q || p(z|x)) = {gap: .6f}") print(f" ELBO + KL(q||p(z|x)) = {elbo + gap: .6f}") print(f" log p(x) = {log_px: .6f}") print(f" |误差| = {abs(elbo + gap - log_px): .3e}") # ── 蒙特卡洛波动:ELBO 的估计量是无偏的,不是恒小于 log p(x) ── print(f"\n[2] 蒙特卡洛估计的波动(每项 {MC_REP} 次重复)") print(f" {'L':>5s} {'mean':>12s} {'std':>10s} {'超过 log p(x) 的比例':>22s}") mc = mc_table(st) for L in MC_LS: r = mc[L] print(f" {L:5d} {r['mean']:12.5f} {r['std']:10.5f} {r['frac']:21.1%}") print(f" log p(x) = {log_px:.5f}(上界本身),ELBO(解析) = {elbo:.5f}") # ── IWAE:k 越大越紧 ── print(f"\n[3] IWAE 的 k 样本界(每项 {IW_REP} 次重复)") print(f" {'k':>5s} {'mean L_k':>12s} {'std':>9s} {'log p(x) - L_k':>16s}") iw = iwae_table(st) for k in IW_KS: r = iw[k] print(f" {k:5d} {r['mean']:12.5f} {r['std']:9.5f} {r['gap']:16.5f} SE={r['se']:.5f} 越界={r['frac']:.1%}") print(f" 参考:单样本 MC ELBO 的均值 = {elbo:.5f}(= ELBO 解析值,不随 k 变)") print("\n结论:期望下界 L_k 随 k 单调;单次结果和有限重复均值可以越界或不单调;" "而「取平均再 log」和「log 再取平均」是两件事。") if __name__ == "__main__": main() beta_kl_weight.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """β-VAE 的 KL 加权:一个能解到最优的玩具模型,看潜变量维度怎么死掉。 构造:数据每个坐标独立,x_j ~ N(0, λ_j),潜变量每个坐标也独立, q(z_j | x_j) = N(a_j x_j, s_j²),p(x_j | z_j) = N(b_j z_j, σ²) 这样每个维度就是一份独立副本,最优解可由闭式得到;保留多起点梯度上升用于独立数值对照。 单维目标(要最大化),把 λ 与 σ² 都当作常数: L(a, b, ℓ) = -[λ(1 - b a)² + b² s²] / (2σ²) - β · KL KL = 0.5 · (a² λ + s² - 1 - log s²), s = exp(ℓ) 梯度全部解析,脚本里再用有限差分核对一遍,避免推错。 只依赖 numpy,直接 `python beta_kl_weight.py` 即可运行。 """ import numpy as np SIGMA2 = 0.25 # 解码器观测噪声方差 def kl_of(a, lam, s): """KL(N(a x, s²) || N(0, 1)),对 x ~ N(0, λ) 取期望后的形式。""" return 0.5 * (a ** 2 * lam + s ** 2 - 1.0 - np.log(s ** 2)) def recon_mse(a, b, lam, s): """E_{x,z}[(x - b z)²],重建误差的期望。""" return lam * (1.0 - b * a) ** 2 + b ** 2 * s ** 2 def objective(p, lam, beta, free_bits=None): """目标函数值(越大越好)。free_bits=C 时 KL 项取 max(KL, C)。""" a, b, ell = p s = np.exp(ell) kl = kl_of(a, lam, s) kl_eff = max(kl, free_bits) if free_bits is not None else kl return -(recon_mse(a, b, lam, s)) / (2 * SIGMA2) - beta * kl_eff def grad(p, lam, beta, free_bits=None): """解析梯度 ∇(∂L/∂a, ∂L/∂b, ∂L/∂ℓ)。""" a, b, ell = p s2 = np.exp(2 * ell) g = np.zeros(3) g[0] = lam * b * (1.0 - b * a) / SIGMA2 g[1] = (lam * a * (1.0 - b * a) - b * s2) / SIGMA2 g[2] = -b ** 2 * s2 / SIGMA2 if free_bits is None or kl_of(a, lam, np.exp(ell)) > free_bits: # KL 项对 (a, b, ℓ) 的梯度 g[0] -= beta * a * lam g[2] -= beta * (s2 - 1.0) return g def grad_fd(p, lam, beta, free_bits=None, h=1e-6): """中心差分梯度,用来核对解析梯度有没有推错。""" g = np.zeros(3) for i in range(3): e = np.zeros(3) e[i] = h g[i] = (objective(p + e, lam, beta, free_bits) - objective(p - e, lam, beta, free_bits)) / (2 * h) return g def optimize(lam, beta, free_bits=None, n_init=9, n_iter=12000, seed=0): """多起点梯度上升,返回最好的 (a, b, s)。 必须多起点:(a, b) = (0, 0) 是「维度死亡」解,单起点容易掉进去出不来。 """ rng = np.random.default_rng(seed) best_p, best_v = None, -np.inf for k in range(n_init): if k == 0: p = np.array([0.9, 0.9, np.log(0.6)]) elif k == 1: p = np.array([0.0, 0.0, 0.0]) # 死亡解,也让它试试 else: p = np.array([rng.uniform(-1.5, 1.5), rng.uniform(-1.5, 1.5), rng.uniform(-1.2, 0.5)]) lr = 0.02 for t in range(n_iter): g = grad(p, lam, beta, free_bits) g = np.clip(g, -50.0, 50.0) # 梯度裁剪:b²s²/σ² 那一项能把参数炸飞 p = p + lr * g # 数值保护,别让 exp(ℓ) 或 (a, b) 爆掉 p[0] = float(np.clip(p[0], -8.0, 8.0)) p[1] = float(np.clip(p[1], -8.0, 8.0)) p[2] = float(np.clip(p[2], -4.0, 1.0)) if t % 1000 == 999: lr *= 0.6 # 收尾再磨一遍:β 很小的时候收敛慢,不磨的话 s² 会差到 1e-3 lr = 1e-4 for t in range(6000): g = np.clip(grad(p, lam, beta, free_bits), -50.0, 50.0) p = p + lr * g v = objective(p, lam, beta, free_bits) if v > best_v: best_v, best_p = v, p a, b, ell = best_p return float(a), float(b), float(np.exp(ell)), float(best_v) def closed_form_optimal(lam, beta, free_bits=None): """数据期望 KL 外取 max(KL,C) 的闭式全局最优(省略似然常数)。""" if lam <= 0 or beta <= 0 or (free_bits is not None and free_bits < 0): raise ValueError("lam and beta must be positive; C must be nonnegative") rate = max(0.0, 0.5 * np.log(lam / (beta * SIGMA2))) if free_bits is not None: rate = max(rate, free_bits) t = np.exp(-2 * rate) a, b, s = np.sqrt((1 - t) / lam), np.sqrt(lam * (1 - t)), np.sqrt(t) return float(a), float(b), float(s), float(objective(np.array([a,b,np.log(s)]),lam,beta,free_bits)) def active_unit(a, lam, thr=0.01): """IWAE 5.2 节的 active unit 经验判据:Cov_x(E_q[z]) = a²λ 是否超过阈值。""" return a ** 2 * lam > thr # 文章与配图共用的扫描范围 LAMS = [4.0, 1.0, 0.25, 0.0625] BETAS = [0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0, 12.0] def sweep(lams, betas, verbose=True): """扫 (λ, β) 网格,返回 {(λ, β): (a, b, s, MSE, KL, a²λ)}。""" table = {} if verbose: print("\n[1] β 扫描:每个 (λ, β) 的最优解") print(f"\n {'λ':>7s} {'β':>5s} {'a':>8s} {'b':>8s} {'s':>8s} " f"{'重建MSE':>9s} {'KL':>8s} {'a²λ':>8s} {'存活':>5s}") for lam in lams: for beta in betas: a, b, s, _ = closed_form_optimal(lam, beta) kl = kl_of(a, lam, s) mse = recon_mse(a, b, lam, s) au = a ** 2 * lam table[(lam, beta)] = (a, b, s, mse, kl, au) if verbose: print(f" {lam:7.4f} {beta:5.2f} {a:8.4f} {b:8.4f} {s:8.4f} " f"{mse:9.4f} {kl:8.4f} {au:8.4f} " f"{'是' if active_unit(a, lam) else '死':>5s}") return table def main(): print("=" * 72) print("β-VAE 的 KL 加权:潜变量维度怎么死掉") print(f"σ² = {SIGMA2}(解码器噪声),单维独立副本,闭式全局最优;多起点梯度上升独立对照") print("=" * 72) # ── 先核对解析梯度 ── print("\n[0] 解析梯度 vs 有限差分") for lam, beta in ((4.0, 1.0), (0.25, 4.0)): p = np.array([0.7, -0.4, np.log(0.8)]) ga, gf = grad(p, lam, beta), grad_fd(p, lam, beta) print(f" λ={lam:<5} β={beta:<4} 解析={np.array2string(ga, precision=5)} " f"差分={np.array2string(gf, precision=5)} 最大偏差={np.abs(ga - gf).max():.2e}") print("\n[0b] 独立多起点梯度上升 vs 闭式(数值误差,不当成全局证明)") for lam,beta in [(1.0,1.0),(0.25,4.0),(4.0,0.25)]: numeric=optimize(lam,beta,n_init=3,n_iter=3000,seed=17) exact=closed_form_optimal(lam,beta) print(f" lambda={lam} beta={beta} objective_gap={exact[3]-numeric[3]:.6g} " f"MSE_numeric={recon_mse(*numeric[:2],lam,numeric[2]):.6g} MSE_exact={min(beta*SIGMA2,lam):.6g}") table = sweep(LAMS, BETAS) print("\n[2] 每个 λ 的「死亡阈值」:β 到多大时这个维度不再被用") for lam in LAMS: dead = [beta for beta in BETAS if not active_unit(table[(lam, beta)][0], lam)] if dead: print(f" λ={lam:<7.4f} 从 β={min(dead)} 起死亡" f"(β=1 时 KL={table[(lam, 1.0)][4]:.4f})") else: print(f" λ={lam:<7.4f} 在 β≤{max(BETAS)} 内都存活") # ── 闭式解核对:存活时 s² = βσ²/λ,a²λ = 1 - s²,重建 MSE = βσ² ── print("\n[2b] 闭式解核对(存活的格子才成立)") print(f" {'λ':>7s} {'β':>5s} {'s²闭式':>9s} {'βσ²/λ':>9s} " f"{'a²λ闭式':>9s} {'1-s²':>9s} {'MSE闭式':>9s} {'βσ²':>9s}") worst = 0.0 for lam in LAMS: for beta in BETAS: a, b, s, _, kl, au = table[(lam, beta)] if not active_unit(a, lam): continue s2_pred = beta * SIGMA2 / lam rows = (s ** 2, s2_pred, au, 1 - s ** 2, recon_mse(a, b, lam, s), beta * SIGMA2) print(f" {lam:7.4f} {beta:5.2f} " + " ".join(f"{v:9.5f}" for v in rows)) worst = max(worst, abs(s ** 2 - s2_pred), abs(au - (1 - s ** 2)), abs(recon_mse(a, b, lam, s) - beta * SIGMA2)) print(f" 最大偏差 = {worst:.2e} → 存活时闭式内部一致性(不代替独立数值核验)") print(f" 推论:维度存活条件 s² < 1 ⟺ β < λ/σ²,即该维的信噪比要盖过 KL 的价") for lam in LAMS: print(f" λ={lam:<7.4f} 理论死亡阈值 β* = λ/σ² = {lam / SIGMA2:.2f}") # ── free bits:数据期望KL外的截断惩罚;非信息下限保证 ── print("\n[3] free bits 对照(β=4,KL 项取 max(KL, C))") print(f" {'λ':>7s} {'C=0(纯 β)':>22s} {'C=0.05':>22s} {'C=0.2':>22s}") for lam in LAMS: cells = [] for C in (None, 0.05, 0.2): a, b, s, _ = closed_form_optimal(lam, 4.0, free_bits=C) cells.append(f"a={a:5.3f} KL={kl_of(a, lam, s):5.3f} " f"MSE={recon_mse(a, b, lam, s):5.3f}") print(f" {lam:7.4f} " + " ".join(f"{c:>22s}" for c in cells)) print("\n[4] 一个直觉:β 大了以后,模型宁可不要这个维度") print(" λ=0.0625 时这个坐标的数据方差本来就小于解码器噪声 σ²=0.25,") print(" 用它换来的重建收益抵不上 KL 的代价,最优解就是把它整条关掉。") if __name__ == "__main__": main()
2026年09月25日
7 阅读
0 评论
0 点赞
2026-09-25
AIGC 基本功|FlashAttention 为什么不需要存下注意力矩阵-FlashAttn
AIGC 基本功|FlashAttention 为什么不需要存下注意力矩阵-FlashAttn 所属方向:注意力与位置编码 | 难度:进阶 | 前置知识:自注意力机制的计算与显存账本(本篇是它的直接后继) 关键词:FlashAttention、online softmax、tiling、IO感知、显存优化 01. 为什么需要它 上一篇《自注意力机制的计算与显存账本》结尾留了一个没解的结:N=4096 的视频 DiT,按保守教学分配模型计,单层注意力为 2.39 GiB 激活,无重算训练时按同一假设累加 32 层就是 76.5 GiB,尚未包含权重、梯度、优化器与 FFN。推理的临时激活则不能直接乘层数。当时摆出了两条出路:稀疏化(算得少,但要看清赔的是什么质量)和 FlashAttention。这一篇把后者讲透。 先把一个流传很广的说法钉死:FlashAttention 不是近似注意力。它算出来的就是标准的 softmax 注意力,在实数算术下等价,浮点下允许舍入差异(第 04 节有实测:本次 float64 实验最大绝对误差不超过 8.9e-16;这组输入的舍入量级,非通用误差保证)。它快的理由也不神秘——附录 io_ledger.py 给出的同类型元素教学账本,单头 d=64、N=4096 时: 朴素实现要在 HBM(显存)上搬 130 MiB:分数矩阵 S 写一次读一次,softmax 权重 P 写一次读一次,4N² 次元素搬运; FlashAttention 只搬 17.8 MiB,少了 7.3 倍。 两种实现的主导矩阵乘 FLOPs 相同,但重算和归一化开销不同。在这组 A100 参数的 roofline 模型中,朴素实现算术强度为 31.5 FLOP/byte,低于约 200 的平衡点;这支持优先优化 IO 的方向,不证明所有序列、硬件或近似注意力方法都必然更慢。本文 130 MiB、17.8 MiB 是脚本估算值,不是 GPU 访存计数器实测。 另一个结论更值钱:显式保存概率矩阵 P 的朴素实现仍需二次存储,S 本身通常不必一并保留,该显式物化实现的主导存储随 N 二次增长;FlashAttention 把这两个矩阵整个从显存里删掉了,训练激活从 O(N²) 降到 O(N)——在统一采用六份线性张量、忽略小的 LSE 与工作区的教学预算中,上面那个 76.5 GiB 降为 4.5 GiB。N=32768 的长视频任务,本文教学模型的朴素实现要约 4.54 TiB 激活,超过这里对比的单卡教学预算;融合侧同一账本为 36 GiB,仍未计 FFN、权重、优化器和完整工作区。FlashAttention 为长序列精确注意力减少了重要的容量障碍;具体模型能否训练还受其他层、并行策略和硬件影响。 02. 最小可用理解 三句话讲完核心思想: softmax attention 每行需要两个标量和一个 d 维向量:这行的最大值 m、指数和 l、未归一化加权和 U。把分数按块流进来,每块用一条递推式把这三个量修正一次,全流完之后 U÷l 就是精确的 softmax 注意力输出——中间任何一步都不需要把整行摆在内存里。这叫 online softmax。 显式物化 N×N 经常造成较低的算术强度;受限类型取决于维度、硬件与实现。朴素实现把 N×N 的 S、P 写回显存再读回来;FlashAttention 用 tiling(分块)把它们关在片上 片上存储里算完就扔,在原论文两级存储模型及 d≤M≤Nd 的范围内,HBM 读写量由 Θ(Nd+N²) 变为 Θ(N²d²/M);M 按可存的同类型元素数计,实际 kernel 的混合精度预算另算。 反向传播要用的 S、P 全部重算。除 Q/K/V(或它们的重算来源)外,前向存输出 O 和每行的 logsumexp(都是 O(N)),反向时拿 Q、K 重新过一遍分块流程,把需要的局部 P 重新算出来——用重复计算换显存,这是整个方法里最「敢」的一步。 这张图要看什么:左边朴素实现的 S、P 两个 N×N 是该前向示意中的主要中间量,它们必须写回 HBM 再读回来;右边是同一个 N×N 被切成 B_r×B_c 的小块,K/V 块进 SRAM 常驻、Q 块逐行流过,跨块只有 O_i、l_i、m_i 三个 O(N) 的量一直活着。 03. 数学推导 3.1 出发点:safe softmax 为什么需要先看全一行 设一行分数为 $s_1, \dots, s_N$,softmax 的定义是 $$p_i = \frac{\exp(s_i)}{\sum_{j=1}^{N} \exp(s_j)}$$ 先修要求是矩阵乘、逐行 softmax 和链式法则;这里补齐指数换参考系与 softmax 雅可比,不需要先懂 CUDA。本文先讨论有限分数、至少一个未遮 key、无 dropout 的单头自注意力,$Q,K,V\in\mathbb R^{N\times d}$,$S=QK^\top/\sqrt d$;d 为头宽度。直接计算 $\exp(s)$ 会溢出,fp32 的阈值约为 $\log(\mathrm{finfo}(\mathrm{float32}).\mathrm{max})\approx88.72$。常用 safe softmax——先求这行的最大值 m,再算平移后的指数: $$p_i = \frac{\exp(s_i - m)}{\sum_{j=1}^{N} \exp(s_j - m)}, \qquad m = \max_{1 \le j \le N} s_j$$ 把指数共同平移 -m(数学上等价于分子分母同乘 $\exp(-m)$,实现不真的计算这个乘子),结果不变,但指数的输入全部落在 $(-\infty, 0]$,永不溢出。问题就出在这个 m 上:m 是对整行取的 max。你必须先把 N 个分数全部看过一遍才知道 m 是多少,然后才能开始算 exp——这解释了常见实现先保存分数再做 softmax 的流程;算法并不强制保存整个 N×N,也可重算分数或逐行处理,只是 IO 和效率不同。 而注意力输出对这一行还要再多两个量:分母 $l = \sum_j \exp(s_j - m)$,以及未归一化加权和 $U = \sum_j \exp(s_j - m)\, v_j$($v_j$ 是第 j 个 token 的 Value 向量)。最终输出 O 就是 $U / l$。 所以真正要回答的问题是:如果分数是一块一块到来的(事先不知道后面块里有什么),这三个量还能算吗? 3.2 online softmax 递推式(全文核心) 能。做法是把「以 m 为参考系」改成「以当前的 m 为参考系,m 变了就整体换算」。 设 $J_t$ 为前 t 块中未被遮住的 key 下标集合,至少含一个有限分数;维护三个量:参考系最大值 $m^{(t)}$、分母 $l^{(t)}$、未归一化加权和 $U^{(t)}$,它们满足不变式 $$U^{(t)} = \sum_{j \in J_t} \exp(s_j - m^{(t)})\, v_j, \qquad l^{(t)} = \sum_{j \in J_t} \exp(s_j - m^{(t)})$$ U 表示未归一化加权和,最终输出才记作 O;这样不把中途累加器与归一化结果混为同一量。现在第 $t+1$ 块到了,块内最大值是 $m_{\text{blk}}$。新的全局最大值是 $$m^{(t+1)} = \max(m^{(t)},\ m_{\text{blk}})$$ 关键一步来了:旧累积量是按 $m^{(t)}$ 为参考系记的,而新的不变式要求参考系换成 $m^{(t+1)}$。把不变式里的 $\exp(s_j - m^{(t)})$ 拆成 $\exp(s_j - m^{(t+1)}) \cdot \exp(m^{(t+1)} - m^{(t)})$,旧量的换算系数就是 $\exp(m^{(t)} - m^{(t+1)})$: $$l^{(t+1)} = \exp(m^{(t)} - m^{(t+1)})\, l^{(t)} + \sum_{j \in \text{blk}} \exp(s_j - m^{(t+1)})$$ $$U^{(t+1)} = \exp(m^{(t)} - m^{(t+1)})\, U^{(t)} + \sum_{j \in \text{blk}} \exp(s_j - m^{(t+1)})\, v_j$$ 每一步只是把指数拆成两项相乘再重新合并,等价性是代入即可验证的恒等式;所有块流完后 $U^{(T)}/l^{(T)}$ 与朴素 softmax 在实数算术下相同(差在浮点舍入,第 04 节实测 1e-16 量级)。这个「换参考系」的系数在论文和代码里叫 rescale,跨块最大值、求和与 rescale 都会带来额外的标量操作;它们不改变主导矩阵乘次数。 严谨一点可以正向验证不变式:假设第 t 步的不变式成立,那么 $$U^{(t+1)} = e^{m^{(t)} - m^{(t+1)}} \sum_{j \in J_t} e^{s_j - m^{(t)}} v_j + \sum_{j \in \text{blk}} e^{s_j - m^{(t+1)}} v_j = \sum_{j \in J_{t+1}} e^{s_j - m^{(t+1)}} v_j$$ (第一个等号就是递推式,第二个等号把 $e^{m^{(t)} - m^{(t+1)}}$ 乘进求和号里、指数相加后正好变回 $e^{s_j - m^{(t+1)}}$。)旧块和新块在同一个参考系下合并,不变式保持。$l$ 的证明一字不差,把 $v_j$ 去掉就行。归纳基础是初始状态 $m^{(0)} = -\infty$、$l^{(0)} = 0$、$U^{(0)} = 0$:第一个含有效 key 的块到来时换算系数按 0 处理(代码只在旧最大值有限的位置做减法,避免提前计算未定义差值),三个量直接等于该块的局部值。 三个细节值得停一下: 这个递推对任意分块都成立,块大小可以是 1(一个 key 一个 key 地流),也可以是 128 行。块大小只影响效率,不影响结果。 $\exp(m^{(t)} - m^{(t+1)}) \le 1$ 恒成立(因为 $m$ 单调不减),rescale 永远是在把旧量缩小,不会因正指数而溢出,但有限精度下可能下溢到零、求和也有舍入误差;点积溢出必须另外处理。 $m^{(t)} + \log l^{(t)}$ 就是 logsumexp(LSE)。它是反向传播唯一需要额外记录的东西——记住这个,第 3.3 节要用。 掩码把被遮分数设为 $-\infty$。若此前最大值有限,它的指数贡献为零。若一整行截至当前全被遮住,不能直接算 $-\infty-(-\infty)$;本文约定该行累计量为零,最终全遮行返回零输出、LSE 为 $-\infty$。这是接口约定,全遮行本身没有概率归一化定义。实现用带 where= 的 np.subtract 跳过未定义减法;普通 np.where 会先求值两个分支,不能保证避免 NaN 警告。更进一步,如果一整块都被遮住(Q 块整体在 K 块之前),这块连算都不用算,直接跳过。第 04 节实测这个「整块跳过」在 N=4096、块 64 时省掉 49.2% 的块。 3.3 反向传播:重算换显存 这里的 $O=PV$ 指最终归一化输出,$G=\partial L/\partial O\in\mathbb R^{N\times d}$ 是上游梯度,$A=GV^\top=\partial L/\partial P\in\mathbb R^{N\times N}$。先对一行 softmax 求导,$\delta_{jk}$ 表示 j=k 时为 1、否则为 0: $$\frac{\partial p_j}{\partial s_k}=p_j(\delta_{jk}-p_k)$$ 把损失对各 $p_j$ 的导数 $A_j$ 代入链式法则,得到 $$\frac{\partial L}{\partial s_k}=\sum_j A_jp_j(\delta_{jk}-p_k)=p_k\left(A_k-\sum_j A_jp_j\right)$$ 因此矩阵形式为 $$dS=P\odot\left(A-\operatorname{rowsum}(A\odot P)\right)$$ 行和是 $N\times1$,沿 key 轴广播。继续对 $S=QK^\top/\sqrt d$ 与 $O=PV$ 求导: $$dQ=dS\,K/\sqrt d,\qquad dK=dS^\top Q/\sqrt d,\qquad dV=P^\top G$$ 反向真正巧的一步是先把行和算出来。对第 i 行,用内积线性性可得 $$\delta_i=\sum_j P_{ij}\langle G_i,V_j\rangle=\left\langle G_i,\sum_j P_{ij}V_j\right\rangle=\langle G_i,O_i\rangle$$ 所以 $\delta=\operatorname{rowsum}(G\odot O)$,仅为 N 个数,无需为这项先构造完整 A/P。前向额外保存 O 和 $\mathrm{LSE}_i=m_i+\log l_i$;反向从 Q/K 重算局部分数,再用 $$P_{ij}=\exp(S_{ij}-\mathrm{LSE}_i)$$ 恢复局部概率,把它和局部 $G_iV_j^\top-\delta_i$ 相乘得到 dS,立即累加到 dQ/dK/dV 后丢弃。Q/K/V、输出和梯度仍是 $O(Nd)$,局部 tile 为 $O(B_rB_c)$;只有在 d 与块大小固定时才能简称 O(N)。原论文附录 B 给出分块反向;附录 attention_backward 把上述步骤完整实现,第 4.6 节实际核验梯度。 以上推导无 dropout。若训练使用概率 dropout,令随机乘子为 $R/(1-p)$,p 为丢弃率;dV 和 dP 都要使用同一随机掩码重放,不能直接照抄本教学函数。LSE 仍对应 dropout 之前的 softmax。生产实现还需保留随机生成器状态,见第 05 节。 用重复计算换显存——这笔交易的换算率是:每层每个头多算一遍 QK^T 和一次 exp(具体比例取决于前后向统计口径与实现),对本文保守教学账本,这对应去掉 S、P、浮点 dropout 乘子三项;真实朴素反向通常不用同时保存 S,具体减少几份取决于实现。N=4096、本文简化模型上,单层激活 2.39 GiB → 0.141 GiB,17 倍。 3.4 IO 账:搬运量到底差多少 设单头维度 d,片上存储预算为 M 个元素(不是字节);fp16 下 192 KiB 对应 M=98304。SRAM 里要同时放下 K 块、V 块(各 $B_c \times d$)和 Q 块、输出块(各 $B_r \times d$),论文 Algorithm 1 取(脚本对整数块大小向下取整) $$B_c = \frac{M}{4d}, \qquad B_r = \min(B_c,\ d)$$ NVIDIA Ampere 文档 给出 A100 每个 SM 的 192 KiB 统一 L1/共享内存空间及共享内存配额;它不是全部可用于单个 block 的任意数组。以这一数量级作为教学预算,fp16 下 d=64 时 $B_c = 384$、$B_r = 64$——这是论文 IO 模型的粗略分块预算;真实 kernel 还受 score tile、累加器、寄存器、共享内存配额与 occupancy 约束,192 KiB 也不是每个 block 可独占的共享内存。 朴素实现的搬运量(单位:元素个数):QK^T 读 Q、K 各 Nd、写 S 一次 N²;softmax 读 S 写 P 各 N²;PV 读 P 一次 N²、读 V 写 O 各 Nd。合计 $4Nd + 4N^2$,主导项 4N²。 FlashAttention 的搬运量:K、V 各进 SRAM 一次($2Nd$);Q 块每换一个 K/V 块就要重读一遍,共 $T_c \cdot Nd$($T_c = \lceil N/B_c\rceil$ 是 K/V 块数);输出 O 同理要读出写回各一遍($2 T_c Nd$);l、m 两个 O(N) 的运行量共 $4T_c N$。合计约 $2Nd + 3T_cNd + 4T_cN$,代进去: $$\text{HBM 搬运量} \;\approx\; 3 \cdot \frac{N}{B_c} \cdot Nd \;=\; \Theta\!\left(\frac{N^2 d^2}{M}\right)$$ 两个量级一比:朴素是 $\Theta(N^2)$,Flash 是 $\Theta(N^2 d^2/M)$,比值 $d^2/M$ 在 d=64、fp16、M=98304 个元素(192 KiB) 时约等于 0.04——大 O 比例省略了常数,不能直接当作 25 倍的实际收益;脚本计入 Q/O 反复读写后的模型比值为 7.3 倍。原论文 Theorem 2 的前提为 $d\le M\le Nd$;其中线性项可被主导项吸收。若整份输入已能驻留片上,不能继续把 $N^2d^2/M$ 往下外推到低于读写输入输出的 $\Omega(Nd)$。小 N 时仍应保留线性项与取整。本文计数也忽略初始化、最终归一化、LSE 写回及多级缓存,提供对照账本而非 GPU 字节计数器。 一个预算反例:d=64、Bc=384、Br=64,若 Q/K/V 用 2 字节,而输出累加器、分数 tile 和统计量用 4 字节,仅这些数组就需 221696 字节,超过 192 KiB;这还未算额外概率 tile。因而同类型元素模型推得的 384 不能作为实际混合精度 kernel 的可用配置。真实实现会重用存储、分配寄存器并选不同 tile。 按本文简化模型,朴素注意力的算术强度在 N 远大于 d 时约为 d/2 FLOP/byte(fp16),所以固定 d=64 时接近 32;它会随 d 改变,并非与头维度无关。A100 示例的约 200 FLOP/byte 是 NVIDIA A100 白皮书 中 dense FP16 Tensor Core 峰值312 TFLOP/s与HBM带宽1555 GB/s之比,真实 softmax 的非矩阵乘指令、缓存与调度仍会影响性能。 这张图要看什么:固定 d 与 SRAM 大小时,左图两者的大 N 主导项均为 N²,比值渐近趋于常数;线性项和取整会让有限 N 的比值变化;右图斜率不同(N² 对 N),N 越大两条线离得越远,朴素教学激活曲线越过示例 80 GiB 预算线的位置,FlashAttention 还有几百倍余量。 04. 代码实现 完整脚本在文末附录(flash_online_softmax.py、io_ledger.py、attention_checks.py、make_figures.py)。前三项数值/账本实验主要用 NumPy,梯度对拍还需 PyTorch,配图需 matplotlib。本次 2026-10-14 在 Windows、Python 3.11.9、NumPy 2.4.6、PyTorch 2.10.0+cpu、matplotlib 3.10.7 实跑;BLAS/OMP 线程设为 1。下面表格来自同一次完整运行,不是 GPU 性能复现。 4.1 递推式长什么样:逐块打印 先看一个能盯着看的例子(N=8、D=4、K/V 块大小 B_c=2,Q 整行一起处理): # 局部Sij已设置掩码;m_old/l/O是已见块的累计量。 m_new = np.maximum(m_old, Sij.max(axis=-1)) shifted = np.full_like(Sij, -np.inf) np.subtract(Sij, m_new[:, None], out=shifted, where=np.isfinite(m_new)[:, None]) Pij = np.exp(shifted) delta = np.full_like(m_old, -np.inf) np.subtract(m_old, m_new, out=delta, where=np.isfinite(m_old)) corr = np.exp(delta) l[i0:i1] = l[i0:i1] * corr + Pij.sum(axis=-1) O[i0:i1] = O[i0:i1] * corr[:, None] + Pij @ Vj m[i0:i1] = m_new 和第 3.2 节逐符号对上:m_new 是 $m^{(t+1)}$,corr 是换参考系系数,Pij.sum 是块内指数和,Pij @ Vj 是块内加权和。代码内部 O 保存的是数学记号 U,函数最后才除以 l 得到归一化输出。全遮行不做除法,保持零输出。函数要求相同非空形状 [N,D] 与有限浮点输入;支持尾块、一维 key_valid 和等长因果掩码。float16 输入提升到 float32 累加并输出 float32;float64 保持 float64。它未实现非等长、GQA 或 dropout,不能当作完整 CUDA API 替代品。 实跑的逐块轨迹(第 0 号 query): j= 0 m=+0.2968 l=1.7748 corr=0.0000 j= 2 m=+0.6071 l=3.0317 corr=0.7333 j= 4 m=+0.6071 l=3.9732 corr=1.0000 j= 6 m=+0.6071 l=5.3007 corr=1.0000 看两点:j=2 时新块里出现了更大的分数,m 被抬高、旧累积量被打了个 0.733 的折扣;j=4 之后 m 没再变,corr 恒为 1,rescale 白做——此时乘子在数学上为 1;具体 GPU kernel 是否跳过这些操作要核对实现,不能仅由轨迹推断。最大的单个临时矩阵只有 16 个元素([8, 2]),而 N×N 是 64;这里不等于同时存活临时数组的总元素数;代码对尾块和全遮行的额外检查也占临时空间。 4.2 等价性:和朴素实现在实数算术下等价,浮点下允许舍入差异 attention_naive 是上一篇的标准实现(S、P 都落地),两者对拍: N D causal 最大绝对误差 相对误差 256 64 False 5.551e-16 1.192e-15 256 64 True 8.882e-16 3.520e-16 1024 64 False 2.082e-16 8.962e-16 1024 64 True 6.661e-16 2.332e-16 2048 64 False 3.053e-16 1.205e-15 2048 64 True 6.661e-16 2.856e-16 4096 64 False 1.249e-16 9.209e-16 4096 64 True 7.772e-16 2.426e-16 本次表的最大绝对误差为 1.2e-16 到 8.9e-16;这是此输入、规模与 float64 运算顺序下的结果。这就是「精确注意力」四个字的实测含义:不是「误差很小」,是算法本身和朴素 softmax 完全等价。 4.3 峰值内存:O(N²) 对 O(N),实测 用 tracemalloc 量函数内新分配的峰值内存(float64、D=64、块 64×64;Q/K/V 已提前分配,不计入本表): N naive 实测 naive 理论 2N²·8 flash 实测 比值 1024 16.5 MiB 16.0 MiB 1.2 MiB 13.9x 2048 65.0 MiB 64.0 MiB 2.2 MiB 29.4x 4096 258.0 MiB 256.0 MiB 4.3 MiB 60.0x 8192 1028.0 MiB 1024.0 MiB 8.5 MiB 121.0x 理论列只计 S/P;tracemalloc 表还包含输出、运行量和局部临时数组。NumPy 在此环境向 Python 跟踪器登记数组分配,这些数值与主导项一致;tracemalloc 不等于进程 RSS,也不覆盖所有 BLAS 内部缓冲,更不是 GPU 显存峰值。最后那行是阶数的直接证据:N 翻倍,naive 峰值 ×4(二次),flash 峰值 ×2(线性)。 4.4 时间:CPU 教学循环的实测口径 N naive flash(B=64) flash/naive 1024 13.6 ms 16.2 ms 1.2x 2048 57.5 ms 65.4 ms 1.1x 4096 230.0 ms 266.1 ms 1.2x 每个方法先预热一次,再量 5 次取墙钟中位数,包含 Python 循环与 NumPy分配,不计输入生成。这里 flash 教学循环慢约 1.1–1.2 倍;不同机器或本机再次运行可能变化。可能原因包括:乘加次数一样还多了 rescale 和逐块 exp;一次大矩阵乘被拆成 (N/64)² 个 64×64 小矩阵乘,BLAS 跑不满小块,Python 循环开销也进来了;CPU 也有 SRAM 缓存,但这里的 Python/NumPy 分块没有实现专门的缓存与线程优化,不能把它当作 GPU 内核速度的预测。第三张图展示 IO 模型为何支持在 GPU 上尝试这一优化。 4.5 因果掩码:跳块比例与真实算力要分开 本脚本朴素路径先计算全分数,再遮住未来 key,因此该路径不因掩码节省矩阵乘;这是所测实现的选择,不是所有精确算法的限制。FlashAttention 按「Q 块整体在 K 块之前就整块跳过」处理,N=4096、块 64 时实测: 块总数(非因果) : 4096 块总数(因果跳过): 2080 跳过的块占比 : 49.2% 保留的是含对角线的下三角,块数是 $T(T+1)/2$,占比 $(T+1)/2T$,T 大时趋近一半。对齐、等长且等块时,49.2% 是跳过的 tile 数比例;对角 tile 中仍有被遮元素,softmax/调度还有开销,所以不等于总 FLOPs 或延迟精确减半。时序因果视频注意力还须明确 token 排序与掩码定义。 4.6 反向与边界:实际梯度对拍 attention_checks.py 对 N=1/7/17、float64/float32/float16 输入、因果/非因果、正常/全遮及部分遮 key、三组包含尾块的块大小组合,总计 108 组前向结果与 PyTorch 2.10 的 CPU math SDPA(float64 参考)对拍;最大绝对差为 2.33e-7。float16 输入先提升到 float32,并非模拟 GPU 的每一处精度行为。float64 分块反向对 autograd 最大差为 1.33e-15。另对 N=7、d=5 的 Q/K/V 共 105 个参数做中心差分,步长 1e-5,最大差为 8.74e-11;全遮行零输出约定也参与检验。 还实际提取并执行 FlashAttention v2.8.3 官方测试的 attention_ref 与 construct_local_mask 原函数 AST,CPU float64、不做默认 float32 upcast,对三种长度及两种因果设置核对输出与梯度。这是原始参考函数的检查,未安装或执行 CUDA kernel,不能据此声称 GPU 速度/低精度训练已复现。独立对拍支持实现正确性;第 03 节的恒等式才说明算法为何等价。 05. 工业级实现对照 最小实现讲清了原理,但生产 kernel 和它有四处本质差异,每处都值得知道为什么: 第一,理论 tile 和工程配置要分清。 本次固定核对 flash-attention v2.8.3,入口为 flash_attn_func。flash_fwd_launch_template.h 中有基于头维度、架构、dropout/因果等条件的编译配置分支;Python _get_block_size_n 也说明应匹配 CUDA 块尺寸。静态分支选配置不等于每次运行都在线 autotune。$B_c=M/(4d)$ 只为 IO 分析提供同阶预算,实际还要计 score tile、寄存器、共享内存和 occupancy。本文不判断更新版本的 kernel 细节。 第二,减少非矩阵乘操作。 FlashAttention-2 使用未归一化输出累计,减少 rescale、除法等非矩阵乘工作。corr=1 时数学上无需改变旧值,但不能笼统声称所有 kernel 都按每行最大值是否变化来分支跳过。 第三,前向的结构是「外层 Q、内层 K/V」。 我们按论文 v1 的写法外层遍历 K/V 块;FlashAttention-2 把循环反过来(外层 Q 块),好处是输出 O 常驻寄存器不用反复读写、且不同 Q 块之间天然并行,能吃满更多 SM。论文 v1 的伪代码适合理解递推,v2 的循环结构才是现在 kernel 的样子。 第四,dropout 要重放随机过程。 固定接口保存 rng_state,反向从同一状态重建概率 dropout;flash_bwd_kernel.h 中可核对 Philox seed/offset 与块坐标。仅口头说“种子一样”仍不够:并行划分与计数器映射也必须一致。无需保存完整二次掩码,是特定 kernel 的实现能力,不代表所有朴素框架必定保存浮点掩码。 另一个易错边界是非等长因果对齐。 v2.8.3 接口明确对齐右下角;Nq=5、Nk=2 的保留矩阵为下面五行,前 3 行全遮时输出零。本文 NumPy函数只支持等长,因此不暗中把一种对齐当作另一种。PyTorch 2.10 SDPA 文档 的 is_causal 在非方形时定义左上对齐;跨 API 迁移要显式核对。 0 0 0 0 0 0 1 0 1 1 PyTorch SDPA 可以按环境选优化后端,也有 CPU/math 路径;调用该 API 本身不能证明使用了 CUDA FlashAttention。评估时还要把 dropout_p 显式设为 0,该函数不会因模型切换 eval 就自动关闭传入的 dropout。 原 FlashAttention 论文分别在 BERT-large/MLPerf、GPT-2/HuggingFace、Long Range Arena 上报告 15%、3 倍、2.4 倍端到端收益。这些模型、基线与任务不同,不能把三组数字当作只改变 N 的对照实验,更不能由它们证明收益严格随 N 单调增加。 排查问题时你会想知道「此刻到底在用哪个后端」。PyTorch 留了一个官方口子: import torch from torch.nn import functional as F from torch.nn.attention import sdpa_kernel, SDPBackend # CUDA环境使用;形状[B,H,N,D],输入half,关闭dropout。 q = k = v = torch.randn(1, 2, 128, 64, device="cuda", dtype=torch.float16) with sdpa_kernel(SDPBackend.FLASH_ATTENTION): out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0) # 强制走 flash;如果这个头维度/掩码组合它不支持,应按官方后端约束排查告警或错误;本文CPU实跑明确选择MATH, # 而不是悄悄退回 math 后端——「悄悄降级」正是性能莫名掉一半时最该先查的事 这张图要看什么:横轴是 N,纵轴是「每搬 1 字节做多少次运算」,灰色虚线是机器平衡点(201 FLOP/byte)。本模型中朴素实现的强度渐近接近 32;FlashAttention 的估计从 N=2048 起超过平衡点。这是按矩阵乘峰值做的模型分类,真实瓶颈还受非矩阵乘指令、缓存、并行度和调度影响。 06. 代价与边界 FlashAttention 省下了 HBM 搬运和 N² 显存,赔进去的和没管住的也要说清楚。 代价:重计算和额外归一化操作。 反向要重算局部分数及概率,其中包含矩阵乘和逐元素运算,不能统称为固定 30% 的额外非矩阵乘开销。小序列的收益取决于 kernel、调度和硬件,没有统一的“序列短于512就亏损”阈值。 数值边界:实数算术等价不保证浮点逐位一致。 kernel 内部用 fp16/bf16 存储、fp32 累加,块内的归一化和朴素实现的一次性归一化在浮点上不同。对训练的影响需要结合 dtype、输入尺度与误差测试判断,但如果你在做数值敏感的分析(比如逐 token 概率对比),应建立与输入尺度、dtype 和应用需求匹配的容差,并检查遮罩及非有限值;不能仅凭“精确算法”就把所有数值差异判为正常舍入。 边界的核心一条:它改变存储阶数,但不消除精确稠密注意力的二次计算。 显存从 $O(N^2)$ 降到 $O(N)$,但算力还是 $\Theta(N^2 d)$、在第 3.4 节限定模型中,HBM 搬运为 $\Theta(N^2 d^2/M)$。N=32768 时 FlashAttention 的单头搬运是 1061 MiB——比朴素实现的 8.2 GiB 好得多,但随 N 继续平方增长这一点没变。上下文再往上涨(1M token),可以考虑稀疏注意力、线性注意力、状态空间模型或并行方案;它们各有建模或系统代价,是否合适须按任务验证。FlashAttention 的块结构恰恰是它们的底座:把注意力切成块之后,「整块跳过」才成为可能,第 4.5 节那个 49.2% 推广到任意稀疏模式就是块稀疏注意力。 不该用的场景:需要拿到完整注意力权重做分析或可视化的(P 从头到尾没存在过,想看它就得回到朴素实现);自定义的任意注意力偏置如果 kernel 不支持,绕过去的方法可能把优势吃掉;以及缺乏适配内核的环境;CPU 也有 SRAM 缓存,但本文 NumPy 循环没有实现专门的 CPU cache 优化(第 4.4 节的 CPU 实测就是例子)。 07. 经典论文脉络 这条线的演进关系一句话各说清: Milakov & Gimelshein, 2018(arXiv:1805.02867)Online normalizer calculation for softmax——给出 softmax 的 online normalizer 算法:流式更新 max 和指数和。从原文可核对其减少内存访问的 softmax/TopK 实验;不从这篇论文推断此前所有研究的范围。 Rabe & Staats, 2021(arXiv:2112.05682)Self-attention Does Not Need O(n²) Memory——讨论单 query 的常数额外空间方案,并给出分块的低内存自注意力实现及速度/内存实验;不能把单 query 的额外空间结论写成整段输出总存储 O(1),也不能概括为没有实用价值。 Dao et al., 2022(arXiv:2205.14135)FlashAttention——本篇锚点。补上缺失的一环:把流式更新从「逐 token」改成「逐块」(tiling),配上 IO 复杂度分析和 GPU kernel,报告精确注意力在其测试 GPU/模型上的速度和内存收益。 Dao, 2023(arXiv:2307.08691)FlashAttention-2——循环重排(外层 Q)、削减非矩阵乘 FLOPs、更好的并行度,把 v1 大约 25-40% 的峰值算力利用率推向 50-73%。 这条阅读线从流式归一化、低内存自注意力到 IO 与 GPU 调度,逐步解释同一数学操作为何能有不同的存储和运行代价。它是选定论文的演进关系,不是对所有先行工作或产品的首次性判断。 08. 常见误解 以下几条都值得单独记住,前两条我当初也信过: 「FlashAttention 是近似注意力,所以有精度损失」。错。它是精确算法,和朴素 softmax 数学等价(第 4.2 节实测 1e-16)。真正近似的是 Linformer、Performer 那一族。两者经常被并列讨论,但一个在改算法,一个在改数据的搬运方式。 「它快是因为算了更少的 FLOPs」。反了,它的 FLOPs 略多于朴素实现(第 4.4 节本次CPU教学循环约慢 1.1–1.2 倍)。主要收益来自减少 IO,也来自更好的并行划分与调度——这正是它给所有做系统优化的人的启示:先问数据在哪,再问算了多少。 「显存优化只对前向有用」。训练反向同样受益:除 Q/K/V 外额外存 O 和 LSE(固定 d 时为 O(N)),反向重算 P。本文三份二次张量是教学分配假设,真实反向不一定保存三份;可靠结论是分块重算避免显式保留完整 P。 「块开得越大越快」。块大小受 SRAM 硬约束:超过每个 block 的资源上限可能无法启动或编译;寄存器 spill 等情况也会引入额外访存,整个方法的根基(中间结果不落 HBM)就塌了。真实 kernel 的块大小是约束内的调优问题,不是越大越好。 「有了它就不用再关心 N²」。它改变了激活存储阶数;固定 d/M、在本文模型内,稠密计算与搬运的主导项仍随 N 平方增长。能继续走的长上下文路线是稀疏化(改算力阶数)和状态空间/线性注意力(改注意力本身),FlashAttention 的块结构是它们的载体,不是替代品。 09. 动手验证 跑文末附录的 flash_online_softmax.py(只依赖 numpy),预期输出与正文一致: 第 2 节等价性表:本次非因果最大误差 1.2e-16 到 5.6e-16,因果 6.7e-16 到 8.9e-16——全是舍入级别; 第 3 节内存表:naive 峰值和理论值 $2N^2 \times 8$ 字节吻合,N 翻倍时 naive ×4.0、flash ×2.0。 再运行附录 attention_checks.py,它用 seed=1024、N=1024、d=64、块 64×64,故意把正确函数的 corr = np.exp(delta) 替换成 corr = np.ones_like(delta);这是自动构造的错误反例。最大差实测为 0.275658、平均差为 0.013599。这项对比把“丢掉换参考系会出错”变成可重复实验,不让读者误把旧代码里的 np.where 当作必要写法。 最后改 io_ledger.py 开头的硬件常数(比如把 SRAM_PER_SM 调到 48 KiB,作为更小片上预算的假设,不能据此代表某款消费级卡),重跑看块大小和搬运量比值怎么变——你会看到 SRAM 越小,FlashAttention 相对朴素实现的搬运量优势越小,$d^2/M$ 里的 M 直接控制这一切。 10. 延伸阅读 按知识树的依赖关系,建议按这个顺序继续走: 前置:自注意力机制的计算与显存账本——本篇所有显存账本的出处(2.39 GiB、94% 那几笔账都在那篇里);旋转位置编码 RoPE 与 视频 DiT 里的 3D RoPE——注意力的另外两个必备零件,本篇刻意没有碰位置编码。 后继:视频生成里的稀疏注意力——FlashAttention 的块结构是稀疏模式的执行底座,「整块跳过」从因果掩码推广到任意稀疏图;算子融合与 CUDA Graph——把「少搬内存」推到极端就是融合,FlashAttention 是这个思路最成功的案例;KV Cache 与自回归视频生成——推理时的另一本显存账。 延伸到知识树之外:想读 kernel 源码,从 Dao-AILab/flash-attention 的 flash_attn/flash_attn_interface.py 进,先读 forward 再读 backward;想读原始推导,Milakov & Gimelshein(1805.02867)给出了独立 softmax 在线归一化的直接推导。 附录:完整代码 09 节用到的脚本全文如下(io_ledger.py、flash_online_softmax.py、attention_checks.py、make_figures.py)。复制到本地存成同名文件,按各脚本开头的依赖说明准备环境后即可运行。 io_ledger.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """IO 账本:朴素注意力和 FlashAttention 在 HBM 上到底搬了多少字节。 这篇讲的是「为什么快」。答案不在 FLOPs 上——两者的乘加次数几乎一样—— 而在于 GPU 有两级内存: HBM(显存) 带宽约 1.5 TB/s,容量 40~80 GB SRAM(片上) 带宽约 19 TB/s,但每个 SM 只有 192 KB 朴素实现把 N×N 的分数矩阵写回 HBM 再读出来,等于把数据在这条 13 倍带宽差的 通道上来回搬;FlashAttention 用 tiling 让这些中间结果根本不落 HBM。 下面的模型把每一笔读写都点清楚,参数是可调的,读者可以改硬件常数重算。 硬件数量级取自 FlashAttention 论文第2节与NVIDIA A100白皮书(40GB、非稀疏Tensor Core峰值)。 只依赖 numpy。直接 `python io_ledger.py` 即可运行。 """ import numpy as np # ── 硬件常数(A100 40GB 量级,论文第2节/NVIDIA白皮书)──────────────────── HBM_BW = 1.555e12 # HBM 带宽,字节/秒 SRAM_BW = 19.0e12 # 片上 SRAM 带宽,字节/秒 SRAM_PER_SM = 192 * 1024 # 每个 SM 的片上 SRAM,字节 FLOPS_PEAK = 312e12 # fp16 tensor core 峰值,FLOP/秒 GIB = 1024 ** 3 MIB = 1024 ** 2 # ──────────────────────────────────────────────────────────── # 块大小:SRAM 里能同时放下什么 # ──────────────────────────────────────────────────────────── def block_sizes(d, sram_bytes=SRAM_PER_SM, dtype_bytes=2): """返回 (B_r, B_c):Q 块行数与 K/V 块行数。 SRAM 里要同时放下 K_j、V_j 两个 [B_c, d] 和 Q_i、O_i 两个 [B_r, d]。 论文的取法是 B_c = M / (4d)、B_r = min(B_c, d),这里取整近似: 先让 K_j+V_j 占掉一半 SRAM,Q 块则不超过 d 行(与分数tile的存储预算同阶)。 """ if d <= 0 or dtype_bytes <= 0 or sram_bytes < 4*d*dtype_bytes: raise ValueError("存储预算不足或参数非正") elems = sram_bytes / dtype_bytes B_c = max(1, int(elems // (4 * d))) B_r = min(B_c, d) return B_r, B_c # ──────────────────────────────────────────────────────────── # HBM 读写量(单位:元素个数,乘 dtype_bytes 得字节) # ──────────────────────────────────────────────────────────── def hbm_elems_naive(N, d): """朴素实现:S 和 P 都要落地。 QK^T: 读 Q(Nd) + 读 K(Nd) + 写 S(N²) softmax: 读 S(N²) + 写 P(N²) PV: 读 P(N²) + 读 V(Nd) + 写 O(Nd) """ return 4 * N * d + 4 * N * N def hbm_elems_flash(N, d, B_r, B_c): """FlashAttention:外层遍历 K/V 块,内层遍历 Q 块。 K、V 各读一遍(每个 j 块进 SRAM 后,内层 i 循环里一直复用) Q 每个 j 都要重读一遍:T_c · N·d O 每个 (j,i) 都要读出来再写回去(累加器跨 j 迭代):2 · T_c · N·d l、m 两个 O(N) 的运行量同理:4 · T_c · N """ T_c = int(np.ceil(N / B_c)) return 2 * N * d + 3 * T_c * N * d + 4 * T_c * N def flops_attention(N, d): """两个 N×N 矩阵乘,一次 [M,K]×[K,N] 算 2MKN 个浮点运算。""" return 4 * N * N * d def activation_bytes(N, d, H, B=1, dtype_bytes=2): """单层注意力的训练激活(反向要用的中间张量)。 保守教学模型:6 个 [B,N,d_model] 线性项 + S/P/浮点乘子 3 个 [B,H,N,N]。 此函数 d 是总宽度 d_model,不是其他 IO 函数中的头宽度。 融合侧保留同样六份线性预算,忽略小的 LSE 与工作区;不是框架峰值测量。 """ lin = 6 * B * N * d * dtype_bytes quad = 3 * B * H * N * N * dtype_bytes return lin + quad, lin def roofline(bytes_moved, flops): """算术强度(FLOP/byte)与两个理想下界时间。返回 (强度, 内存时间, 算力时间, 瓶颈)。""" intensity = flops / bytes_moved t_mem = bytes_moved / HBM_BW t_comp = flops / FLOPS_PEAK bound = "内存受限" if t_mem > t_comp else "算力受限" return intensity, t_mem, t_comp, bound # ──────────────────────────────────────────────────────────── # 报表 # ──────────────────────────────────────────────────────────── def report_io(): d, b = 64, 2 # 单头维度 64,fp16 B_r, B_c = block_sizes(d) print("=" * 74) print("1. SRAM 块大小与 HBM 读写量(单头,d=64,fp16)") print("=" * 74) print(f" SRAM {SRAM_PER_SM / 1024:.0f} KB / SM,fp16 下能放 " f"{SRAM_PER_SM / b:.0f} 个元素") print(f" → B_c = M/(4d) = {B_c},B_r = min(B_c, d) = {B_r}\n") print(f"{'N':>7} {'naive HBM':>12} {'flash HBM':>12} {'比值':>8} " f"{'naive 强度':>11} {'flash 强度':>11}") rows = [] for N in (1024, 2048, 4096, 8192, 16384, 32768): nb = hbm_elems_naive(N, d) * b fb = hbm_elems_flash(N, d, B_r, B_c) * b fl = flops_attention(N, d) i_n, _, _, bound_n = roofline(nb, fl) i_f, _, _, bound_f = roofline(fb, fl) rows.append((N, nb, fb, i_n, i_f, bound_n, bound_f)) print(f"{N:>7} {nb / MIB:>10.1f} MiB {fb / MIB:>10.1f} MiB " f"{nb / fb:>7.1f}x {i_n:>9.1f} {i_f:>9.1f} ") print(f"\n 机器平衡点(峰值算力/带宽)= {FLOPS_PEAK / HBM_BW:.0f} FLOP/byte") print(" 仅本理想模型:强度低于平衡点由HBM项主导,高于时由计算项主导。\n") print(f"{'N':>7} {'naive 瓶颈':>12} {'flash 瓶颈':>12}") for N, nb, fb, i_n, i_f, bn, bf in rows: print(f"{N:>7} {bn:>12} {bf:>12}") print("\n 注意:flash 的强度在 N 大时越过平衡点,模型说它变成算力受限了。") print(" 但真实 kernel 达不到峰值——softmax 的 exp 走的是特殊函数单元,") print(" 不走 tensor core,这个「非矩阵乘开销」正是 FlashAttention-2 之后") print(" 继续优化的地方。峰值资源模型给出理想时间下界,不是实测速度承诺。") def report_activation(): d, H, B, b = 3072, 24, 1, 2 print("\n" + "=" * 74) print("2. 教学激活存储模型(简化 Transformer:d=3072, H=24, B=1, 32 层, bf16)") print("=" * 74) print(f"{'N':>7} {'朴素/层':>12} {'Flash/层':>12} {'比值':>8} " f"{'朴素 32 层':>12} {'Flash 32 层':>13}") for N in (1024, 4096, 8192, 32768): naive, flash = activation_bytes(N, d, H, B, b) print(f"{N:>7} {naive / GIB:>10.2f} GiB {flash / GIB:>10.3f} GiB " f"{naive / flash:>7.0f}x {32 * naive / GIB:>10.1f} GiB " f"{32 * flash / GIB:>11.2f} GiB") n4096, f4096 = activation_bytes(4096, d, H, B, b) print(f"\n N=4096 时,朴素实现单层 {n4096 / GIB:.2f} GiB,其中二次项占 " f"{100 * (1 - f4096 / n4096):.1f}%——") print(" 这正是 attention_basics 那篇里「三个 N×N 吃掉 94%」的那一格。") print(" FlashAttention 删掉的就是这一格,剩下的是 O(N) 的线性项。") def report_speed_limit(): d, b = 64, 2 B_r, B_c = block_sizes(d) print("\n" + "=" * 74) print("3. 理想时间下界:如果只受 HBM 带宽限制,两者各要多久(单头,d=64)") print("=" * 74) print(f"{'N':>7} {'FLOPs':>10} {'naive 内存时间':>16} {'flash 内存时间':>16} " f"{'纯带宽模型比':>9}") for N in (1024, 4096, 16384): nb = hbm_elems_naive(N, d) * b fb = hbm_elems_flash(N, d, B_r, B_c) * b fl = flops_attention(N, d) print(f"{N:>7} {fl / 1e9:>8.2f} G {nb / HBM_BW * 1e3:>14.3f} ms " f"{fb / HBM_BW * 1e3:>14.3f} ms {nb / fb:>8.1f}x") print("\n 此列只比较理想 HBM 时间,并非完整模型训练或实际 kernel 加速上界。") print(" 论文的端到端训练加速与此处单头 IO 模型口径不同,不能直接相比。") if __name__ == "__main__": report_io() report_activation() report_speed_limit() flash_online_softmax.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """online softmax 的正确性与显存代价实测(FlashAttention 的核心那一步)。 对比两种实现: naive —— 先算出完整的 N×N 分数矩阵 S,再逐行 softmax 得到 P,最后算 O = P V flash —— 按块流过 K/V,只维护三个 O(N) 的运行量:输出累加器 O、分母 l、 运行最大值 m(第 03 节递推式的直接实现) 无dropout且同一掩码时两者在实数算术下等价;求和顺序、dtype及临时数组不同。 这份脚本回答三个问题: 1. 递推式写出来的结果,和朴素 softmax 数学等价、数值误差多大?(1.1 节) 2. 峰值内存真的差一个 N 吗?(用 tracemalloc 量,不是估的) 3. 那算力呢?——在 CPU + numpy 上速度由实际运行环境决定,这一点必须诚实讲清楚 只依赖 numpy。直接 `python flash_online_softmax.py` 即可运行。 """ import time import tracemalloc import numpy as np NEG = -np.inf # ──────────────────────────────────────────────────────────── # 两个被测实现 # ──────────────────────────────────────────────────────────── def attention_naive(Q, K, V, causal=False): """标准实现:S 和 P 都是完整的 N×N 常驻张量。""" D = Q.shape[-1] S = (Q @ K.T) / np.sqrt(D) # [N, N] ← 第一块 N×N if causal: S = np.where(np.triu(np.ones((Q.shape[0], K.shape[0])), 1) > 0, NEG, S) S -= S.max(axis=-1, keepdims=True) # safe softmax,不改变结果 P = np.exp(S) # [N, N] ← 第二块 N×N P /= P.sum(axis=-1, keepdims=True) return P @ V def attention_flash(Q, K, V, B_r, B_c, causal=False, trace=False, key_valid=None): """等长单头自注意力;输入有限,float16提升到float32累加。 key_valid是一维key保留掩码;全遮行按本文约定返回零输出、LSE=-inf。 支持尾块,不实现dropout、GQA、非等长因果对齐或GPU kernel。 """ Q, K, V = (np.asarray(x) for x in (Q, K, V)) if Q.ndim != 2 or Q.shape != K.shape or Q.shape != V.shape or min(Q.shape) <= 0: raise ValueError("Q/K/V必须为相同非空[N,D]浮点矩阵") if not all(np.issubdtype(x.dtype, np.floating) and np.isfinite(x).all() for x in (Q,K,V)): raise ValueError("仅支持有限浮点输入") if any(not isinstance(b, (int, np.integer)) or isinstance(b, bool) or b <= 0 for b in (B_r,B_c)): raise ValueError("块大小必须为正整数") N, D = Q.shape dtype = np.result_type(Q.dtype, K.dtype, V.dtype, np.float32) Q, K, V = (x.astype(dtype, copy=False) for x in (Q,K,V)) valid = np.ones(N, dtype=bool) if key_valid is None else np.asarray(key_valid) if valid.shape != (N,) or valid.dtype != bool: raise ValueError("key_valid必须为[N]布尔数组") scale = 1.0 / np.sqrt(D) O = np.zeros((N, D), dtype=dtype) l = np.zeros(N, dtype=dtype) m = np.full(N, NEG, dtype=dtype) peak_tmp = 0 for j0 in range(0, N, B_c): Kj, Vj = K[j0:j0+B_c], V[j0:j0+B_c] for i0 in range(0, N, B_r): i1 = min(i0+B_r, N) if causal and i1 <= j0: continue Qi = Q[i0:i1] Sij = (Qi @ Kj.T) * scale if not np.isfinite(Sij).all(): raise ValueError("点积溢出;需缩小输入或使用更高精度") keep = np.broadcast_to(valid[j0:j0+B_c], Sij.shape).copy() if causal: keep &= (j0+np.arange(Sij.shape[1]))[None,:] <= np.arange(i0,i1)[:,None] Sij = np.where(keep, Sij, NEG) peak_tmp = max(peak_tmp, Sij.size) m_old = m[i0:i1] m_new = np.maximum(m_old, Sij.max(axis=-1)) # 对全遮行不计算-inf-(-inf),避免np.where提前求值导致NaN/警告。 shifted = np.full_like(Sij, NEG) np.subtract(Sij, m_new[:,None], out=shifted, where=np.isfinite(m_new)[:,None]) Pij = np.exp(shifted) delta = np.full_like(m_old, NEG) np.subtract(m_old, m_new, out=delta, where=np.isfinite(m_old)) corr = np.exp(delta) l[i0:i1] = l_old = l[i0:i1] * corr + Pij.sum(axis=-1) O[i0:i1] = O[i0:i1] * corr[:,None] + Pij @ Vj m[i0:i1] = m_new if trace: print(f" j={j0:>2} i={i0:>2} m={m_new[0]:+.4f} l={l_old[0]:.4f} corr={corr[0]:.4f}") out = np.zeros_like(O) np.divide(O, l[:,None], out=out, where=l[:,None]>0) LSE = np.full_like(l, NEG) active = l>0 LSE[active] = m[active]+np.log(l[active]) return out, LSE, peak_tmp def attention_backward(Q, K, V, G, B_r=64, B_c=64, causal=False, key_valid=None): """无dropout的分块反向:从O/LSE重算局部P,不存完整P。""" O, LSE, _ = attention_flash(Q,K,V,B_r,B_c,causal=causal,key_valid=key_valid) Q,K,V,G = (np.asarray(x,dtype=O.dtype) for x in (Q,K,V,G)) if G.shape != O.shape or not np.isfinite(G).all(): raise ValueError("G必须为有限[N,D]梯度") N,D=Q.shape valid=np.ones(N,dtype=bool) if key_valid is None else np.asarray(key_valid) dQ,dK,dV = (np.zeros_like(x) for x in (Q,K,V)) row_delta=(G*O).sum(axis=-1) for i0 in range(0,N,B_r): i1=min(i0+B_r,N) for j0 in range(0,N,B_c): j1=min(j0+B_c,N) keep=np.broadcast_to(valid[j0:j1],(i1-i0,j1-j0)).copy() if causal: keep &= np.arange(j0,j1)[None,:] <= np.arange(i0,i1)[:,None] S=(Q[i0:i1]@K[j0:j1].T)/np.sqrt(D) logP=np.full_like(S,NEG) np.subtract(S,LSE[i0:i1,None],out=logP,where=keep & np.isfinite(LSE[i0:i1,None])) P=np.exp(logP) dP=G[i0:i1]@V[j0:j1].T dS=P*(dP-row_delta[i0:i1,None]) dQ[i0:i1] += dS@K[j0:j1]/np.sqrt(D) dK[j0:j1] += dS.T@Q[i0:i1]/np.sqrt(D) dV[j0:j1] += P.T@G[i0:i1] return dQ,dK,dV def peak_bytes(fn, *args, **kwargs): """跑一次 fn,返回 (结果, 峰值字节数)。用 tracemalloc 实测。""" tracemalloc.start() tracemalloc.reset_peak() out = fn(*args, **kwargs) _, peak = tracemalloc.get_traced_memory() tracemalloc.stop() return out, peak def timed(fn, *args, repeat=5, **kwargs): durations = [] fn(*args, **kwargs) # 预热一次,不计入统计 out = None for _ in range(repeat): t0 = time.perf_counter() out = fn(*args, **kwargs) durations.append(time.perf_counter() - t0) return out, float(np.median(durations)) # ──────────────────────────────────────────────────────────── # 1. 递推式长什么样:一个能逐块打印的最小例子 # ──────────────────────────────────────────────────────────── def demo_recurrence(): rng = np.random.default_rng(0) N, D, B_c = 8, 4, 2 Q = rng.normal(size=(N, D)) K = rng.normal(size=(N, D)) V = rng.normal(size=(N, D)) print("=" * 68) print("1. online softmax 的递推过程(N=8, D=4, K/V 块大小 B_c=2)") print("=" * 68) print(" Q 块固定为整行(B_r=N),K/V 分成 4 块依次流过;") print(" 每个K/V块打印第 0 号 query 的 m / l / corr,看它们怎么被逐次修正:\n") _, _, peak = attention_flash(Q, K, V, B_r=N, B_c=B_c, trace=True) print(f"\n 最大的单个临时矩阵只有 {peak} 个元素(并非临时内存总和) = [{N}, {B_c}],而 N×N = {N * N}") # ──────────────────────────────────────────────────────────── # 2. 等价性:和朴素 softmax 逐位对得上吗 # ──────────────────────────────────────────────────────────── def demo_exactness(): print("\n" + "=" * 68) print("2. 数值等价性(float64,非因果 / 因果两种掩码)") print("=" * 68) print(f"{'N':>6} {'D':>4} {'causal':>7} {'最大绝对误差':>14} {'相对误差':>12}") for N in (256, 1024, 2048, 4096): D = 64 rng = np.random.default_rng(N) Q = rng.normal(size=(N, D)) K = rng.normal(size=(N, D)) V = rng.normal(size=(N, D)) for causal in (False, True): ref = attention_naive(Q, K, V, causal) got, _, _ = attention_flash(Q, K, V, B_r=64, B_c=64, causal=causal) diff = np.abs(ref - got).max() rel = diff / np.abs(ref).max() print(f"{N:>6} {D:>4} {str(causal):>7} {diff:>14.3e} {rel:>12.3e}") print("\n 误差量级是浮点舍入(1e-15),不是近似——FlashAttention 是精确算法。") # ──────────────────────────────────────────────────────────── # 3. 峰值内存:是不是真的差一个 N # ──────────────────────────────────────────────────────────── def demo_memory(): print("\n" + "=" * 68) print("3. 峰值内存实测(tracemalloc,float64,D=64,B_r=B_c=64)") print("=" * 68) print(f"{'N':>6} {'naive 实测':>12} {'naive 理论 2N²·8':>18} " f"{'flash 实测':>12} {'比值':>8}") peaks = {} for N in (1024, 2048, 4096, 8192): D = 64 rng = np.random.default_rng(N) Q = rng.normal(size=(N, D)) K = rng.normal(size=(N, D)) V = rng.normal(size=(N, D)) _, p_naive = peak_bytes(attention_naive, Q, K, V) _, p_flash = peak_bytes(attention_flash, Q, K, V, 64, 64) peaks[N] = (p_naive, p_flash) print(f"{N:>6} {p_naive / 2**20:>10.1f} MiB " f"{2 * N * N * 8 / 2**20:>16.1f} MiB " f"{p_flash / 2**20:>10.1f} MiB {p_naive / p_flash:>7.1f}x") g_n = peaks[8192][0] / peaks[4096][0] g_f = peaks[8192][1] / peaks[4096][1] print(f"\n N 从 4096 翻到 8192:naive 峰值 ×{g_n:.1f},flash 峰值 ×{g_f:.1f}") print(" 这就是 O(N²) 和 O(N) 的区别:前者翻两倍(4×),后者跟着翻倍(2×)。") print(" 朴素实现要同时留住 S 和 P 两个 N×N(理论列就是 2N²·8 字节,和实测对得上);") print(" flash 只留 B_r×B_c 的块,剩下的是 O(N·D) 的输出累加器。") # ──────────────────────────────────────────────────────────── # 4. 时间:CPU教学循环的实际耗时 # ──────────────────────────────────────────────────────────── def demo_time(): print("\n" + "=" * 68) print("4. 墙钟时间(CPU + numpy,线程配置以运行环境为准;预热1次+5次中位数)——反直觉的一项是这个") print("=" * 68) print(f"{'N':>6} {'naive':>10} {'flash(B=64)':>13} {'flash/naive':>12}") for N in (1024, 2048, 4096): D = 64 rng = np.random.default_rng(N) Q = rng.normal(size=(N, D)) K = rng.normal(size=(N, D)) V = rng.normal(size=(N, D)) _, t_naive = timed(attention_naive, Q, K, V) _, t_flash = timed(attention_flash, Q, K, V, 64, 64) print(f"{N:>6} {t_naive * 1e3:>9.1f} ms {t_flash * 1e3:>12.1f} ms " f"{t_flash / t_naive:>11.1f}x") print("\n 以上比值是本机本次计时,不能据此预测 GPU;可能影响速度的因素包括:") print(" 1. 乘加次数一模一样,还额外多了每块的 rescale 和逐元素 exp;") print(" 2. 一次大矩阵乘被拆成 (N/B)² 次 64×64 的小矩阵乘,") print(" BLAS 在小块上根本跑不满,Python 循环开销也进来了;") print(" 3. 最关键的:它省的是**内存搬运**,不是 FLOPs,") print(" CPU 同样有 SRAM 缓存,但本示例没有专门优化缓存和线程,") print(" 因此分块可能省容量,却被 Python 循环与小 GEMM 开销抵消。") print(" 第 05 节讲GPU实现如何优化IO与调度,本文不实测GPU性能。") # ──────────────────────────────────────────────────────────── # 5. 因果掩码:统计本实现跳过的全遮块比例 # ──────────────────────────────────────────────────────────── def demo_causal_blocks(): print("\n" + "=" * 68) print("5. 因果掩码下实际算了多少块(N=4096, B_r=B_c=64)") print("=" * 68) N, B = 4096, 64 T_r = T_c = N // B total = 0 for j in range(T_c): for i in range(T_r): if i * B + B <= j * B: # 整个 Q 块都在 K 块之前 → 全遮,跳过 continue total += 1 print(f" 块总数(非因果) : {T_r * T_c}") print(f" 块总数(因果跳过): {total}") print(f" 跳过的块占比 : {100.0 * (1 - total / (T_r * T_c)):.1f}%") print("\n 朴素实现就算加了因果掩码,N×N 的矩阵照样得整块算完再遮,") print(" 本脚本朴素路径不跳块;flash省去全遮块,但并非精确减半(严格说是") print(" (T²+T)/2 / T² ≈ 一半多一点的对角块保留)。") if __name__ == "__main__": np.set_printoptions(precision=4, suppress=True) demo_recurrence() demo_exactness() demo_memory() demo_time() demo_causal_blocks() attention_checks.py """CPU数学对拍与反例;依赖numpy和torch。与GPU速度复现无关。""" import inspect import json from pathlib import Path import numpy as np import torch from torch.nn import functional as F from torch.nn.attention import sdpa_kernel, SDPBackend from flash_online_softmax import attention_flash, attention_backward def main(): torch.set_num_threads(1) rng=np.random.default_rng(42) rows=[];grad_errors=[] for N in (1,7,17): D=5 for dtype in (np.float64,np.float32,np.float16): Q,K,V=(rng.normal(size=(N,D)).astype(dtype) for _ in range(3)) for causal in (False,True): for masked in (False,True): valid=np.ones(N,dtype=bool) if masked: valid[:min(3,N)]=False keep=np.broadcast_to(valid,(N,N)).copy() if causal: keep &= np.arange(N)[None,:] <= np.arange(N)[:,None] # Math后端、float64参考用于隔离输入量化和分块舍入。 ts=[torch.tensor(x.astype(np.float64)[None,None],requires_grad=True) for x in (Q,K,V)] with sdpa_kernel(SDPBackend.MATH): ref=F.scaled_dot_product_attention(*ts,attn_mask=torch.tensor(keep),dropout_p=0.0) G=rng.normal(size=(N,D));expected=torch.autograd.grad((ref*torch.tensor(G[None,None])).sum(),ts) for br,bc in ((1,1),(3,2),(8,11)): out,lse,_=attention_flash(Q,K,V,br,bc,causal=causal,key_valid=valid) error=float(np.max(np.abs(out-ref.detach().numpy()[0,0]))) assert error < (2e-12 if dtype==np.float64 else 2e-6) assert np.all(out[~keep.any(-1)]==0) and np.isneginf(lse[~keep.any(-1)]).all() rows.append(dict(N=N,dtype=str(np.dtype(dtype)),causal=causal,masked=masked,block=[br,bc],max_abs=error)) if dtype==np.float64: gs=attention_backward(Q,K,V,G,br,bc,causal=causal,key_valid=valid) err=max(float(np.max(np.abs(g-e.numpy()[0,0]))) for g,e in zip(gs,expected)) assert err<2e-12;grad_errors.append(err) Q,K,V=(rng.normal(size=(7,5)) for _ in range(3));G=rng.normal(size=(7,5));valid=np.array([False,False,True,True,False,True,True]) grads=attention_backward(Q,K,V,G,3,2,causal=True,key_valid=valid) fd_errors=[];step=1e-5 for x,g in zip((Q,K,V),grads): for idx in np.ndindex(x.shape): old=x[idx];x[idx]=old+step;plus=np.sum(attention_flash(Q,K,V,3,2,causal=True,key_valid=valid)[0]*G) x[idx]=old-step;minus=np.sum(attention_flash(Q,K,V,3,2,causal=True,key_valid=valid)[0]*G);x[idx]=old fd_errors.append(abs((plus-minus)/(2*step)-g[idx])) assert max(fd_errors)<1e-8 # 仅为展示错误:删除rescale,其他代码保持一致。不要用于计算正确注意力。 src=inspect.getsource(attention_flash).replace('corr = np.exp(delta)','corr = np.ones_like(delta)') ns=dict(np=np,NEG=-np.inf);exec(src,ns) rng=np.random.default_rng(1024);Q,K,V=(rng.normal(size=(1024,64)) for _ in range(3)) correct=attention_flash(Q,K,V,64,64)[0];bad=ns['attention_flash'](Q,K,V,64,64)[0];diff=np.abs(correct-bad) res=dict(passed=True,environment=dict(torch=torch.__version__,numpy=np.__version__,device='CPU',threads=1),forward_cases=rows, forward_max_abs=max(x['max_abs'] for x in rows),autograd_max_abs=max(grad_errors),finite_difference_parameters=len(fd_errors),finite_difference_max_abs=max(fd_errors), broken_rescale=dict(seed=1024,N=1024,D=64,max_abs=float(diff.max()),mean_abs=float(diff.mean())),errors=[],warnings=[]) Path(__file__).with_name('validation_results.json').write_text(json.dumps(res,indent=2),encoding='utf8') print(json.dumps({k:v for k,v in res.items() if k!='forward_cases'},indent=2)) if __name__=='__main__': with np.errstate(all='raise'): main() make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「FlashAttention」的三张解释图。 数值全部来自同目录下的 io_ledger.py(HBM 读写量、激活显存、算术强度), 改了那个脚本的话这里要跟着重跑,避免图与正文数字不一致。 三张图分别回答: 1. tiling 到底怎么切、什么留在 SRAM、什么留在 HBM 2. 两本账随 N 怎么长(HBM 读写量、教学激活存储模型) 3. 为什么说是「内存受限」变的「不那么受限」(算术强度对机器平衡点) 只依赖 numpy + matplotlib,直接 `python make_figures.py` 即可运行。 """ from pathlib import Path import matplotlib.pyplot as plt import numpy as np from matplotlib.patches import FancyArrowPatch, FancyBboxPatch, Rectangle from io_ledger import ( FLOPS_PEAK, HBM_BW, activation_bytes, block_sizes, flops_attention, hbm_elems_flash, hbm_elems_naive, ) ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["Microsoft YaHei", "SimHei", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, }) C_NAIVE = "#e05263" # 朴素实现:红 C_FLASH = "#2f9e6f" # FlashAttention:绿 C_SRAM = "#f0a03c" # SRAM 高亮:橙 C_GREY = "#94a3b8" C_SKIP = "#e4e9f0" INK = "#182238" MIB = 1024 ** 2 GIB = 1024 ** 3 def _style(ax, title, xlabel=None, ylabel=None): ax.set_title(title, fontsize=13.5, weight="bold", color=INK, pad=10) if xlabel: ax.set_xlabel(xlabel, fontsize=11, color="#475569") if ylabel: ax.set_ylabel(ylabel, fontsize=11, color="#475569") ax.tick_params(colors="#475569", labelsize=10) for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("left", "bottom"): ax.spines[s].set_color("#cbd5e1") ax.set_axisbelow(True) def _box(ax, x, y, w, h, text, fc, fs=10.5, tc="white"): ax.add_patch(FancyBboxPatch( (x, y), w, h, boxstyle="round,pad=0.02,rounding_size=0.12", linewidth=0, facecolor=fc, edgecolor="none")) ax.text(x + w / 2, y + h / 2, text, ha="center", va="center", fontsize=fs, color=tc, weight="bold", linespacing=1.5) def _arrow(ax, x1, y1, x2, y2, label=None, color="#475569"): ax.add_patch(FancyArrowPatch( (x1, y1), (x2, y2), arrowstyle="-|>", mutation_scale=13, linewidth=1.4, color=color, shrinkA=2, shrinkB=2)) if label: ax.text((x1 + x2) / 2, max(y1, y2) + 0.22, label, ha="center", fontsize=9, color=color) def _caption(fig, text): fig.text(0.5, 0.012, text, ha="center", fontsize=9.5, color="#94a3b8") # ──────────────────────────────────────────────────────────── # 图 1:tiling 怎么切 # ──────────────────────────────────────────────────────────── def fig_tiling(): fig = plt.figure(figsize=(13.6, 6.0), facecolor="#fbfcfe") # ── 左:朴素实现 ── ax = fig.add_axes([0.03, 0.09, 0.46, 0.80]) ax.set_xlim(0, 10); ax.set_ylim(0, 7.4); ax.axis("off") ax.text(0, 7.05, "本例朴素前向:物化 S 与 P", fontsize=13.5, weight="bold", color=C_NAIVE) _box(ax, 0.1, 5.1, 1.5, 1.1, "Q, K, V\n[N, d]", "#64748b", fs=10) _box(ax, 2.4, 4.8, 2.0, 1.7, "S = QK^T / √d\n[N, N]", C_NAIVE, fs=11) _box(ax, 5.3, 4.8, 2.0, 1.7, "P = softmax(S)\n[N, N]", C_NAIVE, fs=11) _box(ax, 8.2, 5.1, 1.6, 1.1, "O = PV\n[N, d]", "#64748b", fs=10) _arrow(ax, 1.65, 5.65, 2.35, 5.65) _arrow(ax, 4.45, 5.65, 5.25, 5.65, "读 S 写 P") _arrow(ax, 7.35, 5.65, 8.15, 5.65, "读 P") ax.text(4.4, 4.25, "N=4096, d=64, fp16:S 和 P 各 32 MiB", ha="center", fontsize=9.5, color="#475569") ax.add_patch(Rectangle((0.1, 2.2), 9.7, 1.35, facecolor="#eef2f7", edgecolor="#cbd5e1", linewidth=1.2)) ax.text(4.95, 3.2, "HBM 带宽 1.5 TB/s", ha="center", fontsize=10.5, color="#475569", weight="bold") ax.text(4.95, 2.5, "S 写一次读一次,P 写一次读一次 → 简化计数 4N² 次元素搬运", ha="center", fontsize=9.5, color="#64748b") _arrow(ax, 3.4, 4.75, 3.4, 3.6, color=C_NAIVE) _arrow(ax, 6.3, 4.75, 6.3, 3.6, color=C_NAIVE) ax.text(0.1, 1.35, "代价:数据在这条通道上来回两趟,", fontsize=10.5, color=C_NAIVE, weight="bold") ax.text(0.1, 0.75, "表中带宽数量级不同,非延迟之比", fontsize=10.5, color=C_NAIVE, weight="bold") # ── 右:FlashAttention ── ax2 = fig.add_axes([0.53, 0.09, 0.44, 0.80]) ax2.set_xlim(0, 10); ax2.set_ylim(0, 7.4); ax2.axis("off") ax2.text(0, 7.05, "FlashAttention:按块流过,只留 O(N)", fontsize=13.5, weight="bold", color=C_FLASH) T = 8 gx0, gy0, cell = 1.05, 3.1, 0.46 for i in range(T): for j in range(T): if i + 1 <= j: # 因果掩码下整块跳过(先判,优先级最高) fc, ec = C_SKIP, "#c3ccd8" elif j == 3: # 当前正在处理的 K/V 块列 fc, ec = "#cdeadb", C_SRAM elif i == 5: # 当前 Q 块行 fc, ec = "#dfeaf6", "#94a3b8" else: fc, ec = "#f7fafc", "#dde3ea" ax2.add_patch(Rectangle( (gx0 + j * cell, gy0 + (T - 1 - i) * cell), cell * 0.90, cell * 0.90, facecolor=fc, edgecolor=ec, linewidth=1.1)) gtop = gy0 + T * cell ax2.text(gx0 - 0.42, gy0 + T * cell / 2, "Q 块逐行流过", fontsize=9.5, color="#475569", ha="center", va="center", rotation=90) ax2.add_patch(FancyBboxPatch( (gx0 + 3 * cell - 0.06, gy0 - 0.06), cell * 1.02, T * cell + 0.12, boxstyle="round,pad=0.03,rounding_size=0.1", linewidth=1.6, edgecolor=C_SRAM, facecolor="none", linestyle="--")) ax2.text(gx0 + 3.5 * cell, gy0 - 0.42, "K_j, V_j 常驻 SRAM", ha="center", fontsize=9.5, color=C_SRAM, weight="bold") ax2.text(gx0 + T * cell + 0.45, gtop - 0.35, "一块 = [B_r, B_c]\n= [64, 64]", fontsize=9.5, color="#475569", va="top", linespacing=1.5) ax2.text(gx0 + T * cell + 0.45, gtop - 1.55, "灰色块:因果掩码下\n整块跳过,连算都不算", fontsize=9, color="#94a3b8", va="top", linespacing=1.5) ax2.text(gx0 + T * cell + 0.45, gtop - 2.95, "片上空间受限:\n驻留 K/V 与 Q/O 块,\n还要容纳分数和统计量", fontsize=9, color=C_SRAM, va="top", linespacing=1.5) by = 1.85 ax2.text(0.15, by + 0.52, "跨 K/V 块一直复用的三个量:", fontsize=10.5, color=INK, weight="bold") for k, (label, color) in enumerate([ ("O_i 未归一化累加器 [N, d]", C_FLASH), ("l_i 分母 [N]", "#7fb3d5"), ("m_i 运行最大值 [N]", "#c9a227")]): ax2.add_patch(Rectangle((0.15, by - k * 0.46), 4.6, 0.28, facecolor=color, edgecolor="none")) ax2.text(4.9, by - k * 0.46 + 0.14, label, fontsize=9.5, color="#475569", va="center") ax2.text(0.15, 0.12, "每来一块就 rescale 一次(乘 e^{m旧−m新}),最后 O/l 收尾", fontsize=9.5, color="#64748b") _caption(fig, "图 1:切的是「K/V 块 × Q 块」这两层循环,不是把注意力切开。" "左边突出两份二次中间量,右边只维护线性的累加量(固定 d)。") fig.savefig(OUT / "tiling.png", facecolor="#fbfcfe") plt.close(fig) print(" figures/tiling.png") # ──────────────────────────────────────────────────────────── # 图 2:两本账随 N 怎么长 # ──────────────────────────────────────────────────────────── def fig_curves(): d, b = 64, 2 B_r, B_c = block_sizes(d) Ns = np.array([1024, 2048, 4096, 8192, 16384, 32768]) naive_io = np.array([hbm_elems_naive(N, d) * b for N in Ns]) / MIB flash_io = np.array([hbm_elems_flash(N, d, B_r, B_c) * b for N in Ns]) / MIB fig, axes = plt.subplots(1, 2, figsize=(13.2, 5.1), facecolor="#fbfcfe") fig.subplots_adjust(bottom=0.20, top=0.86, wspace=0.25) ax = axes[0] ax.loglog(Ns, naive_io, "o-", color=C_NAIVE, lw=2.2, ms=6, label="朴素实现(∝N²)") ax.loglog(Ns, flash_io, "s-", color=C_FLASH, lw=2.2, ms=6, label="FlashAttention(∝N²/M)") ax.annotate(f"{naive_io[2] / flash_io[2]:.1f}×", xy=(Ns[2], naive_io[2]), xytext=(Ns[2] * 1.6, naive_io[2] * 3.0), fontsize=11, color=C_NAIVE, weight="bold", arrowprops=dict(arrowstyle="->", color=C_NAIVE, lw=1.3)) ax.legend(fontsize=10, frameon=False, loc="upper left") ax.grid(True, which="both", color="#eef2f7", lw=0.9) _style(ax, "HBM 读写量(单头 d=64, fp16)", "序列长度 N", "MiB") ax = axes[1] d2, H, B = 3072, 24, 1 lay = 32 naive_act = np.array( [activation_bytes(N, d2, H, B)[0] * lay for N in Ns]) / GIB flash_act = np.array( [activation_bytes(N, d2, H, B)[1] * lay for N in Ns]) / GIB ax.loglog(Ns, naive_act, "o-", color=C_NAIVE, lw=2.2, ms=6, label="朴素实现(∝N²)") ax.loglog(Ns, flash_act, "s-", color=C_FLASH, lw=2.2, ms=6, label="FlashAttention(∝N)") ax.axhline(80, color="#64748b", ls="--", lw=1.4) ax.text(Ns[0] * 1.15, 92, "示例预算:80 GiB(非硬件规格)", fontsize=9.5, color="#64748b") ax.annotate("76.5 GiB", xy=(4096, naive_act[2]), xytext=(4096 * 1.8, naive_act[2] * 2.2), fontsize=10.5, color=C_NAIVE, weight="bold", arrowprops=dict(arrowstyle="->", color=C_NAIVE, lw=1.3)) ax.annotate("4.5 GiB", xy=(4096, flash_act[2]), xytext=(4096 * 0.55, flash_act[2] * 6.0), fontsize=10.5, color=C_FLASH, weight="bold", arrowprops=dict(arrowstyle="->", color=C_FLASH, lw=1.3)) ax.legend(fontsize=10, frameon=False, loc="upper left") ax.grid(True, which="both", color="#eef2f7", lw=0.9) _style(ax, "教学激活存储模型(简化视频 DiT,32 层)", "序列长度 N", "GiB") _caption(fig, "图 2:两张都是对数轴,斜率就是复杂度阶数。" "左图的大 N 主导项均为 N²,有限 N 时比例受线性项和取整影响;" "右图斜率不同(N² 对 N),N 越大差距越离谱。") fig.savefig(OUT / "io_curve.png", facecolor="#fbfcfe") plt.close(fig) print(" figures/io_curve.png") # ──────────────────────────────────────────────────────────── # 图 3:算术强度 vs 机器平衡点 # ──────────────────────────────────────────────────────────── def fig_intensity(): d, b = 64, 2 B_r, B_c = block_sizes(d) Ns = np.array([1024, 2048, 4096, 8192, 16384, 32768]) naive_i = np.array([flops_attention(N, d) / (hbm_elems_naive(N, d) * b) for N in Ns]) flash_i = np.array([flops_attention(N, d) / (hbm_elems_flash(N, d, B_r, B_c) * b) for N in Ns]) balance = FLOPS_PEAK / HBM_BW fig, ax = plt.subplots(figsize=(11.4, 5.0), facecolor="#fbfcfe") fig.subplots_adjust(bottom=0.22, top=0.86) ax.axhspan(0, balance, color="#fdf1f2", zorder=0) ax.axhline(balance, color="#64748b", ls="--", lw=1.6, zorder=3) ax.text(Ns[-1] * 2.1, balance * 0.70, f"机器平衡点 {balance:.0f} FLOP/byte\n(峰值算力 ÷ HBM 带宽)", fontsize=10, color="#475569", va="center", linespacing=1.5) ax.text(Ns[0] * 1.02, balance * 0.40, "线下:本 roofline 模型由 HBM 项主导(实际瓶颈需测量)", fontsize=10, color=C_NAIVE) ax.text(Ns[0] * 1.02, balance * 1.32, "线上:本模型计算项主导", fontsize=10, color=C_FLASH) ax.semilogx(Ns, naive_i, "o-", color=C_NAIVE, lw=2.4, ms=7, label="朴素实现", zorder=4) ax.semilogx(Ns, flash_i, "s-", color=C_FLASH, lw=2.4, ms=7, label="FlashAttention", zorder=4) for x, v in zip(Ns, naive_i): ax.annotate(f"{v:.0f}", (x, v), textcoords="offset points", xytext=(0, -16), ha="center", fontsize=9, color=C_NAIVE) for x, v in zip(Ns, flash_i): ax.annotate(f"{v:.0f}", (x, v), textcoords="offset points", xytext=(0, 9), ha="center", fontsize=9, color=C_FLASH) ax.set_xticks(Ns) ax.set_xticklabels([f"{n // 1024}K" for n in Ns]) ax.set_xlim(Ns[0] * 0.75, Ns[-1] * 4.5) ax.set_ylim(0, balance * 1.85) ax.legend(fontsize=10.5, frameon=False, loc="center right") ax.grid(axis="y", color="#eef2f7", lw=0.9) _style(ax, "算术强度:每搬 1 字节能做多少次运算(单头 d=64, fp16)", "序列长度 N", "FLOP / byte") _caption(fig, "图 3:朴素实现的强度几乎不随 N 变(一直贴在 32 附近)," "在此模型中由 HBM 项主导;FlashAttention 的估计抬过平衡点," "并不单独证明实际 kernel 已充分利用计算单元。") fig.savefig(OUT / "intensity.png", facecolor="#fbfcfe") plt.close(fig) print(" figures/intensity.png") if __name__ == "__main__": print("生成配图:") fig_tiling() fig_curves() fig_intensity()
2026年09月25日
5 阅读
0 评论
0 点赞
2026-09-25
AIGC 每日速读|2026-09-25|快手开源修图登顶,GEdit仍输GPT
今日 AIGC 论文速览 今日共 10 篇 · 图像编辑与生成安全 2 篇 · 统一预训练与视频运动 2 篇 · 多模态联合生成 2 篇 · 音频生成与后训练 2 篇 · 推理加速与视频评测 2 篇 重点论文标题列表 KwaiMind(快手):180万对数据炼电商修图 UVU(清华大学):视觉监督前移,RefCOCO涨7.7 ⚡ Unite-Audio(东京科学大学):117M音频生成,CLAP登顶 MotionSpec(卡迪夫大学):频谱轨迹约束视频运动 RecCAR(巴伊兰大学):反向注意力补齐,音画更同步 今日论文速览 1. KwaiMind:180万对数据炼电商修图 KwaiMind Technical Report | 快手 | arXiv:2609.26375 关键词:图像编辑,电商,CTR奖励,Ecom-Bench,在线强化学习 ⚠️ 前序问题:电商修图不是通用编辑:商品主体要一模一样、促销文字要渲染准确、图还得有点击欲。通用编辑模型在这三件事上都没专门优化,而电商数据的采集清洗又贵又杂,缺一套能稳定产出高质量配对数据的流水线。 本文贡献:快手发布 KwaiMind:多模态 DiT 底座,Agent 数据引擎维护约 180 万高质量编辑对;继续预训练加 SFT 后走偏好优化与在线 RL,用 VLM 通用裁判加 CTR、文字渲染、商品一致性三类专用奖励训练专精策略,再经 on-policy 蒸馏合并成单一模型;配套发布覆盖 11 类商业编辑任务的 Ecom-Bench。 Overview of the end-to-end data pipeline. Coordinator Agent routes samples through filtering, generation, captioning, and post-filtering with bounded feedback loops and human review. 实验效果:开源编辑器里综合最强:ImgEdit 4.15、GEdit 5.79、REDEdit 英/中 3.75/3.73,Ecom-Bench 视觉质量第一;CTR 引导优化把生成图预测 CTR 超过原图的比例从 12.16% 提到 37.41%,线上 A/B 实际 CTR 相对提升约 2.44%。 Overall comparison on general image editing benchmarks: ImgEdit, GEdit, and the English and Chinese splits of REDEdit. Hatched bars denote closed-source models. 批判点评:最强的限定词是开源:图上闭源模型全面更高——ImgEdit 被 Seedream 4.0 的 4.27 和 GPT-Image-2 的 4.20 压着,GEdit 上 Nano Banana 2 拿 7.02、GPT-Image-2 拿 7.10,比 KwaiMind 的 5.79 高出一大截,REDEdit 两个语种同样如此。真正立住的卖点其实是 CTR 这条商业指标,而不是编辑质量本身。 2. UVU:视觉监督前移,RefCOCO涨7.7 UVU: Improving Multimodal Understanding via Vision-Language Unified Autoregressive Paradigm | 清华大学;腾讯优图实验室;南京大学;格拉斯哥大学 | arXiv:2609.27915 关键词:统一自回归,像素级codebook,连续视觉编码,视觉监督,预训练 ⚠️ 前序问题:多模态大模型的细粒度视觉理解长期靠稀疏文本监督撑着,已有的视觉监督大多加在后训练阶段,那时视觉表征已基本定型,监督信号只能当辅助约束。要重塑感知骨干,得把视觉监督搬进预训练。 本文贡献:提出 UVU 视觉语言统一自回归范式:不用向量量化,SigLIP-2 连续视觉特征直接进 LM 解码器做下一 token 预测;设计大规模迭代层次聚类算法构建 20 万词表的像素级视觉 codebook,让图像 patch 与文本 token 在预训练期共享统一监督,模型由此内化视觉重建能力。 Overview of the UVU vision-language unified autoregressive framework. Continuous visual features from SigLIP-2 are projected and integrated with textual embeddings for autoregressive next-token prediction in an LM decoder, generating pixel-level image tokens. 实验效果:同为 3B/Qwen2.5 底座,RefCOCO 从 84.1 涨到 91.8、LISA 57.4→71.7、CVBench-3D 71.9→76.6、BLINK 46.9→52.8;相对去掉视觉监督的自版 UVU*,RefCOCO +6.2、LISA +12.1。注意力热图显示 UVU 更聚焦目标区域。 Comparison of visual attention heatmaps under three paradigms: from left to right, without visual supervision, AR + VQ, and UVU (Ours). 批判点评:对手没输干净:SEEDB 74.1 仍低于 LLaVA-OV 的 75.4,MMStar 55.0 输给 56.7,ScienceQA 91.3 远低于 GLM-4v 的 96.7;MME 2201.9 对 LLaVA-OV 2146.3 的领先也只有 2.6%。统一目前只覆盖理解侧,生成侧数据还没进训练,作者自己也在文中承认。 3. Unite-Audio:117M音频生成,CLAP登顶 UNITE-AUDIO: Joint Learning of Continuous Tokenization and Latent Flow Matching for Text-to-Audio Generation | 东京科学大学;清华大学;武汉大学;东京大学;蚂蚁集团 | arXiv:2609.28206 关键词:文本生成音频,联合训练,flow matching,Flow-GRPO,连续表征 ⚠️ 前序问题:文本生成音频的主流做法是两阶段:先训重建导向的音频 tokenizer 再冻结,然后在固定隐空间训生成模型。但为重建优化的表征未必适合生成,生成目标无法反过来塑造隐空间。 本文贡献:首个把连续音频表征学习和隐空间 flow matching 联合训练的 TTA 系统:掩码波形过在线编码器,EMA 编码器提供停止梯度的目标特征,重建与自监督生成预测耦合,让生成目标直接塑造隐空间;再用 Flow-GRPO 后训练强化文本条件对齐。 Overview of the proposed framework: (a) the reconstruction path and (b) the generation path. Blue denotes regions with gradient propagation, while gray denotes regions without gradient propagation. 实验效果:117M 隐空间生成模型(不含编解码器)在 AudioCaps-886 上 CLAP 0.534 全表最高(EzAudio 0.496、TangoFlux-RL 0.480),KL 1.057 最低,RTF 0.062;NFE 压到 4 时 RTF 仅 0.009 仍保 CLAP 0.512。 Comparison between conventional separate tokenizer-generator training and our joint single-stage training. 批判点评:分布指标难看:FD_OpenL3 84.76 是 EzAudio 38.26 的两倍多,FD_16k 43.25 也落后 GenAU 的 31.94,IS 11.34 不及 TangoFlux-RL 的 12.20——语义对齐赢了,音频分布保真还差得远;首个联合训练的含金量取决于你更看重 CLAP 还是 FD。 4. MotionSpec:频谱轨迹约束视频运动 MotionSpec: Spectral Trajectory Supervision for Motion-Consistent Video Generation | 卡迪夫大学;中国科大;华东师范大学 | arXiv:2609.28095 关键词:视频生成,运动一致性,傅里叶谱,光流,微调目标 ⚠️ 前序问题:文生视频单帧越来越真,运动却常失真:动作推进不一致、复杂动作下结构崩坏。标准生成目标对运动本身几乎没有专门约束,运动演化处于欠约束状态。 本文贡献:提出 MotionSpec 运动监督:谱轨迹一致性 STC 构造锚点相对的稠密运动轨迹,经时间傅里叶变换成运动谱体积,对齐幅度与相位同时约束运动强度和时序组织;辅以局部光流一致性 LFC 稳定相邻帧过渡,两者随时间步加权加进微调目标。 Overview of the pipeline. At each training step, the video generative model takes the prompt, timestep, noise level, and real data as inputs. 实验效果:在 Wan2.1 上微调,VMBench 均分 61.57→65.07,其中物体完整性 OIS 48.34→56.53、时序一致性 TCS 97.09→98.82;VideoJAM-Bench 运动分 85.3→93.61。定性对比里 Wan2.1 的骑车段后半程自行车几乎消失,MotionSpec 保留完整。 Qualitative Results. On the left are the results of Wan2.1, and on the right are the results of ours. 批判点评:外观分几乎没动:App 仅 81.95→82.20(+0.25);PAS 上单独用 LFC(24.27)反而高于完整版(23.74),说明两个目标存在拉扯。评测全部在 Wan2.1 一个底座上做,换模型的泛化性未知。 5. RecCAR:反向注意力补齐,音画更同步 All modalities are equal, but video is more equal: Closing the Cross-Attention Gap in Joint Video Generation | 巴伊兰大学;NVIDIA | arXiv:2609.27901 关键词:联合视频生成,跨模态注意力,KL正则,音画同步,人体解剖 ⚠️ 前序问题:联合视频-动作/视频-音频生成的扩散 Transformer 里存在不对称:伴随模态到视频的对应很强,但反过来约束视频生成的对应路始终弱——视频听动作指令的能力远弱于动作看视频的能力,人体解剖崩坏、音画不同步由此而来。 本文贡献:把两个方向的跨模态对应表示成视频 token 上的分布,定义反向对应差距;提出 RecCAR,KL 正则把弱方向对齐到以视频到模态对应为固定参考的强方向,即插即用地微调 EchoMotion 这类联合生成模型。 Illustration of asymmetric reciprocal cross-modal correspondence. For a motion token corresponding to the head, the VM correspondence correctly localizes the head region in the video, while the reciprocal MV correspondence is dispersed. 实验效果:视频-动作联合生成 Human Anatomy 0.69→0.75,VMBench 多数指标超 EchoMotion、CoMoVi、FlowMo;视频-音频生成音画失同步从 0.804 降到 0.752。散点图上每个 prompt 的反向差距在微调后一致下降。 Qualitative comparison of RecCAR against EchoMotion, CoMoVi, and FlowMo across three representative scenes. 批判点评:增益以自家基线为参照:0.75 的解剖分意味着仍有约四分之一样本不达标,音频侧 0.752 的失同步也谈不上小;teaser 里对比的 LTX-2 只做了定性展示,未给出与闭源模型在同一协议下的量化差距。 6. GestureFAR:9.3毫秒每token的流式手势 GestureFAR: Streaming Co-Speech Gesture Generation with Flow Autoregression | 罗切斯特大学;东京大学;加州大学圣克鲁兹分校;加州大学洛杉矶分校;Meta | arXiv:2609.21576 关键词:手势生成,流式,flow自回归,蒸馏,实时交互 ⚠️ 前序问题:流式陪聊手势生成此前靠离散运动 token 自回归,把高维连续运动压进有限 codebook,真实感和多样性都受损;而连续方法又难以保证逐 token 因果,实时交互卡在延迟上。 本文贡献:提出流式手势生成框架 GestureFAR:因果 VAE 把全身动作编码成可流式的连续隐变量,Transformer 对音频-运动上下文自回归,每 token 配 flow-matching 头从连续分布采样下一潜变量;再冻结因果骨干,用一致性加分布匹配目标把多步 flow 头蒸馏成单步前向。 Overview of GestureFAR. GestureFAR generates co-speech gestures from streaming audio by autoregressing over continuous motion latents. 实验效果:BEAT2 流式组 FGD 3.08 全面最佳(LiveGesture 4.57、MIBURI 8.06),BC 0.741 接近离线最优;1 步蒸馏后每 token 9.3ms,是教师 8 步 24.4ms 的 2.6 倍速、MIBURI 62.9ms 的 6.8 倍速。 Qualitative comparison on BEAT2. We visualize generated full-body gestures from different methods under the same speech inputs, with the semantically salient words highlighted in red. 批判点评:BC 0.741 仍低于 MIBURI 的 0.790 和 LiveGesture 的 0.794,多样性 13.24 也略逊离线系;单步蒸馏换速度后 FGD 反而比教师还好(3.08 vs 3.17),这个反常更像评测协议的产物而非能力证明,文中未解释。 7. DeltaS:状态漂移选KV,六个基准+2.1 DeltaS: Reading the Gated Linear Attention State for KV Cache Eviction in Streaming Video | Maum AI;首尔大学;延世大学 | arXiv:2609.27470 关键词:KV缓存淘汰,流式视频,线性注意力,混合架构,训练无关 ⚠️ 前序问题:视频语言模型转向线性/全注意力混合架构后,线性注意力状态固定大小,但全注意力 KV 缓存仍随流式视频无限增长;流式场景下问题还没来就得决定淘汰谁,现有基于位置、注意力或 KV 本身的信号都拿不到这个先验。 本文贡献:提出训练无关的 DeltaS:读门控 delta 线性注意力的循环状态,用一块帧内状态的归一化变化量(状态漂移)衡量该块带来多少新信息,漂移大的块对应 KV 留存;信号计算只占前向的 1.9%,无需代理查询。 Overview of DeltaS. Each video chunk updates the recurrent linear-attention states, whose normalized changes yield a shared score for retaining the corresponding KV entries across full-attention layers. 实验效果:预算与留存策略对齐的受控比较里,状态漂移信号全面胜过位置/注意力/KV 三类基线;六个长视频基准平均超最强无查询基线 2.1 分,最长基准 LVBench 超 5.6 分。 Performance margin as a function of KV retention rate. DeltaS generally shows larger margins at lower retention rates. 批判点评:增益随留存率升高而收窄:留存 token 超过约四成后在 EgoS 上反而落后基线 0.7 分,高预算场景收益有限;hybrid 架构前提也把它限死在门控 delta 线性注意力这一类骨干上。 8. DriftAudio:一步生成后训练FAD降34% DriftAudio: Marginal Drifting for Distributional Post-Training of One-Step Text-to-Audio Generators | 悉尼科技大学 | arXiv:2609.27598 关键词:一步生成,分布后训练,文本生成音频,Drifting,FAD ⚠️ 前序问题:一步文本生成音频模型把推理成本打下来之后,生成分布仍欠火候;常规后训练依赖逐条件配对的真实样本,自由文本条件下一个 prompt 往往只有一两条真实参考,逐条件 Drifting 根本做不了。 本文贡献:提出 DriftAudio 边缘分布后训练:把 Drifting 从逐条件搬到边缘音频分布上做,在冻结的 PANNs 特征空间里用重采样真实样本、滚动生成的 FIFO 样本库和当前批次共同估计漂移场,得到 detached 训练目标只更新生成器,一步推理流程原样保留。 Overview of DriftAudio. The one-step generator remains text-conditioned, while Drifting is performed on the marginal audio distribution in feature space. 实验效果:从 MeanAudio 出发 FAD 降 33.9%、FD 降 17.6%,KL 与 CLAP 同步改善;从 FdAudio 出发 FAD 1.22→0.99,一步生成的分布质量明显抬升。 Mean and covariance components of FAD for the four evaluated models. Numbers above the bars denote total FAD. 批判点评:不是白捡的:从 FdAudio 出发时 IS 和 CLAP 出现回退,说明边缘分布对齐和条件保真之间存在交换;方法绑定冻结的 PANNs 特征空间,换更强的音频评价骨干是否还成立未验证。 9. InGuard:嵌入层安检,省一半去噪步 InGuard: Towards Generalized Inner Guardrail for Safe Text-to-Image Generation | 阿里巴巴 AAIG | arXiv:2609.27620 关键词:文生图安全,内嵌护栏,嵌入改写,潜空间检测,RevGen ⚠️ 前序问题:T2I 的安全护栏都装在外面:前置 prompt 分类器加后置图像分类器,两者不用模型自身表征——prompt 检查准头有限,图像检查要等全量去噪烧完才跑,而且违规 prompt 只能一刀切拒绝,本可改写成安全输出的也被扔掉。 本文贡献:提出 InGuard 内嵌护栏:文本编码器嵌入上直接做风险三级分类(不安全/风险/良性),SAGE 对风险 prompt 做嵌入空间软门控改写使其输出安全图而非拒答,去噪中途用一步干净潜变量估计做潜空间检测,发现风险即刻提前终止;配套 RevGen Safety 基准,1 万条 prompt 由真实图反生成并注入受控 IP 角色。 Overview of our InGuard framework. Three complementary components work in sequence: prompt embedding risk classification, SAGE embedding-space safety enhancement, and latent detection with early termination. 实验效果:五个开源 T2I 模型上安全率 97.9%-98.8%,追平或超过外置护栏;良性扰动少 57.5%-73.5%,参数少约 3.7 倍,50%-55.6% 的去噪步被跳过;潜空间检测在 44%-50% 计算量处就逼近图像级 F1。 Pretraining effects on Avg F1 across five LDMs; latent detection vs denoising-step fraction used as a compute proxy. 批判点评:SAGE 改写不是万能的:论文自己承认增强失败时只能靠潜空间检测兜底拦截,等于承认存在漏改样本;RevGen 由自家反生成流水线构造,跨基准的代表性有待第三方复检。 10. CineGuard:11种运镜的抄袭检测 Did You Steal My Shot? Pioneering Camera Motion Plagiarism Detection in Generative Videos | 河海大学 | arXiv:2609.22267 关键词:运镜抄袭,生成视频,涡度,基准,版权保护 ⚠️ 前序问题:运镜是影视级 IP,但生成视频用一句 prompt 就能复刻高价值运镜;现有相似度检测都盯着画面内容,训练数据把运镜和内容缠在一起,传统光流又表达不了复杂相机运动——运镜抄袭至今无人管。 本文贡献:构建首个运镜分析基准 CineFlow:11 种运镜风格的模拟数据集把运镜与内容解耦;提出 CineGuard 检测网络,在光流上叠加流体力学涡度作为平移不变线索,3D 卷积提 tubelet 嵌入后过 ViT 主干加全局相机适配器,对比学习聚拢同运镜视频。 Architecture of the motion plagiarism detection network. The input flow is first augmented with vorticity, then a 3D convolution extracts tubelet embeddings. 实验效果:抄袭检测超最强基线 3.02 倍;在 Veo 3.1 与即梦的商业生成视频上,余弦相似度 0.56-0.83,基线全部落在 0.35 以下,Precision@5 最高 0.73 对基线约 0.27。 Performance of plagiarism detection on generative videos. Left: cosine similarity between videos. Right: Precision@5 for retrieving videos with the same motion label. 批判点评:局限写在图里:FPV Drone 这类复合运镜相似度掉到 0.56,是全表最低,混合运镜仍是软肋;基准是简单几何场景模拟出来的,真实影视素材上的迁移只有商业视频小样本佐证。 趋势观察 编辑模型的主战场从通用能力挪向行业指标 KwaiMind 把 CTR 预测直接做成奖励信号训进编辑模型,线上 A/B 实测 +2.44%;InGuard 则把安全检查搬进生成管线内部,省下一半去噪步。两条线指向同一件事:图像生成的竞争正在从「生成质量」转向「商业约束下的生成质量」,奖励函数和护栏的构造方式开始比骨干网络更值钱。 视觉监督正在从后训练前移到预训练 UVU 用像素级 codebook 把视觉监督直接塞进预训练,RefCOCO 一项涨 7.7 分;MotionSpec 则在视频微调里用频谱约束补上运动监督的缺口。共同逻辑:等表征定型再补监督太晚了,感知质量的上限在预训练阶段就被决定。 一步生成与流式生成进入质量补课期 DriftAudio 给一步 TTA 做边缘分布后训练,FAD 降三分之一;GestureFAR 用单步蒸馏把手势生成压到每 token 9.3ms;Unite-Audio 把去噪步数压到 4 还能保住 CLAP。速度战的下一程是质量战——谁能在一步、几步的预算里把分布差距补回来,谁就拿到实时产品化的门票。 混合架构的两组记忆开始学会协作 DeltaS 证明线性注意力的循环状态可以反过来指挥全注意力 KV 缓存的淘汰,六个长视频基准平均 +2.1;Unite-Audio 则让生成目标直接塑造 tokenizer 隐空间,打破「先冻结表征再训生成」的铁律。两篇都在挑战模块化管线里各训各的默认假设。 人工智能炼丹君 整理 | 2026-09-25 更多 AIGC 论文解读,关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注
2026年09月25日
22 阅读
0 评论
0 点赞
2026-09-25
AIGC 基本功|性能建模与 Profiling:算力、带宽与显存账本-Roofline
AIGC 基本功|性能建模与 Profiling:算力、带宽与显存账本-Roofline 所属方向:推理加速 | 难度:进阶 | 前置知识:混合精度与数值稳定性、自注意力机制的计算与显存账本 关键词:性能建模、Profiling、Roofline、算术强度、memory bandwidth、FLOPs、MFU、PyTorch Profiler、Nsight 01. 为什么需要它 先看本次 CPU 复核得到的两段计算:GEMM 约 98.57 MFLOP,逐元素 add 约 100 MFLOP;二者运算量相近,最短耗时却分别为 1.029 ms 和 23.623 ms,差 22.9 倍。环境为 Windows / AMD64、NumPy 2.4.6、OpenBLAS,线程请求值为 1。machine_probe.py 会记录实际 NumPy 构建配置和全部计时样本。 这里 add 复用 200 万元素的三份 fp32 数组,工作集 24 MB,重复 50 轮;累计逻辑流量 1200 MB,并不等于实测 DRAM 搬运量。GEMM 通过矩阵复用获得更高逻辑算术强度,但本实验没有硬件计数器,不能把 22.9 倍差距全部归因于 DRAM 带宽。它足以提醒我们:只比较 FLOPs,不能预测耗时。 再看一个优化评审:想把 LayerNorm 换成更快的低精度计算单元。若一行数据确实能驻留片上、只读输入并写输出,本文简化账本的强度约为 1.5 个教学操作/Byte;是否由带宽主导,还要核对标量、归约、平方根的执行资源,不能直接与 BF16 Tensor Core 的 ridge point 比较。只有在已确认流量主导、其他资源不变的模型里,单独提高计算峰值的收益才是 1.00。 小 batch GEMM 的权重量化常能减少字节,但反量化、格式元数据和内核效率也要记账。 这篇的目的就是把这套「先算账再动手」的方法补齐:两个账本——速度账(算力 vs 带宽)和容量账(显存四项)——加上一套实测手段(profiler)。量化、算子融合、KV cache 管理、并行切分,所有推理优化节点的收益判断都站在这篇的地基上。 02. 最小可用理解 阅读前只需会矩阵形状、乘加计数和单位换算;混合精度一篇解释 dtype 与字节数,注意力一篇解释 Q/K/V。本篇会补足异步计时、缓存层级与活跃 token 的必要背景,不要求先懂 CUDA 内核。 三句话讲完核心: 任何算子的耗时有一个下界,由两种资源里更慢的那个决定:$T \ge \max(F/P_{\text{peak}},\ D/\beta)$。$F$ 是算子要做的浮点运算数,$P_{\text{peak}}$ 是硬件与该指令类型匹配的峰值算力;$D$ 是算子要搬运的字节数,$\beta$ 是带宽。算得再快也快不过「数据没到」。 算术强度 $I = F/D$ 决定卡在哪:与 ridge point $I^{\ast} = P_{\text{peak}}/\beta$ 比较,$I \lt I^{\ast}$ 表示该二资源模型的流量项主导,优先检查少搬字节的机会(融合、量化、FlashAttention);$I \gt I^{\ast}$ 表示计算项主导,再检查少算或提高同口径吞吐的机会。两种分类都须用实际流量和热点执行资源核验。 显存容量是另一本独立的账:权重 + KV cache + 激活 + 额外开销四项加总,才决定一张卡能塞多少并发。「权重放得下就能跑」只覆盖了四项里的第一项。 03. 数学推导 3.1 时间下界为什么取 max 一个算子要做 $F$ 个浮点运算(FLOP,口径:一次乘加记 2 个 FLOP),硬件每秒最多做 $P_{\text{peak}}$ 个——就算计算单元一刻不停,也至少要 $F/P_{\text{peak}}$ 秒。同理,算子要把 $D$ 字节的数据在内存和计算单元之间搬个来回,总线每秒最多搬 $\beta$ 字节——至少要 $D/\beta$ 秒。无论计算和搬运是否重叠,总时间都不能短于这两个最低时间,所以可取 max 作为下界。若进一步用等号近似真实时间,才需要充分重叠和接近峰值利用率等假设: $$T \ge \max\left(\frac{F}{P_{\text{peak}}},\ \frac{D}{\beta}\right)$$ 注意这是下界:真实 kernel 还有 kernel launch、同步、缓存未命中、TLB miss 等额外开销,只要 $F$、$D$ 和对应物理峰值口径正确,实际耗时不能突破这个下界。后面的实验用经验标定值和逻辑数组流量代入,此时算出的只是参考模型估计;计时波动或缓存会产生超过 100% 的达成率,并不违反物理下界。差距也不能全算成可消除的优化空间。 3.2 算术强度与 ridge point 定义算术强度: $$I = \frac{F}{D}$$ 物理含义:每从内存搬 1 字节数据,能换来多少次浮点运算。它取决于实现与所选存储层级;实际缓存命中和重复加载又与硬件有关,不能视为完全与机器无关。再定义机器的 ridge point: $$I^{\ast} = \frac{P_{\text{peak}}}{\beta}$$ 物理含义:给定指令资源与存储层级下,「算」和「搬」模型时间一样长的分界点,单位都是 FLOP/Byte,所以可以比。把 $I$ 与 $I^{\ast}$ 代回 3.1 的下界: $I \lt I^{\ast}$(模型流量项主导):$D/\beta$ 那一项更大,$T \approx D/\beta$,模型性能上界 $F/T \le I \cdot \beta$——性能与算力峰值无关,只跟 $I$ 成正比,这就是 roofline 图上那条斜线; $I \ge I^{\ast}$(模型计算项主导):$T \approx F/P_{\text{peak}}$,性能封顶在 $P_{\text{peak}}$,这就是平顶。 $P_{\text{peak}}$ 必须匹配运算类型。本文采用 A100 40GB SXM 白皮书 Table 1的 BF16 dense Tensor Core 312 TFLOP/s、HBM2 带宽 1555 GB/s,得到 $I^{\ast}=200.6$;同表 FP32 CUDA core 只有 19.5 TFLOP/s,对应 12.5 FLOP/Byte。A100 的表中没有原生 FP8 Tensor Core。保留这些旧卡作为易复算例子,数字不是对所有算子的统一峰值。 H100 SXM 官方规格给带稀疏脚注的 BF16 1979、FP8 3958 TFLOP/s;除去结构化稀疏的二倍口径后约为 989.5、1979。代码取 BF16 的约数 989、带宽 3.35 TB/s。下面若对 add、LayerNorm、softmax 统一代入 $P=312$ TFLOP/s,只是在演示一个假想统一操作吞吐的教学模型,不能据此诊断 A100 上的标量算子。真实不同资源的写法见 3.6。 3.3 给几类算子记账 GEMM:$A[M,K] \times B[K,N] \to C[M,N]$。每个输出元素要做 $K$ 次乘加,共 $M N K$ 次,乘加各记一次: $$F_{\text{gemm}} = 2MNK,\qquad D_{\text{gemm}} = b\,(MK + KN + MN)$$ $b$ 是每元素字节数(bf16 取 2)。这是每份输入读一次、输出写一次且不读取旧 C 的最低流量示例,未计分块重读、write allocate 与输出累加。固定 K/N 时增大 M 可摊薄权重读取,提高模型 I;大 batch 更可能接近计算侧、小 batch 更可能受流量影响,但缓存命中、矩阵布局和内核调度可能改变实际瓶颈。 逐元素算子:$n$ 个元素各做 1 次运算,读 2 份写 1 份: $$F = n,\qquad D = 3 b n,\qquad I = \frac{1}{3b}$$ $I$ 是常数(fp32 下约 0.083,bf16 下约 0.17),与 $n$ 无关。在固定 dtype 与访存模型下,增大 n 本身不会提高 I,因而不会像增加 GEMM 的复用维度那样跨越 roofline 分界。 简化 LayerNorm:共 n 行、每行 d 个元素(本节不加偏置,只保留 gamma),先求均值、减均值、平方后再求均值,最后除标准差并乘 gamma;忽略每行 eps 加法/平方根等小量,按教学规则每个元素约六次操作。若整行及统计量能复用、gamma 读取等小量忽略,最低流量示例为一次读与一次写: $$F_{\text{ln}}\approx6nd,\qquad D_{\text{ln,min}}\approx2bnd,\qquad I_{\text{ln,min}}\approx\frac{3}{b}$$ bf16 时得到约1.5个教学操作/Byte。本文 NumPy 分步实现会反复读写大数组,共约10遍,fp32 下 $I\approx6/(10\times4)=0.15$;广播 gamma 和行统计量按可复用小数组另计。两种流量口径不同,不应把最低流量模型冒充 NumPy 的物理实测流量。 softmax 的读写口径:本文的三阶段教学实现先读 X 求 max,再读 X、计算并保存指数数组 E(同时归约求和),最后读 E、归一化并写 P。因此是 3 读 2 写,约 $D=5bnd$;这里假设“保存 E 与求和”已融合,并不是第 4.2 节 NumPy 各语句的完整流量。若整行能驻留片上存储,融合实现可约 1 读 1 写,二者的理想流量比为 2.5。 在线归一化原论文 Algorithm 2的 Safe softmax 不保存 E,输出时重新计算 exp,因此是 3 读 1 写;Algorithm 3 的 online 版本先流式求归一化量,再重读输入写概率,为 2 读 1 写,其对照是 4/3 的流量比。它们与本文保存 E 的基线不同,也与 Triton 整行融合示例的驻留前提不同,不能混用这些比值。 本文与 roofline_model.softmax 对保存 E 或整行驻留的模型统一用约 $5nd$ 作为教学操作计数(max、减法、exp、求和、除法各按一次,忽略每行的小量);它不是五类操作都拥有相同吞吐的物理断言。max、exp 与除法并非都能跑在 Tensor Core 上;重新计算 exp 的 Safe 版本、online 递推及实际分块内核都需另数操作。 单头注意力($N$ 是 token 数,$D_{h}$ 是每头维度):两次 $N \times N \times D_{h}$ 的矩阵乘加 softmax: $$F_{\text{attn}} = 4N^{2}D_{h} + 5N^{2}$$ 朴素实现的分数矩阵 $S$ 和概率矩阵 $P$ 都要落显存,二次项访存约 $4N^{2}$;FlashAttention 原论文 Algorithm 1 与 IO 分析只让当前分数/概率瓦片驻留片上,不保存完整 $N^{2}$ 矩阵;若仅统计每份输入读取一次与输出写回一次,可得到不可避免的数据流量下界;真实分块内核会反复读入 K/V 或 Q 等数据: $$D_{\text{朴素}} = b\,(4N^{2} + 4ND_{h}),\qquad D_{\text{flash,min}} = b \cdot 4ND_{h}$$ 代入 $N=4096$、$D_{h}=128$、bf16(脚本 roofline_model.py 实算):朴素 $D = 138.41\ \text{MB}$、$I = 62.7$,落在 A100($I^{\ast} = 200.6$)的带宽受限区;理想最低流量 $D_{\mathrm{flash,min}} = 4.19\ \text{MB}$、$I = 2068$,跳进算力受限区。教学函数令两种路径使用相同操作计数(×1.000),参考模型从 0.089 ms 降到 0.028 ms(×3.2);真实 online softmax 有额外归约/重缩放操作,不能据此说实际 FLOPs 完全不变。这是按最低流量与假想统一吞吐计算的理想示例,33 倍不是实际 FlashAttention 访存或速度的测量;有限片上存储下还需分块 IO 模型。FlashAttention 那篇的完整递推在知识树的下一节点展开,这里先用 roofline 把它的收益定位清楚。 图 1:这张图要看三样——散点的横坐标是各算子的算术强度 $I$,点越靠右越「算得过来」;按 I 与 I 的相对位置给点作模型分类;点到经验参考折线的距离并不能单独证明真实瓶颈;FlashAttention 会改变访存量和横坐标,无法直接把这段垂直差距视为它的加速收益。* 3.4 模型估计随资源改变的比值 把算力翻倍($P_{\text{peak}} \to 2P_{\text{peak}}$),加速比是: $$S_{\text{算力}} = \frac{\max(F/P_{\text{peak}},\ D/\beta)}{\max(F/(2P_{\text{peak}}),\ D/\beta)}$$ 若模型流量项主导,分子分母都是 $D/\beta$,$S_{\text{算力}} = 1$。算力受限时加速至多为 2;当 $1\lt I/I^{\ast}\lt 2$ 时,算力翻倍会遇到带宽上限,加速为 $I/I^{\ast}$。带宽翻倍对称地反一次。roofline_model.py 第四节用统一教学计数代入 $P=312$ TFLOP/s、$\beta=1555$ GB/s,下面是该假想二资源模型的结果;只有 BF16 dense GEMM 的计算峰值与 A100 规格直接匹配,其他行仍需 3.6 的资源区分: 算子 I/I* 算力×2 的加速 带宽×2 的加速 逐元素 add [16M] 0.001 1.00x 2.00x LayerNorm [4096, 3072] 0.007 1.00x 2.00x softmax 朴素三遍 [4096, 4096] 0.002 1.00x 2.00x GEMM 512x4096x4096 2.041 2.00x 1.00x attention 朴素 [N=4096,D=128] 0.312 1.00x 2.00x attention ideal-min [N=4096,D=128] 10.307 2.00x 1.00x 这些是模型估计时间的比值,不是对真实优化加速比的无条件上限。一个实现若同时消除了调度开销、缓存重读或中间分配,实际加速可以大于单独改变 P 或 beta 的模型比值;反过来,开销不降也可能达不到该比值。图 2 用于提出测量假设,不能代替 profile。 图 2:这张图要看什么——横轴 $I/I^{\ast}=1$ 那条竖线就是分界线:线左边算力翻倍的理想收益为 1,带宽翻倍收益在 1 到 2 之间;只有 I/I≤0.5 时完整得到 2 倍。线右边对称,I/I≥2 才完整得到算力翻倍的 2 倍。投入硬件或投入算子融合之前,先看自己在哪一侧。 3.5 显存的容量账 速度账之外另有一本容量账。下面以未分块 prefill、同时处理 BS 个 token 为例,显存分四项: $$M_{\text{total}} = N_{\text{params}} \cdot b_{w} + 2BSLd_{\text{kv}}b_{\text{kv}} + N_{\text{active}}a + \rho M_{\text{sum}}$$ 逐项说物理含义:$N_{\text{params}}$ 是参数量、$b_{w}$ 是每参数字节数(fp16 为 2)——这一项与并发无关,是常数;$B$ 是并发序列数、$S$ 是序列长度、$L$ 是层数、$d_{\text{kv}}$ 是每层 KV 总维度(GQA 模型用实际的 KV head 数乘头维度)——KV cache 每个 token 每层都要存一份 K 和一份 V,所以是 $2BSLd_{\text{kv}}b_{\text{kv}}$ 字节;$b_{\text{kv}}$ 是 KV 的每元素字节数,与权重 dtype 独立;KV 总容量随并发线性增长;$a$ 是每 token 的激活峰值(推理不保留整层中间结果,但当前层十几份临时张量要同时活着,$d_{\text{model}}=4096$、fp16 时假设同时存活 16 份这种张量得到 128 KiB/token;这是容量练习的显式假设,未做模型显存峰值测量,不能视为经验通用常数);$\rho$ 是额外开销率,$M_{\text{sum}}$ 是前三项之和。这里 $N_{\text{active}}$ 是同时处理的活跃 token 数:未分块 prefill 取 BS,逐 token decode 通常取 B,分块 prefill 取本轮实际 chunk token 数。KV cache 仍按已保存上下文 token 数计。变长序列应把 BS 换成各序列实际长度之和;权重 fp16 并不要求 KV 也用 fp16,代码显式支持 kv_b 和 active_tokens。第 04 节把未分块 prefill 的教学账本代进一个 7B 模型。 3.6 同一 kernel 的运算可能不走同一资源 把不同工作量分开:$F_{\text{tc}}$ 为 Tensor Core 的乘加 FLOP,$F_{\text{scalar}}$ 为普通浮点管线 FLOP,$N_{\exp}$ 为指数次数,$D_{\ell}$ 为存储层级 $\ell$ 的实测字节数。分别匹配峰值 $P_{\text{tc}}$、$P_{\text{scalar}}$、指数吞吐 $R_{\exp}$ 和该层带宽 $\beta_{\ell}$,仍可逐项取最大: $$T\ge\max\left(\frac{F_{\text{tc}}}{P_{\text{tc}}},\frac{F_{\text{scalar}}}{P_{\text{scalar}}},\frac{N_{\exp}}{R_{\exp}},\max_{\ell}\frac{D_{\ell}}{\beta_{\ell}}\right)$$ 这是更细的乐观下界;如果操作有依赖或共享执行资源,还可能需要把某些时间相加。没有 $R_{\exp}$ 实测或明确规格时,就不能从 $5nd$ 推出 softmax 的硬件计算时间。Hierarchical Roofline 原文解释为什么每一级存储都需要自己的流量与带宽,不能拿逻辑数组流量直接配 DRAM 带宽。 一个纯算术反例(单位为教学单位,非任何GPU规格):Tensor工作量10、吞吐100,标量工作量90、吞吐10,搬100字节、带宽100。误把100次操作全除以100得到计算时间1;分资源后下界为 max(0.1,9,1)=9。resource_time 实跑验证了这个9倍差距。这说明“运算数少”与“该类资源能做多快”必须一起定义。 04. 代码实现 保留原来的三个数值脚本与四幅图,并增加可独立运行的 profiler_probe.py。数值实验只需 Python 3.10+ 与 NumPy;图表另外需要 matplotlib,Profiler 对照另需 PyTorch(本次 2.10.0+cpu)。依次运行 machine_probe.py、roofline_model.py、memory_ledger.py、profiler_probe.py、make_figures.py,文末附录包含完整代码。 原稿的大工作集默认占用较多内存,尤其十亿元素 add 还曾通过 float64 临时数组初始化。本版默认用有界工作集;machine_probe.py --large 和 memory_ledger.py --large 保留原 32M 向量、4096 方阵与 64 MiB batch 权重实验,运行前需确认资源。默认模式不能证明 DRAM 流量、GPU 内核瓶颈或服务吞吐。 本次快照时间 2026-10-10 09:12:57;Windows / AMD64、NumPy 2.4.6、OpenBLAS、fp32,BLAS/OpenMP线程请求值均为 1(只设置本次进程)。经验参考 $\beta=55.74$ GB/s、$P=151.75$ GFLOP/s、$I^{\ast}=2.72$ FLOP/Byte。它们是所选 shape 的最短样本结果,不是设备峰值。 4.1 标定两个经验参考 add 标定用 200 万个 fp32 元素、三份数组,总工作集 24 MB,累计每轮逻辑读写也是 24 MB;是否触及 DRAM,取决于缓存层级与系统状态。方阵乘用 1024×1024,三份数组工作集 12 MiB,$F=2n^{3}$,不能将这个 shape 的经验参考视作所有矩阵形状的可达峰值。 预热一次后,add 测 8 次,GEMM 测 5 次;代码同时保存全部样本、min、median、max。参考折线采用最短时间,以观察接近最佳状态的实现;端到端延迟应另测典型请求、并发和分位数,不能套用这个 min。 add:最短 0.431 ms,中位数 0.612 ms GEMM:最短 14.152 ms,中位数 14.367 ms 经验 P = 151.75 GFLOP/s,add 逻辑字节率 beta = 55.74 GB/s,I* = 2.72 比值超过 100% 不代表违反物理定律:逻辑 D 未用硬件计数器测量,而基线 P/beta 只来自一个 shape 和访问模式。本机算子若运行得比参考更好,也会超过这条经验折线。 4.2 把真实算子打上 roofline 关键测试对象是 matmul-softmax-matmul 微基准;此处省略 1/√D 缩放,不是完整的生产 attention。下面按 NumPy 语句累计逻辑读写:右侧 1~10 是分数/概率矩阵的累计访问遍数;它不包含 Q/K/V/O 的线性项,也不是十次完整 DRAM 扫描的测量。实际 BLAS 分块、CPU缓存、write allocate 都可能改变物理流量: def attention_naive(): np.matmul(q, kt, out=S) # 写 S 1 遍 np.max(S, axis=-1, keepdims=True, out=rowmax) # 读 S 2 遍 np.subtract(S, rowmax, out=S) # 读写 S 4 遍 np.exp(S, out=S) # 读写 S 6 遍 np.sum(S, axis=-1, keepdims=True, out=rowsum) # 读 S 7 遍 np.divide(S, rowsum, out=S) # 读写 S 9 遍 np.matmul(S, v, out=O) # 读 S 写 O 10 遍 本次默认模式的真实输出(经验 $I^{\ast}=2.72$;“模型位置”只指该参考折线): 算子 | I | 模型位置 | 估计 ms | 最短 ms | 中位 ms | 参考比值 逐元素 add(1 遍) | 0.08 | 带宽 | 0.431 | 0.400 | 0.511 | 108% add+relu 两遍(预分配) | 0.10 | 带宽 | 0.718 | 1.032 | 1.465 | 70% add+relu 写成一行(有中间数组) | 0.10 | 带宽 | 0.718 | 2.530 | 2.970 | 28% LayerNorm [1024,1024] 预分配 | 0.15 | 带宽 | 0.753 | 1.261 | 1.290 | 60% attention 朴素 [N=512,D=128] | 11.75 | 算力 | 0.893 | 1.191 | 1.211 | 75% GEMM 128x1024x1024 | 51.20 | 算力 | 1.769 | 1.978 | 2.116 | 89% GEMM 512x1024x1024 | 128.00 | 算力 | 7.076 | 7.020 | 7.320 | 101% GEMM 1024x1024x1024 | 170.67 | 算力 | 14.152 | 14.017 | 14.306 | 101% 逐元素、LayerNorm 与注意力的 exp/max 被各按一个教学操作计数,因此表中的参考 P 只可作粗略标尺,不能认为这些指令都能达到 GEMM 吞吐。分配与归约是否主导,还需更细的 profile。 本次朴素 attention 的模型位置是“算力”,参考比值 75%;原稿 arm64、大 shape 快照的位置曾是“带宽”。这正好构成反例:同一个算法换 shape、BLAS、线程和经验参考,模型分类便可能改变;仅凭这张总计时表无法确证哪级缓存、哪类指令才是瓶颈。 图 1 画全部 8 行,图文使用同一快照。 4.3 显存账本与 batch 拐点 未分块 prefill 的容量教学账代一个 7B 模型(32 层、$d_{\text{kv}} = 4096$、权重和 KV 都取 fp16): 模型:7B,fp16 权重 = 13.0 GiB,32 层,d_kv = 4096 KV cache 单价:512 KiB / token (一条 4096 长的序列 = 2.00 GiB) batch 权重 KV cache 激活 额外开销 合计 假设 78 GiB 可用预算 1 13.0G 2.00G 0.50G 4.66G 20.20G 装得下 8 13.0G 16.00G 4.00G 9.91G 42.95G 装得下 16 13.0G 32.00G 8.00G 15.91G 68.95G 装得下 24 13.0G 48.00G 12.00G 21.91G 94.95G OOM 32 13.0G 64.00G 16.00G 27.91G 120.95G OOM 两个读数:权重是常数,batch 再大都是 13.0 GiB;KV cache 按 512 KiB/token 线性涨,batch 16 时(32 GiB)已经是权重的 2.5 倍。此处 30% 是人为设定的额外开销/前三项总量比率,不是实测额外开销率;它也不同于 PagedAttention 论文中 KV cache 已分配容量的浪费比例。同一张卡只把本教学模型的额外开销率改成 4%,batch 24 从 94.95 GiB 降到 75.96 GiB,从 OOM 变装得下:不换卡、不改模型,这个假设账本跨过一个 batch 档位;真实分页收益需只对可优化的 KV 分配项建模,并检查实际可用显存。例如 batch=24 若只把“KV 额外开销/有效KV”从 30% 改为 4%,而权重、激活不加该比例,合计从 87.44 GiB 降到 74.96 GiB;这是另一个显式教学账本,仍不是 PagedAttention 实测复现。 容量账决定“能开多大”,速度账用于比较吞吐与延迟。默认 batch 微基准固定 1024×1024 的 fp32 权重,只有 4 MiB;原稿 4096×4096、64 MiB 可通过 --large 恢复。本次每个 B 预热一次后测 15 次,吞吐是 B 除以最短 GEMM 时间,不包含请求排队、前后处理或完整模型: B | F MFLOP | D MB | I | 模型位置 | 最短 ms | 中位 ms | 吞吐 K token/s 1 | 2.097 | 4.202 | 0.50 | 带宽 | 0.059 | 0.065 | 16.9 2 | 4.194 | 4.211 | 1.00 | 带宽 | 0.382 | 0.391 | 5.2 4 | 8.389 | 4.227 | 1.98 | 带宽 | 0.346 | 0.347 | 11.6 8 | 16.777 | 4.260 | 3.94 | 算力 | 0.406 | 0.416 | 19.7 16 | 33.554 | 4.325 | 7.76 | 算力 | 0.502 | 0.512 | 31.9 32 | 67.109 | 4.456 | 15.06 | 算力 | 0.711 | 0.802 | 45.0 64 | 134.218 | 4.719 | 28.44 | 算力 | 1.107 | 1.138 | 57.8 128 | 268.435 | 5.243 | 51.20 | 算力 | 1.939 | 2.043 | 66.0 256 | 536.871 | 6.291 | 85.33 | 算力 | 3.604 | 3.722 | 71.0 B=1 到 256,最短延迟增长 61.0 倍,GEMM-only 吞吐增长 4.2 倍;没有理由预设“batch 越大延迟就一定平坦”。B=1 的参考比值 128% 也证明经验折线不是硬件上界。4 MiB 权重可能被缓存复用,D 表只记逻辑流量,不能直接把 D/时间当成 DRAM 带宽。 另作两个容量反例:在 batch=24、4096上下文下,把 KV 从 2 字节改成 1 字节,KV 项从 48 GiB 降至 24 GiB;仅改变权重 dtype 并不会自动改变 KV dtype。相同 KV 上下文做逐 token decode,激活项若按 128 KiB/活跃token 的假设计数,则从 prefill 的 12 GiB 变为约 0.00293 GiB;实际 attention workspace、分配器和运行库仍要另计。memory_ledger.ledger 支持这两种参数,见 09 节。 图 3:这张图要看什么——左边是实测延迟,右边是实测吞吐,曲线不保证出现理想平坦段;竖虚线是首个跨过经验 ridge point 的采样 B,模型边界在它与前一采样值之间;延迟变化可能有多个原因,不能从该图猜测具体库路径。 图 4:这张图要看什么——左图是显存四项的堆叠:蓝色权重是常数,红色 KV cache 随并发线性涨,越过黑色可用线就是 OOM;右图是只改一个参数(额外开销率 30%→4%)的效果,batch 24 从 OOM 变装得下。 05. 工业级实现对照 PyTorch v2.10.0 torch/profiler/profiler.py#profile是本次核验的固定源码版本。附录 profiler_probe.py 定义了完整 x/w,CPU 也能运行;若 CUDA 可用,才加 CUDA activity 并在独立计时前后同步。预热与延迟测量放在采集区外,减少把初始化和采集成本当作 kernel 时间的风险。它和本文账本的对应关系是: key_averages 聚合算子事件,不等于直接聚合 GPU kernel;嵌套事件的总时间可能重复包含子调用,不能把所有行的 total time 简单相加。CPU 本次按 self_cpu_time_total 排序;CUDA 分支按 self_device_time_total 排序,该分支尚未在本次 CPU 机器实测。 with_flops=True 在该版本只估计支持的矩阵乘与二维卷积等算子;本次 aten::mm 报 524288,与 $2\times16\times128\times128$ 一致,输出 shape 为 [16,128]。它不是全模型完整 FLOPs 计数,也不是实测指令数。 profile_memory=True 跟踪张量分配/释放。本次 mm 的 self CPU memory 为 8192 字节,恰好等于输出大小;这是该事件的分配增量,不是 CPU RSS 或设备显存峰值,也不涵盖预先创建的权重、KV、allocator reserve 与所有 workspace。 record_shapes 等开关会增加开销并保留张量引用,可能影响优化;定位后应关闭采集单独计时。代码导出 Chrome trace,后续才能分析热点、子事件与时间线。 PaLM §4.1 与 Appendix B定义的训练 MFU 排除激活重算 FLOPs,只计模型所需前向与反向运算。使用 MFU(Model FLOPs Utilization)时应先定义模型理论运算量、时间窗口、卡数与匹配硬件峰值,典型形式为 $\mathrm{MFU}=F_{\text{model}}/(T_{\text{wall}}N_{\text{gpu}}P_{\text{peak}})$。本文不从这一次 mm 采集计算模型 MFU;分布式通信、并发重叠和计时范围都可能改变分母。 Nsight 工具的分工:Nsight Systems 官方 User Guide对应 kernel 空隙、同步和通信时间线;Nsight Compute 官方 Profiling Guide给热点 kernel 的指令吞吐、缓存层级、occupancy 与 stall 等指标。Speed Of Light 的两个百分比是利用率摘要,不能只比较它们便确证瓶颈。先用时间线选热点,再用与该 kernel 指令/存储层级相匹配的计数器检验模型。本次没有 NVIDIA GPU,不宣称完成 nsys/ncu 实测。 SDPA 会按输入和硬件选择后端,不能断言所有 diffusers/transformers 默认用某个 FlashAttention 实现。本文 2068 个教学操作/Byte 来自理想一次读写与统一计数,真实 FlashAttention 仍有分块重读、标量归约和有限片上容量约束。 06. 代价与边界 roofline 是模型,模型有假设。四条主要假设和不成立时的样子: 把下界当时间近似时,假设计算与访存充分重叠。真实 kernel 在算和搬之间来回切换,还有 launch 和同步的开销。小工作集与短 kernel 常会受这些开销影响,但总计时不足以定量区分每一项,本次 B=1 结果也不能单独证明 launch 开销主导。 假设峰值算力是一个数。实际有 fp16/bf16/fp8/稀疏好几档,差 6 倍以上;按 BF16 dense 口径,H100 SXM 相比 A100 40GB SXM 的 ridge point 从约 201 升到 295——在同一 GEMM 逻辑流量与指令口径的模型里,更多强度值会落进带宽侧,拿旧卡的直觉做判断会错。 假设 $D$ 与缓存无关。本文数组账本的 $D$ 是逻辑读写,缓存命中会让逻辑 D/时间大于 DRAM 标称带宽;这并不代表物理DRAM超速,必须先测相应层级流量。锚点论文 Hierarchical Roofline(arXiv:2009.05257)就是把单一 roofline 扩展成每级缓存一条,用于定位数据移动发生在哪一层。 假设算子孤立。decode 阶段 GEMM 的「权重」每层都要重新读一遍,全局账(整个模型、整个请求)和单算子账结论可能不同;通信算子(allreduce)的账本里延迟和消息数占大头,照搬本文公式会算错。 什么时候不用它:动态 shape、算子间强耦合(融合边界在变)、通信密集的分布式场景——这些先上 profiler 看时间线,先选热点与明确测量口径,再用 roofline;分层、分资源模型也可以帮助分析这些复杂负载,只是本文简单二资源式不够。坦诚标注:本文当前数值实测来自 Windows CPU 的 NumPy 与一个 PyTorch CPU 小矩阵;原大工作集形状保留为可选模式,数字本身不可迁移,理想分段趋势也可能被缓存、内核选择和调度打破。 07. 经典论文脉络 Roofline: An Insightful Visual Performance Model for Multicore Architectures(Williams et al., CACM 2009;2008 Berkeley 技术报告版本):提出算术强度与 ridge point,一根折线把「算力受限/带宽受限」变成可判定的题。一切性能建模的原点。 Hierarchical Roofline Performance Analysis for Deep Learning Applications(Yang et al., 2020):把 roofline 按缓存层级展开,回答「多搬的字节发生在哪一级存储」,是本文锚点论文,也补了单一 roofline 在深度学习负载上最大的盲区。 FlashAttention(Dao et al., 2022):IO-aware 的代表作——保持精确注意力的数学结果与主导矩阵乘量级,靠 tiling + online softmax 避免完整 N×N 中间量落 HBM;标量重缩放和重复读取仍有代价,是否改变实际瓶颈须看形状与内核。roofline 视角下「优化带宽」的教科书案例。 Mixed Precision Training(Micikevicius et al., 2017):换 dtype 同时改两本账——$b$ 变小省字节,算力单元换挡提峰值。哪半边有收益取决于受限类型(已发长文专门讲数值稳定那半)。 Efficient Memory Management with PagedAttention(Kwon et al., 2023):在论文比较的服务负载中显著减少 KV cache 分配浪费;这些比例不能直接乘到权重、激活和所有显存上,虚拟内存的分页思想搬进显存管理。知识树里 KV cache 一篇的主角。 五篇连起来是一条线:先有判定工具(roofline),再按受限类型各给一把钥匙——算力侧(混合精度)、带宽侧(FlashAttention)、容量侧(PagedAttention)。 08. 常见误解 “FLOPs 少就更快”。本次开头的约100MFLOP实验已给出反例;它支持“只数运算不够”,不能单独证明DRAM带宽造成了全部时间差。 “显存够放权重就能跑”。未分块 prefill 的显式教学参数给 batch=24 的 94.95 GiB,超过假设78 GiB预算;KV、活跃激活与额外分配都需记账。模型没有实际在GPU上运行,OOM是容量模型预测。 “新卡算力翻倍,推理就必定提速”。仅当匹配的计算资源主导且可利用新增吞吐才有希望;量化还会改变字节、反量化成本和数值误差,不能把它等同于只改P。 “MFU低说明实现差”。小batch、长上下文、通信或排队可能限制吞吐;模型理论FLOPs、每步时间和硬件峰值口径不同,也会改变MFU。先拆测量窗口与资源,才能评价实现。 “图上的带宽分类就是硬件瓶颈证明”。经验GEMM P和逻辑D只构成参考模型;本次4MiB权重可被缓存反复用,不能把表里的D直接当HBM/DRAM流量。换工作集大小、读取硬件计数器和时间线后,再检验假设。 “FlashAttention把整个N²放进SRAM”。它保留瓦片和行统计量;有限片上存储导致分块重读,理想4ND字节只是一次读写的下界。 09. 动手验证 依次执行五个完整脚本,比较快照的最短值、中位数和全部样本;不预设某行参考比值必须低于50%或接近100%。改变工作集、shape与线程设置,每次重新标定并重画四图。若观察到分类或延迟变化,先记录变化,再用工具确认原因。 先运行 machine_probe.py,再用 ledger(..., kv_b=1) 对照 kv_b=2;KV项应分别24/48 GiB。再加 active_tokens=24 比较decode激活,注意它没有改变已保存的KV上下文。 memory_ledger.ledger( 7e9, 2, 32, 4096, 24, 4096, 131072, 0.30, kv_b=1) 在 roofline_model.py 中只把例子的P改为1979e12、保留beta和字节数,观察模型边界右移;这是控制变量练习,不能称为真实H100换卡结果。真实换卡需要同时核对带宽、指令类型、dtype、缓存和内核。 运行 profiler_probe.py,核对 [16,128]×[128,128]、524288 FLOP和8192字节输出,打开trace查看算子事件;CPU结果不能用来证明CUDA分支或GPU显存表现。 10. 延伸阅读 按知识树的依赖关系,从这篇出发有三个方向: 往上游:自注意力机制的计算与显存账本——本文 3.3 节那笔注意力账的完整推导;混合精度与数值稳定性——换 dtype 的另一本账(数值范围与累加精度)。 往下游(受限类型各一把钥匙):FlashAttention(带宽侧)、KV Cache 与自回归视频生成(容量侧)、算子融合与 CUDA Graph(launch/中间量侧)、量化(字节/计算侧)。 目录:完整知识树见博客目录页 AIGC 基本功知识树,按依赖顺序排好了先修课。 继续阅读 FlashAttention 为什么不需要存下注意力矩阵:本文 3.3 节的最低流量示例省略了真实分块重读,那篇把 online softmax 的递推式一步步推出来。 附录:完整代码 09 节用到的脚本全文如下(machine_probe.py、roofline_model.py、profiler_probe.py、memory_ledger.py、make_figures.py)。复制到本地存成同名文件,按各脚本开头的依赖说明准备环境后即可运行。 machine_probe.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """在同一台机器上把 roofline 的「峰值」和「落点」都实测出来。 只依赖 numpy。直接 `python machine_probe.py` 即可运行,约 20 秒。 为什么要自己测一遍 ------------------ roofline_model.py 里的 A100/H100 峰值是厂商标称值,读者手边不一定有 GPU。 本脚本统一记录教学操作数和逻辑数组流量;GEMM乘加按2计数,exp/max各按1。 add有效字节率与GEMM经验算力只构成参考折线,不是所有算子的硬件物理下界。 同一台机器、BLAS和线程设置可以比较本次结果,但仍须硬件计数器才能诊断瓶颈。 每个算子的 D 都按实现里实际的读写遍数来数: LayerNorm 预分配版 10 遍(mean / 减均值 / 平方 / 再 mean / 除 / 乘 gamma) attention 朴素预分配 10 遍(写 S / max / 减 / exp / sum / 除 / 读 S 写 O) 遍数不是拍脑袋来的,是照着下面每一行 numpy 数出来的。 输出的每个数字都来自 `time.perf_counter()` 的真实计时。 """ import argparse import io import os import statistics from contextlib import redirect_stdout import json import platform import time from pathlib import Path import numpy as np FP32 = 4 # 本机跑 fp32,BLAS 后端在快照 environment 中实际记录 _LAST_TIMING = {} SNAPSHOT = Path(__file__).resolve().parent / "probe_results.json" def _best(fn, repeat: int): """预热后取repeat次最短值,另存全部样本和中位数;不代表请求延迟分位数。""" fn() # 预热:第一次要分配 / 触发缺页 samples = [] for _ in range(repeat): t = time.perf_counter() fn() samples.append(time.perf_counter() - t) _LAST_TIMING.clear() _LAST_TIMING.update(samples_seconds=samples, median_seconds=statistics.median(samples), min_seconds=min(samples), max_seconds=max(samples), repeats=repeat) return min(samples) # ── 一、标定可达带宽:二元 add 模式(读两份写一份)─────────────────── def measure_bandwidth(n: int = 2_000_000, repeat: int = 8) -> float: """返回实测可达带宽(Byte/s)。 默认每份 fp32 数组 8 MB;large模式128 MB;应核对目标机器缓存大小,这里测量此访问模式的有效带宽。 z = x + y:读 2n、写 n,共 3n 个元素。 """ x = np.ones(n, dtype=np.float32) y = np.ones(n, dtype=np.float32) z = np.empty(n, dtype=np.float32) dt = _best(lambda: np.add(x, y, out=z), repeat) return 3 * n * FP32 / dt # ── 二、标定峰值算力:大矩阵乘 ──────────────────────────────────── def measure_matmul_peak(n: int = 1024, repeat: int = 5) -> float: """返回实测可达算力(FLOP/s)。 大方阵乘能摊薄访存;测得的是此shape/BLAS/线程数的经验参考,非物理峰值。 FLOPs = 2·n³(口径同 roofline_model.py:一次乘加记 2 个浮点运算)。 """ rng = np.random.default_rng(0) a = rng.standard_normal((n, n), dtype=np.float32) b = rng.standard_normal((n, n), dtype=np.float32) c = np.empty((n, n), dtype=np.float32) dt = _best(lambda: np.matmul(a, b, out=c), repeat) return 2 * n ** 3 / dt def calibrate(verbose: bool = True, large: bool = False): n = 32_000_000 if large else 2_000_000 side = 4096 if large else 1024 bw = measure_bandwidth(n=n) bw_timing = dict(_LAST_TIMING) fl = measure_matmul_peak(n=side) peak = dict(peak_flops=fl, peak_bw=bw, vector_elements=n, matrix_side=side, bandwidth_timing=bw_timing, matmul_timing=dict(_LAST_TIMING)) if verbose: print("=" * 80) print(f"本机标定({platform.machine()},numpy {np.__version__},fp32)") print("=" * 80) print(f" 实测可达带宽 beta = {bw / 1e9:8.2f} GB/s") print(f" 实测可达算力 P = {fl / 1e9:8.2f} GFLOP/s") print(f" ridge point I* = {fl / bw:8.1f} FLOP/Byte") print(" 注:这是该shape下的经验参考,不是所有算子的物理上限;") print(" 换 BLAS 后端、换 dtype、换线程数都会变。口径一致才有可比性。") return peak def build_ops(large: bool = False): """构造待实测的算子。每个元素 = (名字, F, D, 函数, repeat)。""" rng = np.random.default_rng(0) ops = [] # ── 逐元素:1 遍 vs 2 遍 vs「看起来融合了」 ── n = 32_000_000 if large else 2_000_000 a = rng.standard_normal(n, dtype=np.float32) b = rng.standard_normal(n, dtype=np.float32) o1 = np.empty(n, dtype=np.float32) tmp = np.empty(n, dtype=np.float32) def add_relu_two_pass(): # 5n 个元素:tmp 读写各一次 np.add(a, b, out=tmp) np.maximum(tmp, 0, out=o1) def add_relu_one_line(): # 写法像融合了,numpy 照样分配中间数组 np.maximum(np.add(a, b), 0, out=o1) ops.append(("逐元素 add(1 遍)", 1.0 * n, FP32 * 3 * n, lambda: np.add(a, b, out=o1), 8)) ops.append(("add+relu 两遍(预分配)", 2.0 * n, FP32 * 5 * n, add_relu_two_pass, 8)) ops.append(("add+relu 写成一行(有中间数组)", 2.0 * n, FP32 * 5 * n, add_relu_one_line, 8)) # ── LayerNorm:预分配的 10 遍版 ── ln_n, ln_d = (8192, 4096) if large else (1024, 1024) x = rng.standard_normal((ln_n, ln_d), dtype=np.float32) g = rng.standard_normal(ln_d, dtype=np.float32) mu = np.empty((ln_n, 1), dtype=np.float32) var = np.empty((ln_n, 1), dtype=np.float32) xc = np.empty_like(x) xc2 = np.empty_like(x) def layernorm_buffered(): np.mean(x, axis=-1, keepdims=True, out=mu) # 读 x 1 遍 np.subtract(x, mu, out=xc) # 读 x 写 xc 3 遍 np.multiply(xc, xc, out=xc2) # 读 xc 写 xc2 5 遍 np.mean(xc2, axis=-1, keepdims=True, out=var) # 读 xc2 6 遍 np.sqrt(var + 1e-5, out=var) # 小量 np.divide(xc, var, out=xc) # 读 xc 写 xc 8 遍 np.multiply(xc, g, out=xc) # 读 xc 写 xc 10 遍 ops.append((f"LayerNorm [{ln_n},{ln_d}] 预分配", 6.0 * ln_n * ln_d, FP32 * 10 * ln_n * ln_d, layernorm_buffered, 5)) # ── 朴素注意力:S 落盘,10 遍 N² 访存 ── N, D = (2048 if large else 512), 128 q = rng.standard_normal((N, D), dtype=np.float32) k = rng.standard_normal((N, D), dtype=np.float32) v = rng.standard_normal((N, D), dtype=np.float32) kt = np.ascontiguousarray(k.T) S = np.empty((N, N), dtype=np.float32) O = np.empty((N, D), dtype=np.float32) rowmax = np.empty((N, 1), dtype=np.float32) rowsum = np.empty((N, 1), dtype=np.float32) def attention_naive(): # 本微基准省略 1/sqrt(D) 缩放,只测 matmul-softmax-matmul 的执行开销。 np.matmul(q, kt, out=S) # 写 S 1 遍 np.max(S, axis=-1, keepdims=True, out=rowmax) # 读 S 2 遍 np.subtract(S, rowmax, out=S) # 读写 S 4 遍 np.exp(S, out=S) # 读写 S 6 遍 np.sum(S, axis=-1, keepdims=True, out=rowsum) # 读 S 7 遍 np.divide(S, rowsum, out=S) # 读写 S 9 遍 np.matmul(S, v, out=O) # 读 S 写 O 10 遍 ops.append((f"attention 朴素 [N={N},D={D}]", 4.0 * N * N * D + 5.0 * N * N, FP32 * (10 * N * N + 4 * N * D), attention_naive, 5)) # ── GEMM 三个规模 ── side = 4096 if large else 1024 for m in ((512, 2048, 4096) if large else (128, 512, 1024)): ma = rng.standard_normal((m, side), dtype=np.float32) mb = rng.standard_normal((side, side), dtype=np.float32) mc = np.empty((m, side), dtype=np.float32) ops.append((f"GEMM {m}x{side}x{side}", 2.0 * m * side * side, FP32 * (m * side + side * side + m * side), lambda ma=ma, mb=mb, mc=mc: np.matmul(ma, mb, out=mc), 5)) return ops def probe(peak: dict, verbose: bool = True, large: bool = False): ops = build_ops(large=large) I_star = peak["peak_flops"] / peak["peak_bw"] if verbose: print("\n" + "=" * 80) print("三、真实算子打到这张 roofline 上:模型估计 vs 实测") print("=" * 80) hdr = (f"{'算子':<32s}{'I':>9s}{'受限':>8s}{'T估计':>10s}{'T实测':>10s}" f"{'达成率':>9s}{'实测算力':>14s}") print(hdr) print("-" * len(hdr)) rows = [] for name, F, D, fn, repeat in ops: dt = _best(fn, repeat) I = F / D bound = "算力" if I >= I_star else "带宽" t_pred = max(F / peak["peak_flops"], D / peak["peak_bw"]) eff = t_pred / dt rows.append(dict(name=name, I=I, bound=bound, t_pred=t_pred, t_real=dt, eff=eff, flops=F, bytes=D, gflops=F / dt / 1e9, timing=dict(_LAST_TIMING))) if verbose: print(f"{name:<32s}{I:>9.2f}{bound:>8s}{t_pred * 1e3:>8.2f}ms" f"{dt * 1e3:>8.2f}ms{eff * 100:>8.0f}%" f"{F / dt / 1e9:>11.1f} GF/s") if verbose: print("\n 参考比值 = 模型估计 / 实测,回答的是「这个实现离 roofline 还有多远」。") print(" 低达成率可能来自指令、缓存、调度或同步;单凭此比值不能诊断额外访存。") return rows def demo_same_flops(): """约100MFLOP,有限工作集重复访问对比,不是DRAM压力测试。""" print("\n" + "=" * 80) print("四、约100MFLOP,两种算子在有限工作集上差多少时间") print("=" * 80) rng = np.random.default_rng(3) target = 1.0e8 # GEMM:2·M·K·N = target,取 K=N=1024 M = int(target / (2 * 1024 * 1024)) ma = rng.standard_normal((M, 1024)).astype(np.float32) mb = rng.standard_normal((1024, 1024)).astype(np.float32) mc = np.empty((M, 1024), dtype=np.float32) t_gemm = _best(lambda: np.matmul(ma, mb, out=mc), 10) f_gemm = 2.0 * M * 1024 * 1024 # 逐元素:n 个元素各算 1 次,F = n # 复用有限工作集,避免十亿元素及float64临时数组造成多GB瞬时分配。 chunk, rounds = 2_000_000, 50 n_elem = chunk * rounds a = rng.standard_normal(chunk, dtype=np.float32) b = rng.standard_normal(chunk, dtype=np.float32) o = np.empty(chunk, dtype=np.float32) def chunked_add(): for _ in range(rounds): np.add(a, b, out=o) t_elem = _best(chunked_add, 5) print(f" GEMM : F = {f_gemm / 1e9:.3f} GFLOP, " f"D = {FP32 * (M * 1024 + 1024 * 1024 + M * 1024) / 1e6:.1f} MB, " f"实测 {t_gemm * 1e3:.3f} ms({f_gemm / t_gemm / 1e9:.0f} GFLOP/s)") print(f" 逐元素 : F = {n_elem / 1e9:.3f} GFLOP, " f"D = {FP32 * 3 * n_elem / 1e6:.1f} MB, " f"实测 {t_elem * 1e3:.3f} ms({n_elem / t_elem / 1e9:.0f} GFLOP/s)") print(f" → FLOPs 基本相同,耗时差 {t_elem / t_gemm:.0f} 倍。") print(" 同样运算量不代表同样时间:这里逐元素算子的运算数并未") print(f" 减少;逻辑 I={1 / (FP32 * 3):.2f} FLOP/Byte,24MB工作集反复访问,不能确证DRAM瓶颈。") return dict(gemm_flops=f_gemm, gemm_seconds=t_gemm, add_flops=n_elem, add_seconds=t_elem, add_working_set_bytes=3 * chunk * FP32, add_logical_bytes=FP32 * 3 * n_elem, chunk=chunk, rounds=rounds, ratio=t_elem / t_gemm) def main(): parser = argparse.ArgumentParser() parser.add_argument("--large", action="store_true", help="恢复原大shape实验;先确认可用内存") args = parser.parse_args() peak = calibrate(large=args.large) rows = probe(peak, large=args.large) demo = demo_same_flops() config = io.StringIO() with redirect_stdout(config): np.show_config() environment = dict(platform=platform.platform(), numpy=np.__version__, thread_settings={k: os.environ.get(k) for k in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS")}, numpy_config=config.getvalue(), preset="large" if args.large else "bounded") # 把这次标定存成快照:make_figures / memory_ledger 直接读它, # 保证「图上的数字」和「正文引用的数字」来自同一次运行。 SNAPSHOT.write_text(json.dumps( dict(peak=peak, rows=rows, demo=demo, environment=environment, note=f"{platform.machine()} / numpy {np.__version__} / fp32", timestamp=time.strftime("%Y-%m-%d %H:%M:%S")), ensure_ascii=False, indent=1), encoding="utf-8") print(f"\n标定快照已存到 {SNAPSHOT.name},make_figures / memory_ledger 会复用它。") print(f"这台机器的 roofline:P = {peak['peak_flops'] / 1e9:.0f} GFLOP/s, " f"beta = {peak['peak_bw'] / 1e9:.1f} GB/s, " f"I* = {peak['peak_flops'] / peak['peak_bw']:.1f} FLOP/Byte") if __name__ == "__main__": main() roofline_model.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """Roofline 模型的最小实现:把「算力账」和「带宽账」合成一条上界曲线。 只依赖 numpy(其实只是用来排版,算法本身是纯标量运算)。 直接 `python roofline_model.py` 即可运行,约 1 秒。 符号与文章第 03 节一一对应 -------------------------- F 浮点运算数(FLOP)。口径:一次乘加(MAC)记 2 FLOP D 教学逻辑流量(Byte):按数组读写计;并非硬件DRAM计数器 I 算术强度 I = F / D,单位 FLOP/Byte P_peak 峰值算力(FLOP/s) beta 峰值带宽(Byte/s) I_star ridge point = P_peak / beta,只判定所给模型的主导资源 关键结论(代码算完会打印): T >= max(F / P_peak, D / beta) 把 F 降为一半或把 D 降为一半时的模型比值,见 optimize_headroom() """ import json from pathlib import Path def _local_spec(): """本机实测峰值:优先读 machine_probe.py 存下的标定快照。 快照不存在时明确要求先运行 machine_probe.py。 这样 SPECS 里的「本机实测」永远是同一次标定,不会和正文引用的数字分叉。 """ snap = Path(__file__).resolve().parent / "probe_results.json" if snap.exists(): p = json.loads(snap.read_text(encoding="utf-8"))["peak"] return dict(peak_flops=p["peak_flops"], peak_bw=p["peak_bw"]) raise FileNotFoundError("请先运行 machine_probe.py 生成本机快照;不能用别人的峰值兜底") # ── BF16/FP8 dense GEMM 厂商峰值;不能视为逐元素/exp吞吐。 # A100 40GB SXM: https://images.nvidia.com/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf # H100 SXM: https://www.nvidia.com/en-us/data-center/h100/ (带*稀疏值除2) # 注意 H100 BF16 dense 算力增幅大于带宽增幅,ridge point 比 A100 高: # 带宽涨 2.15 倍、bf16 算力涨约 3.17 倍 → 更多 kernel 落在带宽受限区。 SPECS = { "A100-40GB SXM (bf16)": dict(peak_flops=312e12, peak_bw=1555e9), "H100-80GB SXM (bf16)": dict(peak_flops=989e12, peak_bw=3350e9), "H100-80GB SXM (fp8)": dict(peak_flops=1979e12, peak_bw=3350e9), # 下面这一条不是标称值,是 machine_probe.py 在同一台机器上实测出来的 "本机实测(见 machine_probe.py)": _local_spec(), } BF16 = 2 # bytes per element def roofline_time(F: float, D: float, peak_flops: float, peak_bw: float): """roofline 时间下界:算力时间与带宽时间取 max(假设两者完美重叠)。""" t_compute = F / peak_flops t_memory = D / peak_bw return max(t_compute, t_memory), t_compute, t_memory def resource_time(terms, moved_bytes, bandwidth): """按不同资源给下界,terms 为 [(工作量, 同口径峰值), ...]。 例:Tensor Core FLOP、FP32 CUDA-core FLOP、exp次数分别使用对应吞吐。 只取max仍是假定资源可完美重叠的乐观下界,不能冒充实测耗时。 """ if moved_bytes < 0 or bandwidth <= 0: raise ValueError("字节数非负、带宽为正") if any(work < 0 or rate <= 0 for work, rate in terms): raise ValueError("工作量非负、同口径吞吐为正") return max([moved_bytes / bandwidth] + [work / rate for work, rate in terms]) def bound_of(F: float, D: float, peak_flops: float, peak_bw: float) -> str: I = F / D I_star = peak_flops / peak_bw return "算力受限" if I >= I_star else "带宽受限" def optimize_headroom(F: float, D: float, peak_flops: float, peak_bw: float): """把运算量 F 或搬运量 D 减半,理想模型可快多少? 答案与算术强度有关:带宽受限时只减少运算量的模型收益为 0,算力受限时最多 2 倍。 返回 (F 减半的加速比, D 减半的加速比),分别等价于只将峰值算力或带宽翻倍。 """ def speedup(F2, D2): t0, _, _ = roofline_time(F, D, peak_flops, peak_bw) t1, _, _ = roofline_time(F2, D2, peak_flops, peak_bw) return t0 / t1 return speedup(F / 2, D), speedup(F, D / 2) # ── 几类算子的 F / D 账本 ────────────────────────────────────────── def gemm(M: int, K: int, N: int, b: int = BF16): """矩阵乘 [M,K] @ [K,N] -> [M,N]。读 A、读 B、写 C。""" F = 2.0 * M * K * N D = b * (M * K + K * N + M * N) return F, D def elementwise(n: int, b: int = BF16, n_pass: int = 1, n_read: int = 1): """逐元素算子:读 n_read 份、写 1 份;n_pass 表示这样读写几轮。 二元运算(如 add)是 n_read=2:读两份输入、写一份输出。 """ F = 1.0 * n * n_pass # 每个元素算 1 次 D = b * (n_read + 1) * n * n_pass return F, D def layernorm(n: int, d: int, b: int = BF16): """简化 LayerNorm(无 beta):两次归约、减均值、平方、除标准差、乘 gamma,约 6 次/元素。 访存 = 读 x 一遍 + 写 y 一遍(统计量是 O(n) 的小量,忽略)。""" F = 6.0 * n * d D = b * 2 * n * d return F, D def softmax(n: int, d: int, b: int = BF16, fused: bool = True): """softmax:exp / 减最大值 / 归一化,约 5 次运算/元素。 fused=True 假设整行可驻留片上存储,读 1 遍写 1 遍;不是独立 online normalizer 的通用 IO fused=False 保存 E 的三阶段模型:读X求max;读X写E同时归约和;读E写P,共3读2写 不是原论文 Algorithm 2 的重算 exp、3读1写版本;F=5nd不适用于后者或online递推 """ F = 5.0 * n * d D = b * 2 * n * d if fused else b * (3 * n * d + 2 * n * d) return F, D def attention(N: int, D: int, b: int = BF16, flash: bool = False): """单头注意力。N 是 token 数,D 是每头维度。 flash=False(朴素):分数矩阵 S 和权重 P 都要落 HBM S = Q K^T 写 N^2 P = softmax(S) 读 N^2 写 N^2 O = P V 读 N^2 → 二次项访存 ≈ 4 N^2 b flash=True(强缓存假设下的最低流量,不是真实 FlashAttention IO):Q/K/V 读进来、O 写出去,只保留一块分数/概率瓦片,不分配完整 N^2 矩阵 → 访存 ≈ 4 N D b """ F = 4.0 * N * N * D + 5.0 * N * N # 两个 N×N×D 的矩阵乘 + softmax if flash: D_bytes = b * (3 * N * D + N * D) # 读 Q K V,写 O else: D_bytes = b * (4 * N * N + 4 * N * D) # 上面再加 S/P 的读写 return F, D_bytes def fmt(x: float, unit: str = "") -> str: for u, s in (("T", 1e12), ("G", 1e9), ("M", 1e6), ("K", 1e3)): if x >= s: return f"{x / s:.2f} {u}{unit}" return f"{x:.2f} {unit}" def main(): print("=" * 78) print("一、三组 GEMM 规格与本机经验标定(不是所有算子的统一物理峰值)") print("=" * 78) print(f"{'设备':<28s}{'P_peak':>14s}{'beta':>14s}{'I* = P/beta':>16s}") ridges = {} for name, s in SPECS.items(): I_star = s["peak_flops"] / s["peak_bw"] ridges[name] = (s, I_star) print(f"{name:<28s}{fmt(s['peak_flops'], 'FLOP/s'):>16s}" f"{fmt(s['peak_bw'], 'B/s'):>14s}{I_star:>12.1f} FLOP/B") a100 = SPECS["A100-40GB SXM (bf16)"] h100 = SPECS["H100-80GB SXM (bf16)"] print(f"\n A100 -> H100:带宽 ×{h100['peak_bw'] / a100['peak_bw']:.2f}," f"算力 ×{h100['peak_flops'] / a100['peak_flops']:.2f}," f"ridge point {a100['peak_flops'] / a100['peak_bw']:.0f} -> " f"{h100['peak_flops'] / h100['peak_bw']:.0f}(更高 = 更多算子落入带宽受限区)") # ── 二、算子账本 ──────────────────────────────────────────────── N, D = 4096, 128 # 一个典型的 DiT / LLM 单头规模 cases = [ ("逐元素 add [16M]", elementwise(16_000_000, n_read=2)), ("LayerNorm [4096, 3072]", layernorm(4096, 3072)), ("softmax fused [4096, 4096]", softmax(4096, 4096, fused=True)), ("softmax 朴素三遍 [4096, 4096]", softmax(4096, 4096, fused=False)), ("GEMM 512x4096x4096", gemm(512, 4096, 4096)), ("GEMM 4096x4096x4096", gemm(4096, 4096, 4096)), ("attention 朴素 [N=4096,D=128]", attention(N, D, flash=False)), ("attention ideal-min [N=4096,D=128]", attention(N, D, flash=True)), ] print("\n" + "=" * 78) print("二、统一操作计数的教学模型(b=2,P=312T,beta=1555G)") print("=" * 78) hdr = (f"{'算子':<32s}{'F (教学操作)':>14s}{'D (Byte)':>14s}" f"{'I':>10s}{'受限':>10s}{'T_roof':>12s}") print(hdr) print("-" * len(hdr)) for name, (F, Db) in cases: I = F / Db T, tc, tm = roofline_time(F, Db, a100["peak_flops"], a100["peak_bw"]) print(f"{name:<32s}{fmt(F):>14s}{fmt(Db, 'B'):>14s}" f"{I:>10.2f}{bound_of(F, Db, a100['peak_flops'], a100['peak_bw']):>10s}" f"{T * 1e3:>10.3f} ms") # ── 三、FlashAttention 把算术强度抬了多少 ──────────────────────── print("\n" + "=" * 78) print("三、物化 S/P 与理想最低 IO 对比(单头 N=4096, D=128, bf16)") print("=" * 78) F_naive, D_naive = attention(N, D, flash=False) F_flash, D_flash = attention(N, D, flash=True) T_naive, _, _ = roofline_time(F_naive, D_naive, a100["peak_flops"], a100["peak_bw"]) T_flash, _, _ = roofline_time(F_flash, D_flash, a100["peak_flops"], a100["peak_bw"]) print(f" FLOPs : {fmt(F_naive)} -> {fmt(F_flash)} " f"(×{F_flash / F_naive:.3f},几乎没变)") print(f" 访存 : {fmt(D_naive, 'B')} -> {fmt(D_flash, 'B')} " f"(×{D_flash / D_naive:.4f},省 {D_naive / D_flash:.1f} 倍)") print(f" 算术强度 : {F_naive / D_naive:.1f} -> {F_flash / D_flash:.1f} FLOP/B " f"(×{(F_flash / D_flash) / (F_naive / D_naive):.1f})") print(f" 受限类型 : {bound_of(F_naive, D_naive, **a100)} -> " f"{bound_of(F_flash, D_flash, **a100)}") print(f" roofline : {T_naive * 1e3:.3f} ms -> {T_flash * 1e3:.3f} ms " f"(×{T_naive / T_flash:.2f})") print(f" 注意:模型比值来自「少搬 {D_naive / D_flash:.0f} 倍字节」," "此处给各运算统一计数和统一P,仅为教学模型;真实内核的标量吞吐、exp及重读需另计。") # ── 四、砍算力 vs 砍带宽,谁的收益大 ──────────────────────────── print("\n" + "=" * 78) print("四、教学资源模型中,只翻倍 P 或 beta 的时间估计比值") print("=" * 78) print(f"{'算子':<32s}{'I/I*':>10s}{'算力×2 的加速':>16s}{'带宽×2 的加速':>16s}") I_star = a100["peak_flops"] / a100["peak_bw"] for name, (F, Db) in cases: s_flops, s_bw = optimize_headroom(F, Db, a100["peak_flops"], a100["peak_bw"]) print(f"{name:<32s}{(F / Db) / I_star:>10.3f}" f"{s_flops:>14.2f}x{s_bw:>14.2f}x") print("\n 读法:带宽受限的行(I/I* < 1)里「算力×2」那一列几乎都是 1.00——" "\n 在只改变计算峰值且访存不变的理想模型里,该列为 1。") # ── 五、GEMM 什么时候从带宽受限翻到算力受限 ───────────────────── print("\n" + "=" * 78) print("五、GEMM [M,4096] x [4096,4096]:M 多大才开始算力受限") print("=" * 78) K = Nout = 4096 print(f"{'M':>8s}{'I (FLOP/B)':>14s}{'受限':>12s}") prev = None for M in (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 4096): F, Db = gemm(M, K, Nout) b = bound_of(F, Db, a100["peak_flops"], a100["peak_bw"]) print(f"{M:>8d}{F / Db:>14.2f}{b:>12s}") if prev == "带宽受限" and b == "算力受限": print(f" ↑ 拐点:M 在 ({M // 2}, {M}] 之间," f"I 越过 ridge point {I_star:.0f}") prev = b print("\n 推论:小 batch 推理(M=1~8)里的 GEMM 是带宽受限的," "这时候做 fp8 量化省的是字节、不是算力。") if __name__ == "__main__": main() profiler_probe.py """最小 PyTorch Profiler 对照:无需模型权重;默认CPU,可选CUDA。""" import json import statistics import time from pathlib import Path import torch from torch.profiler import ProfilerActivity, profile, record_function def main(): torch.manual_seed(7) torch.set_num_threads(1) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") x = torch.randn(16, 128, device=device) w = torch.randn(128, 128, device=device) activities = [ProfilerActivity.CPU] if device.type == "cuda": activities.append(ProfilerActivity.CUDA) def synchronize(): if device.type == "cuda": torch.cuda.synchronize() with torch.inference_mode(): for _ in range(5): torch.mm(x, w) synchronize() samples = [] for _ in range(10): synchronize() start = time.perf_counter() torch.mm(x, w) synchronize() samples.append(time.perf_counter() - start) # profiler区域用于观察事件;上面单独计时用于延迟,避免混入采集开销。 with profile(activities=activities, record_shapes=True, profile_memory=True, with_flops=True, acc_events=True) as prof: with record_function("one_matmul"): y = torch.mm(x, w) synchronize() events = prof.key_averages() sort = "self_device_time_total" if device.type == "cuda" else "self_cpu_time_total" print(events.table(sort_by=sort, row_limit=8)) mm = next(event for event in events if event.key == "aten::mm") expected = 2 * 16 * 128 * 128 assert mm.flops == expected assert y.shape == (16, 128) and torch.isfinite(y).all() result = dict(torch_version=torch.__version__, device=str(device), x_shape=list(x.shape), w_shape=list(w.shape), y_shape=list(y.shape), estimated_mm_flops=mm.flops, expected_mm_flops=expected, mm_self_cpu_memory_bytes=mm.self_cpu_memory_usage, y_bytes=y.numel() * y.element_size(), standalone_seconds=samples, median_seconds=statistics.median(samples), event_keys=[event.key for event in events], note="CPU实跑时不宣称CUDA内核验证;算子分配增量不是设备显存总占用") here = Path(__file__).resolve().parent prof.export_chrome_trace(str(here / "profiler_trace.json")) (here / "profiler_results.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(result, ensure_ascii=False, indent=2)) if __name__ == "__main__": main() memory_ledger.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """显存账本 + batch 扫描:容量这本账怎么算,以及它怎么反过来决定速度。 只依赖 numpy(第二部分要真跑计时)。直接 `python memory_ledger.py` 即可运行,约 10 秒。 第一部分是「容量账」:权重 / KV cache / 激活 / 额外开销四项加总,看 假设 78 GiB 可用预算 能塞多少条并发请求。第二部分是「速度账」:batch 变大之后,同一份权重被更多 token 摊薄,算术强度抬上去。I与I*仅标记模型分类交界,计时曲线不一定出现相同拐点。 """ import argparse import statistics import json import time from pathlib import Path import numpy as np GIB = 1024 ** 3 def _load_peak(): """优先读 machine_probe.py 存下的标定快照(保证口径一致); 未运行时明确报错,避免误用另一台机器的峰值。""" snap = Path(__file__).resolve().parent / "probe_results.json" if snap.exists(): peak = json.loads(snap.read_text(encoding="utf-8"))["peak"] return peak["peak_flops"], peak["peak_bw"] raise FileNotFoundError("请先运行 machine_probe.py;不使用其他机器标定兜底") P_PEAK, BETA = _load_peak() I_STAR = P_PEAK / BETA def kv_cache_bytes(n_tokens: int, n_layers: int, d_kv: int, b: int = 2) -> int: """KV cache 字节数:每个 token 每层都要存一份 K 和一份 V。 n_tokens 序列长度(或并发请求的总 token 数) n_layers 层数 L d_kv 每层的 KV 总维度 = n_kv_heads × d_head(GQA 时用实际的 KV head 数) b 每个元素的字节数(fp16/bf16 = 2,fp8 = 1) """ return 2 * n_tokens * n_layers * d_kv * b def ledger(n_params: float, b_w: int, n_layers: int, d_kv: int, batch: int, seq: int, act_per_token: int, overhead_rate: float, kv_b: int = 2, active_tokens: int | None = None): """推理显存账本的四项,单位 GiB。 act_per_token 每个 token 的激活峰值(字节)。推理不保留整层的中间结果, 但当前层的十几份临时张量要同时活着:d_model=4096、fp16 时 一份 [d] 张量 8 KiB,取 16 份 ≈ 128 KiB / token。 这是16份张量的教学假设,不是实测或通用经验值。 overhead_rate 教学额外开销 / (权重 + KV + 激活),不是实测碎片率。 本函数按未分块 prefill 计 batch*seq 个活跃 token; 逐 token decode 的激活应改为 batch*act_per_token,KV 仍取 batch*seq。 """ w = n_params * b_w if min(n_params, b_w, n_layers, d_kv, batch, seq, act_per_token, overhead_rate, kv_b) < 0: raise ValueError("容量输入不得为负") active = batch * seq if active_tokens is None else active_tokens if active < 0: raise ValueError("活跃token数不得为负") kv = kv_cache_bytes(batch * seq, n_layers, d_kv, b=kv_b) act = active * act_per_token frag = (w + kv + act) * overhead_rate return dict(weights=w / GIB, kv=kv / GIB, act=act / GIB, frag=frag / GIB, total=(w + kv + act + frag) / GIB) def part1(): print("=" * 80) print("一、显存账本:一张 假设 78 GiB 可用预算,能塞多少条并发请求") print("=" * 80) # Llama-2-7B 的配置:32 层、32 个 KV head、d_head=128 → d_kv = 4096 N_PARAMS, N_LAYERS, D_KV = 7e9, 32, 4096 SEQ = 4096 per_token = kv_cache_bytes(1, N_LAYERS, D_KV) print(f" 模型:7B,fp16 权重 = {N_PARAMS * 2 / GIB:.1f} GiB," f"{N_LAYERS} 层,d_kv = {D_KV}") print(f" KV cache 单价:{per_token / 1024:.0f} KiB / token " f"(一条 {SEQ} 长的序列 = {per_token * SEQ / GIB:.2f} GiB)") print() hdr = (f"{'batch':>6s}{'权重':>9s}{'KV cache':>10s}{'激活':>9s}" f"{'额外开销':>9s}{'合计':>9s}{'78GiB预算':>10s}") print(hdr) print("-" * len(hdr)) for batch in (1, 4, 8, 16, 24, 32): # 朴素预分配:按最大长度全预留,教学假设:以权重+KV+激活总和的 30% 估计额外开销,不是实测碎片比例 naive = ledger(N_PARAMS, 2, N_LAYERS, D_KV, batch, SEQ, act_per_token=131072, overhead_rate=0.30) ok = "装得下" if naive["total"] < 78 else "OOM" print(f"{batch:>6d}{naive['weights']:>8.1f}G{naive['kv']:>9.2f}G" f"{naive['act']:>8.2f}G{naive['frag']:>8.2f}G" f"{naive['total']:>8.2f}G{ok:>10s}") # 换成低开销假设(教学假设,非 PagedAttention 实测):额外开销率降到 4% print("\n 同一张卡,把额外开销率从 30% 降到 4%(低开销假设):") print(f"{'batch':>6s}{'朴素合计':>11s}{'低开销假设':>11s}{'多出来的并发':>16s}") for batch in (8, 16, 24, 32): naive = ledger(N_PARAMS, 2, N_LAYERS, D_KV, batch, SEQ, 131072, 0.30) low_overhead = ledger(N_PARAMS, 2, N_LAYERS, D_KV, batch, SEQ, 131072, 0.04) extra = "" if naive["total"] > 78 >= low_overhead["total"]: extra = "从 OOM 变装得下" print(f"{batch:>6d}{naive['total']:>10.2f}G{low_overhead['total']:>10.2f}G" f"{extra:>16s}") print("\n 读法:权重那一项是常数,batch 再大也不变;KV cache 才是随并发") print(" 线性增长的那一项,batch 16 时它已经比权重还大。") print(" 「显存够放权重就能跑」这句话漏掉了后面三项。") # 用 roofline_model 里同款口径:下界 = max(F/P, D/beta)。 # P_PEAK / BETA / I_STAR 已在文件顶部从标定快照读出,这里不再覆盖, # 否则快照刷新后正文引用的数字会和实际跑出来的对不上。 def batch_sweep(repeats: int = 15, d: int = 1024): """固定一份「权重」,扫 batch,返回每一档的实测延迟/吞吐/算术强度。""" rng = np.random.default_rng(0) W = rng.standard_normal((d, d), dtype=np.float32) # 默认4 MiB,large模式64 MiB的权重 W_bytes = W.nbytes FP32 = 4 def run(B): x = rng.standard_normal((B, d), dtype=np.float32) y = np.empty((B, d), dtype=np.float32) np.matmul(x, W, out=y) # 预热 samples = [] for _ in range(repeats): t = time.perf_counter() np.matmul(x, W, out=y) samples.append(time.perf_counter() - t) return min(samples), samples rows = [] for B in (1, 2, 4, 8, 16, 32, 64, 128, 256): F = 2.0 * B * d * d D = W_bytes + 2 * B * d * FP32 I = F / D dt, samples = run(B) rows.append(dict(B=B, F=F, D=D, I=I, lat=dt, thr=B / dt, bound="算力" if I >= I_STAR else "带宽", t_pred=max(F / P_PEAK, D / BETA), eff=max(F / P_PEAK, D / BETA) / dt, matrix_side=d, median_seconds=statistics.median(samples), samples_seconds=samples)) return rows, W_bytes def part2(d=1024): """batch 扫描实测:延迟、吞吐、算术强度,以及拐点在哪。""" print("\n" + "=" * 80) print("二、batch 扫描实测:吞吐什么时候不再涨") print("=" * 80) rows, W_bytes = batch_sweep(d=d) (Path(__file__).resolve().parent / "batch_results.json").write_text( json.dumps(dict(rows=rows, weight_bytes=W_bytes, peak=dict(peak_flops=P_PEAK, peak_bw=BETA), matrix_side=d), indent=2), encoding="utf-8") print(f" 固定「权重」W 形状 [{d}, {d}],fp32 = {W_bytes / 2 ** 20:.0f} MiB;" f"batch B 就是一次喂进去的 token 数") print(f" 经验参考(来自 machine_probe 的标定快照):" f"P = {P_PEAK / 1e9:.0f} GFLOP/s,beta = {BETA / 1e9:.1f} GB/s," f"I* = {I_STAR:.1f}") print() hdr = (f"{'B':>6s}{'F (GFLOP)':>12s}{'D (MB)':>10s}{'I':>9s}{'受限':>8s}" f"{'延迟 ms':>10s}{'吞吐 K/s':>11s}{'估计 ms':>10s}{'达成率':>9s}") print(hdr) print("-" * len(hdr)) prev_bound = None for r in rows: print(f"{r['B']:>6d}{r['F'] / 1e9:>12.3f}{r['D'] / 1e6:>10.2f}" f"{r['I']:>9.2f}{r['bound']:>8s}{r['lat'] * 1e3:>10.3f}" f"{r['thr'] / 1e3:>9.1f}K{r['t_pred'] * 1e3:>10.3f}" f"{r['eff'] * 100:>8.0f}%") if prev_bound == "带宽" and r["bound"] == "算力": print(f" ↑ 拐点:B 从 {r['B'] // 2} 到 {r['B']} 之间" f"越过 ridge point {I_STAR:.1f}") prev_bound = r["bound"] # 只标记延迟变化;不从总计时猜测库实现路径或实际硬件瓶颈。 for r0, r1 in zip(rows, rows[1:]): if r1["B"] == 2 * r0["B"] and r1["lat"] > r0["lat"] * 1.5: print(f" ! 反常:B={r0['B']} 只要 {r0['lat'] * 1e3:.2f} ms," f"B={r1['B']} 却要 {r1['lat'] * 1e3:.2f} ms") base = rows[0] top = rows[-1] print(f"\n 延迟:B=1 时 {base['lat'] * 1e3:.3f} ms,B={top['B']} 时 " f"{top['lat'] * 1e3:.3f} ms(涨 {top['lat'] / base['lat']:.1f} 倍)") print(f" 吞吐:B=1 时 {base['thr'] / 1e3:.1f} K/s,B={top['B']} 时 " f"{top['thr'] / 1e3:.1f} K/s(涨 {top['thr'] / base['thr']:.1f} 倍)") print(f" 单样本成本:B={top['B']} 时把 {W_bytes / 2 ** 20:.0f} MiB 权重的逻辑读取摊薄到 {top['B']} 个样本," f"每个样本只摊 {W_bytes / top['B'] / 2 ** 20:.2f} MiB") print(" → 增大 batch 可摊薄权重访问;延迟是否平坦、吞吐能增加多少仍需实测。") print(" → 上面标 ! 的行是 roofline 看不见的东西:同样的公式、更大的 batch," "可能涉及内核选择、调度或计时波动,单凭本表不能确诊。") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--large", action="store_true", help="batch权重恢复4096x4096") args = parser.parse_args() part1() part2(d=4096 if args.large else 1024) make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「性能建模与 Profiling」一文的四张解释图。 图里的计时优先读取已保存的实测快照,避免重绘时图文使用不同测量: * roofline.png 峰值与落点来自 machine_probe.calibrate() / probe() * batch_scaling.png 来自 memory_ledger.batch_sweep() * memory_ledger.png 来自 memory_ledger.ledger() * opt_gain.png 来自 roofline_model.optimize_headroom() 只依赖 numpy + matplotlib,直接 `python make_figures.py` 即可运行,约 40 秒。 改了另外三个脚本,这里要重跑,避免图与正文数字对不上。 """ import json import sys from pathlib import Path import matplotlib.pyplot as plt import numpy as np HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) import machine_probe as MP # noqa: E402 import memory_ledger as ML # noqa: E402 import roofline_model as RM # noqa: E402 OUT = HERE.parent / "figures" try: # exist_ok=True 是必须的:目录已存在时 pathlib 会抛 FileExistsError, # 少数沙箱环境连 exist_ok=True 的 mkdir 也一并拦,这里再兜一层。 OUT.mkdir(exist_ok=True) except PermissionError: if not OUT.is_dir(): raise plt.rcParams.update({ "font.sans-serif": ["Microsoft YaHei", "SimHei", "PingFang SC", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, "savefig.dpi": 160, }) C_MEM = "#e05263" # 带宽受限:红 C_CMP = "#3f7fbf" # 算力受限:蓝 C_ACC = "#2f9e6f" # 好结果:绿 C_GREY = "#94a3b8" # ── 图 1:roofline 曲线 + 实测落点 ──────────────────────────────── def fig_roofline(peak, rows): P = peak["peak_flops"] / 1e9 # GFLOP/s beta = peak["peak_bw"] / 1e9 # GB/s I_star = P / beta # 因为两个都除过 1e9,比值不变 I = np.logspace(-2, 4, 400) perf = np.minimum(P, beta * I) fig, ax = plt.subplots(figsize=(8.4, 5.4)) ax.loglog(I, perf, color="k", lw=2.2) ax.axvline(I_star, color=C_GREY, ls="--", lw=1.4) ax.text(I_star * 1.15, 3, f"ridge point I* = {I_star:.0f}\nFLOP/Byte", color="#475569", fontsize=9, va="bottom") ax.text(0.012, 25, f"GEMM经验参考 P = {P:.0f} GFLOP/s", fontsize=9, color="#475569") ax.text(0.012, P * 3, "斜线 = add经验参考,beta = " f"{beta:.0f} GB/s", fontsize=9, color="#475569") # 落点:用实测时间反推实际性能 short = { "逐元素 add(1 遍)": "逐元素 add", "add+relu 两遍(预分配)": "add+relu 两遍", "add+relu 写成一行(有中间数组)": "add+relu 一行写法", "LayerNorm [8192,4096] 预分配": "LayerNorm", "attention 朴素 [N=2048,D=128]": "朴素 attention", "GEMM 512x4096x4096": "GEMM M=512", "GEMM 2048x4096x4096": "GEMM M=2048", "GEMM 4096x4096x4096": "GEMM M=4096", } offsets = { # 手工调过:右边缘和重叠的标签让位 "GEMM M=4096": (-10, 10), "GEMM M=2048": (10, -4), "GEMM M=512": (12, -14), "朴素 attention": (7, -3), } # 左下角四个点挤在一起,用箭头把标签拉到右边的空地上 arrows = { "逐元素 add": (0.42, 9.5), "LayerNorm": (0.42, 5.6), "add+relu 两遍": (0.42, 3.1), "add+relu 一行写法": (0.42, 1.75), } for r in rows: name = short.get(r["name"], r["name"]) if r["name"].startswith("GEMM "): name = "GEMM M=" + r["name"].split()[1].split("x")[0] elif r["name"].startswith("LayerNorm"): name = "LayerNorm" elif r["name"].startswith("attention"): name = "朴素 attention" y = r["gflops"] ax.scatter([r["I"]], [y], s=52, color=C_MEM if r["bound"] == "带宽" else C_CMP, zorder=5, edgecolor="white", linewidth=0.8) if name.startswith("GEMM M="): m = int(name.split("=")[1]) target = {128: (55, P * 1.8), 512: (280, P * 3.3), 1024: (900, P * 1.8)}.get(m) if target: ax.annotate(name, (r["I"], y), xytext=target, fontsize=8.5, color="#334155", arrowprops=dict(arrowstyle="-", color=C_GREY, lw=0.9)) continue if name in arrows: tx, ty = arrows[name] ax.annotate(name, (r["I"], y), xytext=(tx, ty), fontsize=8.5, color="#334155", arrowprops=dict(arrowstyle="-", color=C_GREY, lw=0.9), va="center") continue dx, dy = offsets.get(name, (7, -3)) ha = "right" if dx < 0 else "left" ax.annotate(name, (r["I"], y), textcoords="offset points", xytext=(dx, dy), fontsize=8.5, color="#334155", ha=ha) ax.set_xlabel("算术强度 I = F / D (FLOP/Byte)") ax.set_ylabel("教学操作计数 / 实测时间 (Gop/s)") ax.set_title("图 1:逻辑流量与经验参考 —— 模型位置不能单独确诊瓶颈", fontsize=11) ax.set_xlim(0.01, 5000) ax.set_ylim(0.3, max(600, P * 6)) ax.grid(alpha=0.25, which="both", ls=":") handles = [plt.Line2D([], [], marker="o", ls="", color=C_MEM, label="模型分类:I < I*"), plt.Line2D([], [], marker="o", ls="", color=C_CMP, label="模型分类:I >= I*")] ax.legend(handles=handles, loc="lower right", fontsize=8.5, framealpha=0.9) fig.tight_layout() fig.savefig(OUT / "roofline.png") plt.close(fig) # ── 图 3:batch 扫描的延迟与吞吐 ────────────────────────────────── def fig_batch(rows): B = [r["B"] for r in rows] lat = [r["lat"] * 1e3 for r in rows] thr = [r["thr"] / 1e3 for r in rows] fig, ax1 = plt.subplots(figsize=(8.4, 5.0)) ax1.plot(B, lat, "o-", color=C_CMP, lw=2, label="延迟(左轴)") ax1.set_xscale("log", base=2) ax1.set_xlabel("batch B(一次喂进去的 token 数)") ax1.set_ylabel("延迟 (ms)", color=C_CMP) ax1.tick_params(axis="y", labelcolor=C_CMP) ax2 = ax1.twinx() ax2.plot(B, thr, "s--", color=C_ACC, lw=2, label="吞吐(右轴)") ax2.set_ylabel("吞吐 (K token/s)", color=C_ACC) ax2.tick_params(axis="y", labelcolor=C_ACC) # 拐点:I 越过 ridge point 的地方 cross = next((r for r in rows if r["bound"] == "算力"), None) if cross: ax1.axvline(cross["B"], color=C_GREY, ls=":", lw=1.5) ax1.text(cross["B"] * 1.05, max(lat) * 0.92, f"越过 ridge point\nB ≈ {cross['B']}", fontsize=9, color="#475569") # 库的反常:延迟不随 batch 单调 b8 = next((r["lat"] for r in rows if r["B"] == 8), None) spike = next((r for r in rows if r["B"] in (2, 4) and b8 is not None and r["lat"] > 1.5 * b8), None) if spike: ax1.annotate(f"B={spike['B']} 延迟为 B=8 的 {spike['lat']/b8:.1f} 倍\n(原因需 profile)", (spike["B"], spike["lat"] * 1e3), textcoords="offset points", xytext=(14, -6), fontsize=8.5, color=C_MEM) ax1.set_title("图 3:batch 扫描 —— 实测延迟、吞吐与 roofline 模型交界", fontsize=11) ax1.grid(alpha=0.25, ls=":") lines1, lab1 = ax1.get_legend_handles_labels() lines2, lab2 = ax2.get_legend_handles_labels() ax1.legend(lines1 + lines2, lab1 + lab2, loc="upper left", fontsize=9) fig.tight_layout() fig.savefig(OUT / "batch_scaling.png") plt.close(fig) # ── 图 4:显存账本 ──────────────────────────────────────────────── def fig_ledger(): batches = [1, 4, 8, 16, 24, 32] naive = [ML.ledger(7e9, 2, 32, 4096, b, 4096, 131072, 0.30) for b in batches] low_overhead = [ML.ledger(7e9, 2, 32, 4096, b, 4096, 131072, 0.04) for b in batches] fig, (ax, ax2) = plt.subplots(1, 2, figsize=(11.2, 4.8), gridspec_kw={"width_ratios": [1.35, 1]}) x = np.arange(len(batches)) keys = [("weights", "权重", "#3f7fbf"), ("kv", "KV cache", "#e05263"), ("act", "激活", "#e8a33d"), ("frag", "额外开销", "#94a3b8")] bottom = np.zeros(len(batches)) for k, label, color in keys: vals = np.array([d[k] for d in naive]) ax.bar(x, vals, bottom=bottom, label=label, color=color, width=0.62, edgecolor="white", linewidth=0.6) bottom += vals ax.axhline(78, color="k", ls="--", lw=1.3) ax.text(len(batches) - 0.4, 79, "教学假设:78 GiB 可用预算", fontsize=8.5, ha="right", color="#334155", bbox=dict(facecolor="white", alpha=0.85, edgecolor="none", pad=1)) for i, d in enumerate(naive): if d["total"] > 78: ax.text(i, d["total"] + 2, "OOM", ha="center", fontsize=9, color=C_MEM, fontweight="bold") ax.set_xticks(x) ax.set_xticklabels(batches) ax.set_xlabel("并发序列数 batch") ax.set_ylabel("显存 (GiB)") ax.set_title("图 4a:未分块 prefill 教学账本(额外开销率 30%)", fontsize=10.5) ax.legend(fontsize=8.5, ncol=2) ax.grid(axis="y", alpha=0.25, ls=":") w_naive = [d["total"] for d in naive] w_low_overhead = [d["total"] for d in low_overhead] ax2.bar(x - 0.2, w_naive, width=0.4, label="额外开销率30%", color="#e05263") ax2.bar(x + 0.2, w_low_overhead, width=0.4, label="低开销假设", color="#2f9e6f") ax2.axhline(78, color="k", ls="--", lw=1.3) ax2.set_xticks(x) ax2.set_xticklabels(batches) ax2.set_xlabel("并发序列数 batch") ax2.set_ylabel("合计显存 (GiB)") ax2.set_title("图 4b:只把额外开销率从 30% 降到 4%", fontsize=10.5) ax2.legend(fontsize=8.5) ax2.grid(axis="y", alpha=0.25, ls=":") fig.suptitle("图 4:显存账本的四项 —— 权重是常数,KV cache 随并发线性增长", fontsize=11) fig.tight_layout() fig.savefig(OUT / "memory_ledger.png") plt.close(fig) # ── 图 2:砍算力 vs 砍带宽,各有多少钱 ──────────────────────────── def fig_opt_gain(): a100 = RM.SPECS["A100-40GB SXM (bf16)"] ratio = np.logspace(-2, 1.2, 300) sp_flops, sp_bw = [], [] for r in ratio: # 造一个算术强度恰好是 r × I* 的算子:固定 D,F 由 r 决定 D = 1e8 F = r * (a100["peak_flops"] / a100["peak_bw"]) * D s_f, s_b = RM.optimize_headroom(F, D, a100["peak_flops"], a100["peak_bw"]) sp_flops.append(s_f) sp_bw.append(s_b) fig, ax = plt.subplots(figsize=(8.4, 4.8)) ax.semilogx(ratio, sp_flops, color=C_CMP, lw=2.2, label="算力翻 2 倍能拿到的加速") ax.semilogx(ratio, sp_bw, color=C_MEM, lw=2.2, label="带宽翻 2 倍能拿到的加速") ax.axvline(1.0, color=C_GREY, ls="--", lw=1.4) ax.text(1.05, 1.02, "I = I*:分界线", fontsize=9, color="#475569") ax.fill_between(ratio, 0.98, 2.02, where=np.array(ratio) < 1, color=C_MEM, alpha=0.07) ax.fill_between(ratio, 0.98, 2.02, where=np.array(ratio) >= 1, color=C_CMP, alpha=0.07) ax.text(0.05, 1.9, "带宽受限区:\n换更快的算力单元 = 0 收益", fontsize=9, color=C_MEM) ax.text(6, 1.9, "算力受限区:\n加带宽 = 0 收益", fontsize=9, color=C_CMP) ax.set_xlabel("算术强度 / ridge point (I / I*)") ax.set_ylabel("教学模型时间估计比值") ax.set_ylim(0.95, 2.1) ax.set_title("图 2:只改变一个资源的教学模型 —— 实际收益需验证", fontsize=11) ax.grid(alpha=0.25, ls=":") ax.legend(fontsize=9, loc="center right") fig.tight_layout() fig.savefig(OUT / "opt_gain.png") plt.close(fig) def main(): # --only <图名>:只重画指定一张(roofline / batch / ledger / gain), # 避免为了改一张图的标签把所有实测算子重跑一遍、数字跟正文引用分叉。 only = None if "--only" in sys.argv: only = sys.argv[sys.argv.index("--only") + 1] # 优先读 machine_probe.py 存的标定快照,保证图上数字与正文引用同源 if MP.SNAPSHOT.exists(): snap = json.loads(MP.SNAPSHOT.read_text(encoding="utf-8")) peak = snap["peak"] print(f"读取标定快照({snap['timestamp']},{snap['note']}):") else: print("没找到标定快照,现场标定一遍…") peak = MP.calibrate(verbose=False) MP.SNAPSHOT.write_text(json.dumps( dict(peak=peak, rows=MP.probe(peak, verbose=False), note="fresh", timestamp=""))) print(f" P = {peak['peak_flops'] / 1e9:.0f} GFLOP/s, " f"beta = {peak['peak_bw'] / 1e9:.1f} GB/s, " f"I* = {peak['peak_flops'] / peak['peak_bw']:.1f}") if only in (None, "roofline"): rows = snap.get("rows") if MP.SNAPSHOT.exists() and "snap" in locals() else None if not rows: rows = MP.probe(peak, verbose=False) fig_roofline(peak, rows) if only in (None, "batch"): batch_snap = HERE / "batch_results.json" if batch_snap.exists(): sweeps = json.loads(batch_snap.read_text(encoding="utf-8"))["rows"] else: sweeps, weight_bytes = ML.batch_sweep() batch_snap.write_text(json.dumps(dict(rows=sweeps, weight_bytes=weight_bytes, peak=peak), indent=2)) fig_batch(sweeps) if only in (None, "ledger"): fig_ledger() if only in (None, "gain"): fig_opt_gain() if only and only not in ("roofline", "batch", "ledger", "gain"): raise ValueError("unknown figure: " + only) print(f"\n图已写入 {OUT}") if __name__ == "__main__": main()
2026年09月25日
7 阅读
0 评论
0 点赞
2026-09-24
AIGC 基本功|自注意力机制的计算与显存账本-MHA
AIGC 基本功|自注意力机制的计算与显存账本-MHA 所属方向:注意力与位置编码 | 难度:入门 | 前置知识:无(这是知识树注意力方向的根节点) 关键词:自注意力、multi-head attention、QKV、复杂度、显存账本、缩放因子 01. 为什么需要它 先看一个真实会撞上的场景:你拿到一个约 3.62B 参数的简化 Transformer(批量 B=1、宽度 d=3072、头数 H=24、层数 L=32,权重和下面的教学张量均按 bf16 两字节),想在 1024×1024 的图生视频任务上训练,并保留各层激活供反向使用。VAE 八倍下采样、patch size 为 2 之后,单帧 latent 是 128×128、patch 网格是 64×64,单帧序列长度 N=4096(视频若联合多个潜帧,N 还要乘潜帧数)。模型权重 bf16 约 6.75 GiB,80 GB 的 A100/H100 看起来还装得下;但若用下面的显式中间量账本、无激活重算地保留 32 层,权重加注意力激活已经超过容量,还没算 FFN、梯度和优化器状态。这里是解释 OOM 风险的教学构造,不是某张 GPU 的峰值测量。 把账摊开看(怎么算出来的见第 03、04 节):一层注意力在 N=4096 时按下述教学账本计为 2.39 GiB 激活,32 层就是 76.5 GiB——是权重的 11 倍;这个乘法还假设这些层的相应中间量同时为反向传播保留。这说明只看权重大小无法判断是否 OOM;在这套训练存储假设下,激活已经超过预算,而且爆的是其中三个特定张量:分数矩阵、softmax 权重、浮点 dropout 乘子,各占 31.4%,三项合计吃掉单层激活的 94%。 按本文简化层结构,N=4096 时两个 N×N 矩阵乘只占整层前向 FLOPs 的 18.2%;它们与线性项之比在 N=18432 达到 1,只看 attention 投影则交叉点为 6144。这是运算量比例,不是运行时间比例;IO、kernel 形状与融合可能让 FLOPs 较少的部分反而更慢。 因此应同时记录 FLOPs、激活存储假设和实际时间线。N=32768 时,本文教学账本的一层存储为 145 GiB,二次项 FLOPs 占 64%;这些值用于理解增长趋势,不能替代实际框架的显存测量。 02. 最小可用理解 三句话讲完核心思想: 每个 token 经三组可学习的线性投影得到 Query、Key、Value;它们不是输入向量的三份原样拷贝。每个 token 再拿自己的 Query 和所有 token(包括自己)的 Key 做内积打分,分数过 softmax 变成权重,再对所有 Value 加权求和,得到这个 token 的新表示。权重由 Q/K、位置编码和掩码共同决定,输出内容还依赖 V。 计算和显存都分两笔:一笔随 N 线性(QKV 投影、输出投影、FFN),一笔随 N 二次(分数矩阵 S、softmax 权重 P、浮点 dropout 乘子各一份 B×H×N×N)。显存爆炸的几乎都是第二笔,这是显式保存中间量的教学实现假设;融合、重算或不使用 dropout 会改变这笔账。 「O(N²) 是瓶颈」有适用条件:二次项与线性项的比值是 N/(6d),占总量的比例是 N/(6d+N),N 小于 2d 时它连 attention 模块内部的一半都不到。在本例朴素存储方案下,N² 项先成为显存大头;实际耗时瓶颈仍需测量。 两 token 先走一遍。 暂取单头 D=1,某个 query 的数值为 1,两枚 key 分别为 1 和 0,两枚 value 分别为 10 和 0。缩放因子 √D=1,分数为 [1, 0],softmax 权重约 [0.731, 0.269],因此输出是 0.731×10+0.269×0≈7.31。若把两枚 key 都改成 1,分数同为 1,权重各 0.5,输出变成 5。这个反例说明:权重由 Q/K 的相对打分决定,而输出还取决于 V;注意力不是「把 Query 复制出来」,也不是看单个分数的绝对大小。 这张图要看什么:多头切的是 d_model 这个维度(H·D 拆成 H 份),不是把注意力复制 H 份;切分本身零算术开销,但分数矩阵从 1 个 N×N 变成 H 个 N×N,显存乘上 H。 03. 数学推导 3.1 单头:打分、归一化、加权求和 设输入序列 $X \in \mathbb{R}^{N \times d}$,N 是 token 数,d 是每个 token 的向量维度(d_model)。三个投影矩阵 $W_q, W_k, W_v \in \mathbb{R}^{d \times d}$ 把每个 token 映射成查询、键、值: $$Q = X W_q, \quad K = X W_k, \quad V = X W_v$$ 每个 token 的 Query 要和所有 token 的 Key 算相似度,写成矩阵形式就是一次 $N \times D$ 对 $D \times N$ 的矩阵乘,得到分数矩阵 $S \in \mathbb{R}^{N \times N}$,其中 $S_{ij}$ 是第 i 个 token 对第 j 个 token 的打分: $$S = \frac{Q K^{\top}}{\sqrt{D}}$$ 本小节是单头,H=1、D=d,所以 Q/K/V 都是 N×D;下一小节切成 H 头后才有 D=d/H。除以 $\sqrt{D}$ 不是装饰,推导一下就知道:为解释尺度,暂假设每个坐标上的 $q_i$ 与 $k_i$ 相互独立、均值为 0、方差为 1,而且不同坐标的乘积不相关,那么点积的方差是 $$\mathrm{Var}(q \cdot k) = \sum_{i=1}^{D} \mathrm{Var}(q_i k_i) = D$$ 点积的标准差随 $\sqrt{D}$ 线性增长。D=128 时分数的摆动幅度是 D=1 的 11 倍,softmax 拿到这么大的输入会直接饱和:最大的那个分数吃掉几乎全部权重,输出逼近 one-hot。接近 one-hot 时,softmax Jacobian 的多数项会很小;有限 logits 的精确 softmax 通常并非严格 one-hot,浮点舍入可能进一步使梯度消失。在这些假设下,除以 $\sqrt{D}$ 把分数方差归一回 1。训练后的 Q/K 未必满足独立与单位方差,同一 token 的 Q 与 K 也可能相关,所以这不是对所有模型 logits 的精确分布定理;它是控制初始化时维度增长效应的动机。第 04 节 D 从 8 扫到 256 的随机实验只检验这个动机。 分数过 softmax 变成权重(每行归一化,行内竞争): $$P_{ij} = \frac{\exp(S_{ij})}{\sum_{j'} \exp(S_{ij'})}$$ 最后对 Value 加权求和得到输出 $O = P V$,形状和输入一样是 $N \times d$。工程实现里 softmax 前要先减去每行最大值再取指数,防止 $\exp$ 上溢——这不改变结果,因为分子分母同乘了一个常数。 3.2 多头:切的是维度,不是份数 把 d 维切成 H 段,每段 D = d/H 维当作一个独立的「头」,各算各的注意力,最后拼回来过一个输出投影: $$\mathrm{MHA}(X) = \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_H)\, W_o, \qquad \mathrm{head}_h = \mathrm{softmax}\!\left(\frac{Q_h K_h^{\top}}{\sqrt{D}}\right) V_h$$ $Q_h$ 是 $Q$ 的第 h 段 D 列。两个常被搞错的点: 固定 d 时,多头不改变两次矩阵乘的主导 FLOPs。H 个头各做 $N^2 D$ 次乘加,总共 $H \cdot N^2 D = N^2 d$,和不切头(一个 D=d 的单头)一样;softmax、调度等开销仍随头数变化。多头改的是「在多少个独立子空间里同时做注意力」,是表达能力的再分配,不是算力的加倍。 多头增加显存。分数矩阵是按头存的:H 个 $N \times N$。显存里 显式分数存储的 N² 项系数是 H 而不是 1。 3.3 算力账:二次项什么时候过半 一层 transformer 的前向 FLOPs(一次 $[M,K] \times [K,N]$ 矩阵乘算 $2MKN$ 个浮点运算): 四个 d×d 投影(Q、K、V 输入投影 + 输出投影):$8 N d^2$ FFN(升维 4d 再降回):$16 N d^2$ 注意力内部两个 N×N 矩阵乘($QK^{\top}$ 与 $PV$):$4 N^2 d$ 线性项合计 $24 N d^2$,二次项是 $4 N^2 d$,比值等于 $N / (6d)$。令比值等于 1: $$4 N^2 d = 24 N d^2 \quad \Longrightarrow \quad N^{\star} = 6d$$ d=3072 时 $N^{\star} = 18432$。如果只看 attention 模块内部(4 个投影对 2 个 N×N 乘),交叉点是 $N^{\star} = 2d = 6144$。你日常跑的 N=4096 在两条线之下——二次项占整层算力 18.2%,占模块内 40.0%。 3.4 显存账:三个 B×H×N×N 下面采用一个刻意显式的教学分配模型:B 是批量、H 是头数、N 是 token 数、d 是总通道宽度;同时计入 X/Q/K/V/ctx/O,以及 S、P 和一份浮点 dropout 乘子,全部按 bf16 两字节估算。它只数某一层注意力中指定张量的字节总和,不是框架峰值,也不是反向传播的最低存储要求;未计权重、梯度、优化器、FFN、临时工作区及分配器开销。bool mask 通常只占一字节,dropout=0 时可省掉该项,S 也通常不必与 P 同时保留。 线性项:X、Q、K、V、加权和、输出,共 6 个 $[B, N, d]$ 张量(不含 FFN 的话); 二次项:分数矩阵 S、softmax 权重 P、浮点 dropout 乘子,各 $[B, H, N, N]$。 $$\text{单层激活} \approx 6 B N d \cdot 2 + 3 B H N^2 \cdot 2 \;\; \text{字节}$$ 两笔相等解得 $N^{\star} = 2d/H$,d=3072、H=24 时 *N=256**——序列长度刚过几百,显存就已经被 N² 项主导了。算力和显存的交叉点差 72 倍(18432 对 256),这就是「算力瓶颈来得晚、显存瓶颈来得早」的定量出处。 这张图要看什么:左边 N=4096 的堆叠条里红色三个格子(S、P、浮点 dropout 乘子)占 94%,线性项挤在边上几乎看不见;右边是对数轴,两条教学模型的比值随 N 近似线性增长,N=32768 时为 129 倍。融合侧统一按六份线性张量 6BNd·bytes 计,忽略小的 LSE 与内核工作区;这不是具体框架的峰值比。 04. 代码实现 完整脚本在文末附录(mha_minimal.py、attention_memory.py、flops_ledger.py、make_figures.py),只依赖 numpy。下面按执行顺序拆核心片段。形状、随机种子下的数值及静态账本已用附录脚本复跑;CPU 耗时随机器和运行条件变化,后面的表保留原稿 Linux 环境的一次样例,并给出本次复跑作对照。 4.1 前向:六行写完 MHA def mha(X, W_q, W_k, W_v, W_o, H, causal=False): B, N, d_model = X.shape D = d_model // H Q = split_heads(X @ W_q, H) # [B, H, N, D] K = split_heads(X @ W_k, H) # [B, H, N, D] V = split_heads(X @ W_v, H) # [B, H, N, D] S = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(D) # [B, H, N, N] if causal: mask = np.triu(np.ones((N, N), dtype=bool), k=1) S = np.where(mask, -np.inf, S) P = softmax(S, axis=-1) # [B, H, N, N] return merge_heads(P @ V) @ W_o, P # [B, N, d_model] 切头函数值得单独看一眼——对这里连续的输入,它是 reshape 加 transpose 的视图变换;非连续输入的 reshape 可能触发拷贝: def split_heads(t, H): B, N, d_model = t.shape D = d_model // H return t.reshape(B, N, H, D).transpose(0, 2, 1, 3) 配置 B=2、N=8、d_model=32、H=4 时各张量的真实形状: [1] 各张量形状 X (2, 8, 32) 输入 Q/K/V (2, 4, 8, 8) 切头后 [B, H, N, D] S (2, 4, 8, 8) 分数矩阵,N×N 是显存爆炸的源头 O (2, 8, 32) 输出,和输入同形 4.2 验证四条性质 脚本跑出来的校验结果: [3] 性质校验 (a) 权重行和为 1 : max|sum(P)-1| = 2.22e-16 (b) 向量化 vs 四重循环 : max|O - O_loop| = 2.22e-16 allclose=True (c) 因果掩码上三角全 0 : max = 0.00e+00 且下三角行和仍为 1 : max|sum-1| = 2.22e-16 (d) 各头权重并不相同 : 头 0 与 H 头平均的平均绝对差 = 0.0223 头之间的平均离散度 : 0.0206(0 表示所有头完全一样) (b) 是最值得做的一次校验:把公式照定义抄成四重循环(一个 token 一个 token 地打分、归一、加权),结果和向量化版本在浮点精度内完全一致。下标搞反、transpose 方向写错这类 bug,靠肉眼很难发现,靠一个慢十倍但显然正确的对照实现能当场抓住。 4.3 缩放因子的实测 D 扫描表(固定 scale=1,看 softmax 最大权重): 固定 scale=1,扫一遍 D 看 softmax 最大权重怎么变: D 分数标准差 最大权重(未缩放) 最大权重(缩放后) 8 2.8587 0.9089 0.2711 16 3.9246 0.9970 0.1718 32 5.6569 1.0000 0.2427 64 8.0203 1.0000 0.1547 128 11.3903 1.0000 0.1322 256 16.0367 1.0000 0.1619 和 3.1 节的推导对上了:分数标准差就是 $\sqrt{D}$(2.8587 ≈ √8,16.0367 ≈ √256)。本次随机样本在 D=16 时最大权重为 0.9970,之后若干行四舍五入显示 1.0000;这说明可能接近饱和,不能证明所有 token 的梯度严格为零。缩放后最大权重回落到 0.13~0.27,分布活着。另外注意主配置(D=8)下的对比:未缩放最大权重 0.5948,缩放后 0.2589——D 小的时候不缩放也能活,缩放用于控制点积随维度增长的方差;这组随机样本不构成某个 D 阈值的通用结论。 4.4 显存账本实跑 attention_memory.py 在 d_model=3072、H=24、B=1、bf16 下逐项清点: [1] 逐项账本 B=1 N=4096 bf16(2 字节) 分数矩阵 S [B, H, N, N] 805,306,368 31.4% softmax 权重 P [B, H, N, N] 805,306,368 31.4% 浮点 dropout 乘子 [B, H, N, N] 805,306,368 31.4% 输入 X / Q / K / V / ctx / O 各 25,165,824 1.0% 合计 2,566,914,048 2.39 GiB → N² 项共 2.25 GiB,线性项共 144.00 MiB,N² 项占 94.1% 随 N 的增长(一层,不含 FFN): [2] N 增长时一层 MHA 的激活显存(B=1, bf16, 含 浮点 dropout 乘子) N 线性项 N² 项 合计 融合侧教学值 倍数 1024 36.00 MiB 144.00 MiB 180.00 MiB 36.00 MiB 5.0x 2048 72.00 MiB 576.00 MiB 648.00 MiB 72.00 MiB 9.0x 4096 144.00 MiB 2.25 GiB 2.39 GiB 144.00 MiB 17.0x 8192 288.00 MiB 9.00 GiB 9.28 GiB 288.00 MiB 33.0x 16384 576.00 MiB 36.00 GiB 36.56 GiB 576.00 MiB 65.0x 32768 1.12 GiB 144.00 GiB 145.12 GiB 1.12 GiB 129.0x 若训练时按同一教学假设保留每层中间量,且不做重算,才可再乘层数:L=32 层时,N=2048 的激活是 20.25 GiB(权重的 3.0 倍),N=8192 是 297 GiB(权重的 44 倍)。推理时不应把当前层临时激活直接乘 L;需要缓存历史的自回归推理另计 KV cache 随 N 和并发数 B 都线性涨:B=32、N=32768、L=80 时光 KV cache 就要 960 GiB,这就是为什么长上下文服务都把 GQA/MLA 当标配。 4.5 算力账本实跑 flops_ledger.py 的占比表(d=3072,含 FFN): N 线性项 二次项 二次项占比 4096 927.71 GFLOPs 206.16 GFLOPs 18.2% 16384 3.71 TFLOPs 3.30 TFLOPs 47.1% 32768 7.42 TFLOPs 13.19 TFLOPs 64.0% 65536 14.84 TFLOPs 52.78 TFLOPs 78.0% 脚本末尾有一段 CPU 微基准。下表是原稿 Linux 环境的一次实测(numpy float32,d_t=1024、16 头);这不是性能保证,也不能和新环境的时延逐项混用: N 投影 (ms) QK^T (ms) 实测比 FLOPs 比 256 1.64 2.36 1.44 0.25 1024 6.64 60.80 9.16 1.00 4096 21.36 1148.15 53.75 4.00 在原稿样例的 N=256 行,QK^T 的算术量只有投影的四分之一,耗时比却是 1.44。2026-10-08 在 Windows 上重跑同一脚本得到投影 3.39 ms、QK^T 12.37 ms,耗时比 3.65;N=4096 则为 27.13 ms 和 1257.98 ms。两组数值不同,恰好说明时延不是 FLOPs 公式能单独预测的。脚本给出的简化算术强度(AI = FLOPs / 估算访存字节)提示了一种原因:投影约 85~228,QK^T 约 21~31,后者还要写出 H·N² 项。但这些是模型估算,不是 CPU 硬件访存计数;一组耗时不足以排除 BLAS 实现、线程数、缓存命中和矩阵布局。应把它当作 profiler 要验证的假说,而不是由时延表证明的因果结论。FlashAttention 原论文则明确以减少 HBM 与片上 SRAM 之间的 IO 为目标,避免把完整 N×N 权重矩阵落到 HBM;它仍需计算全局注意力的二次项。 这张图要看什么:两条曲线是二次项算力占比随 N 的爬升,红蓝两条竖虚线分别是 N=2d=6144(只算 attention 模块)和 N=6d=18432(算上 FFN);你常用的 N=4096 在两条线左边很远的位置。 05. 工业级实现对照 参考实现(以 2026-09 的 main 分支为准,上游重构频繁): huggingface/transformers → modeling_llama.py:LlamaAttention.forward huggingface/transformers → masking_utils.py / integrations/flash_attention.py:attention 实现分发 生产代码和第 04 节的最小实现有五处不一样,每一处都有理由。 5.1 不落地 N×N:eager / sdpa / flash 三条路 HF 的 attention 实现 attn_implementation 有三档: eager:与本文显式计算 S/P 的思路相同;最小代码没有 dropout,也不代表训练账本的全部分配。好处是 P 可访问,代价是二次存储;具体峰值取决于 dtype、存活期与 dropout。 sdpa:调 PyTorch 的 scaled_dot_product_attention,由 PyTorch 按硬件、dtype、mask 等选择 flash、memory-efficient 或 math 后端;不能仅凭 sdpa 名称断言没有 N×N 中间量。 flash_attention_2:在线 softmax + 分块计算,显存 O(N·d),本文统一教学账本在 N=32768 时两者为 129 倍,真实节省比例需实测。 长序列训练通常优先评估可用的融合后端,但要检查硬件、dtype、掩码形式和是否需要显式注意力权重;不能把 sdpa 这个名字等同于必走 Flash kernel。融合实现常不返回完整 P;可视化注意力图或对 P 加自定义正则时,可能需要 eager、专门的导出路径或改变算法,收益与成本都要实测。 5.2 因果掩码不是加 −inf 的稠密矩阵 最小实现里我建了一个 $N \times N$ 的 bool 矩阵,这本身就又是一笔 N² 显存。支持的融合后端可以用 is_causal=True 或滑动窗口等结构化信息在内核内判断位置关系,从而避免显式的 N×N 因果掩码;但 eager、某些 padding/自定义 mask 或后端回退仍可能构造稠密掩码。是否落地要看实际分发路径,不能仅凭 API 名称保证。自己无条件创建 N×N causal mask,会额外消耗显存。 5.3 QKV 合并成一个投影 三个 $d \times d$ 投影合并成一个 $d \times 3d$(或直接 qkv_proj),一次 GEMM 出 Q、K、V。算术量不变,但少起两次 kernel、权重读取更连续。代价是 PyTorch 里要自己 chunk(3, dim=-1) 拆回来——本次核对的 modeling_llama.py 仍保留独立的 q_proj/k_proj/v_proj;融合 QKV 是另一些架构或执行后端的选择。 5.4 KV 头数可以比 Q 头少:GQA 多头切分时 K、V 的头数用 $H_{kv} < H$(比如 8 对 32),多个 Query 头共享一组 KV。公式的改动只是把 $K_h$ 换成 $K_{\lfloor h/(H/H_{kv})\rfloor}$。它不改变 attention 核心两次矩阵乘的主导 FLOPs,但可以减少 K/V 投影的 FLOPs,省的是 KV cache 和 KV 的显存与带宽——推理时 KV cache 缩到 $H_{kv}/H$(LLaMA-3 70B 的 8/64 为 1/8)。这是在固定 Query 宽度时减少 KV 头数、压缩线性 KV 项的优化,而 FlashAttention 优化的是 N² 那一笔,两者正交,经常一起用。 5.5 buffer 化与 position_ids 无位置编码且无位置相关掩码的自注意力对 token 排列是等变的。位置信息可通过正弦表、可学习嵌入、RoPE 或掩码注入,其中 RoPE 的旋转点积体现相对位移,不应统称绝对位置编码。不同实现对 RoPE 的缓存策略并不相同。以 本次核对的 Hugging Face Llama 源码(2026-10-08)为例,LlamaRotaryEmbedding 把逆频率 inv_freq 存为非持久 buffer,而 forward 用 position_ids 计算当次的 cos/sin;不能把它描述成预先缓存整张 cos/sin 表。LlamaModel.forward 在未提供位置时按已缓存 token 数生成 position_ids,外部也能传入位置,例如 packed 序列需要明确处理段内位置。位置如何进入 Q/K,是 RoPE 那一篇的主题。 06. 代价与边界 把 N² 落地换来了什么,又赔了什么。 朴素实现唯一的优点是 S、P 全程可见:可视化注意力图、或使用依赖完整 P 的蒸馏损失,需要访问相应权重。加到 logits 上的结构化 bias 则不必先物化 P;若内核支持其形式,可在分块时应用。导出完整 P 仍需相应的二次输出存储。工程上常见的折中是:训练用 sdpa/flash,分析时用小 N 的 eager 导出注意力图。 二次项的 FLOPs 收益要按序列长度评估。 若各项运行时间恰好与 FLOPs 成正比,N=4096 时消除占比 18.2% 的二次项,整层理想加速约为 1.22 倍;现实中的 IO 和融合会改变时间占比,因此这不是 FlashAttention 的实测加速上限。 多头的账要两头看。 头数 H 越大,每个头的 D = d/H 越小:每头维度 D 改变会影响子空间容量,但不存在这里能够证明的 D<32 通用质量阈值;固定 d=H·D 时,标准 MHA 的 KV cache 与 KV 投影参数量不因增加 H 而线性增长;若固定 D 则另当别论。所以现代模型反而从「H 越多越好」退到「适度头数 + GQA」:LLaMA-3 70B 用 64 个 Query 头配 8 个 KV 头。 什么时候根本不该用全局自注意力。 像素级 self-attention(把 H×W 个像素当 token)在中等分辨率下 N 就上了万,N² 显存直接不可行——这是 latent diffusion 在压缩空间里计算更经济的原因之一;卷积等其他算子的成本也同时下降。高分辨率密集预测里,窗口注意力、局部注意力是常态而不是妥协。 别只优化注意力。 N 小于 2d 时(d=3072 即 N<6144),attention 模块内部的算力大头是投影;FFN 属于整层的另一部分;显存侧倒是早就归 N² 管。所以「 profiling 之前先改结构」是赌博——performance_profiling 那一篇讲的账本方法就是为此准备的。 07. 经典论文脉络 Attention Is All You Need(Vaswani et al., 2017)——本文锚点。把缩放点积注意力 + 多头定型成今天的形态,丢掉循环结构,端到端只靠注意力。 Neural Machine Translation by Jointly Learning to Align and Translate(Bahdanau et al., 2014)——注意力的史前史:注意力最初是翻译里的一组对齐权重,softmax 那一行的「分布」语义就是从这来的。 Generating Long Sequences with Sparse Transformers(Child et al., 2019)——处理长序列 N² 成本的一条代表性稀疏路线:限制每个位置可看的位置集合,减少需要计算和保存的格子;具体节省取决于稀疏模式与实现。 FlashAttention(Dao et al., 2022)——不动公式、只改执行:在线 softmax + 分块,让 N×N 不落地。精确注意力,不是近似——这一点和稀疏/线性路线有本质区别。 GQA: Training Generalized Multi-Query Transformer Models(Ainslie et al., 2023)——把矛头从 N² 转向 KV cache:KV 头数变少、Q 头分组共享,是许多现代模型降低 KV cache 成本的重要设计。 五篇连起来读的线索:注意力先是「一种对齐手段」(2014),再是「唯一的序列算子」(2017),然后 N² 账单到期,在这五篇代表作里,2019 年的路线砍格子,2022 年的路线改执行而保持精确注意力,2023 年的路线则着眼于自回归推理时线性增长的 KV cache。不同成本可以并存,后出现的方案不意味着前一笔账已消失。 08. 常见误解 「注意力是 O(N²),所以序列不长时它也是最慢的部分。」 N=4096、d=3072 时二次项只占本文层结构前向 FLOPs 的 18.2%,但不能据此猜测热点。应结合算术强度、硬件与 profiler 判断。 「多头注意力算 H 遍,所以比单头慢 H 倍。」 H·N²·D = N²·d,多头和全维单头的算术量完全相同;多的是 H 份 N×N 显存和 H 份小矩阵乘的调度开销,不是 H 倍算力。 「显存不够就是模型太大,换个更小的模型。」 d=3072、L=32 的模型权重 6.75 GiB,N=8192 时仅激活就 297 GiB。先算激活账(6B·N·d·bytes + 3B·H·N²·bytes 乘层数),再决定动不动模型。这套无重算训练账本提示应先比较开 gradient checkpointing、换 sdpa/flash、或降分辨率。 「除以 √D 是可要可不要的数值技巧。」 本次样本中未缩放 logits 更容易接近饱和;D=16 的最大权重实际为 0.9970,显示 1.0000 也不等于数学上梯度严格为零。它控制点积方差随维度增长;其他归一化或初始化设计也能缓解饱和,不能据此宣称不缩放就一定无法训练。 「浮点 dropout 乘子不占显存。」 bool mask 通常每元素一字节,而 bf16 为两字节;本文账本计的是两字节的浮点 dropout 乘子,两者不能混称。使用逐元素 dropout 的 eager 训练通常还需保留随机掩码或等价信息;占用取决于表示与实现,不一定恰好是一份 bf16 张量。FlashAttention 可以保存随机数状态并在分块反向时重建掩码。 09. 动手验证 把附录里的 mha_minimal.py 存下来直接跑(python mha_minimal.py),对照三处输出: 形状链:X (2,8,32) → Q/K/V (2,4,8,8) → S (2,4,8,8) → O (2,8,32)。确认分数矩阵是按头存的 4 个 8×8,不是 1 个。 缩放对照:未缩放最大权重 0.5948,缩放后 0.2589;再把 D 扫描表看一遍,本次 D=16 的未缩放值是 0.9970,后续多行显示值接近 1。 等价性:向量化实现与四重循环实现的 max 误差应为 1e-16 量级(你机器上具体数字可能略有不同,但 allclose 一定是 True)。 然后做两个改动观察变化: 把 H 从 4 改成 1 再改成 8,跑性质校验 (d):在 H=1 时头间离散度按定义为 0(没有第二个头);H=8 的数值要以实际重跑为准。这个随机初始化实验只展示不同投影可产生不同权重,不证明训练后各头必然学到互补功能,也不保证头数越多质量越高。 跑 flops_ledger.py 的 CPU 实测段,找到你机器上「实测比」超过「FLOPs 比」的 N 拐点,和理论交叉点 N*=2d 对一下差多少。 预期最容易误判的是第二个:原稿样例在 N=256 的耗时比是 1.44,本次复跑是 3.65,而 FLOPs 比都是 0.25。这只说明两次运行里时延都没有跟着算术量成比例变化;是否由访存、调度或其他因素主导,还需用 profiler 在目标硬件上验证。 10. 延伸阅读 按知识树的依赖顺序,下一步建议这么走: 旋转位置编码 RoPE 的原理与实现——注意力对位置是盲的,RoPE 是视觉/语言模型目前的主流注入方式;本文的 Q、K 在那里会被旋转一次。 视频 DiT 里的 3D RoPE 与分辨率外推——RoPE 在时间/高/宽三组频率上的拆分,视频生成的位置问题。 性能建模与 Profiling:算力、带宽与显存账本——把本文的两本账变成系统方法,roofline 分析。 序列并行——N² 显存装不下时的横向切分方案。 策略梯度与 PPO 基础 等 RL 方向的文章与本篇无直接依赖,可随时穿插。 附录:完整代码 09 节用到的脚本全文如下(mha_minimal.py、attention_memory.py、flops_ledger.py、make_figures.py)。复制到本地存成同名文件,按各脚本开头的依赖说明准备环境后即可运行。 mha_minimal.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """多头自注意力(MHA)的最小可运行实现:把公式逐行翻译成 numpy。 只依赖 numpy,直接 `python mha_minimal.py` 即可运行。 刻意不用 PyTorch:讲原理时框架的抽象反而是噪声,而且 numpy 谁都能跑。 生产实现见文章第 05 节对 transformers 的引用。 公式符号与代码变量名的对应关系 ------------------------------ X 输入序列 形状 [B, N, d_model] W_q/W_k/W_v 三个输入投影矩阵 形状 [d_model, d_model] W_o 输出投影矩阵 形状 [d_model, d_model] Q, K, V 查询 / 键 / 值 形状 [B, H, N, D] S 缩放后的注意力分数 形状 [B, H, N, N],S = Q K^T / sqrt(D) P softmax 后的权重 形状 [B, H, N, N] O 注意力输出 形状 [B, N, d_model] 运行后你会看到:每一层的真实 shape、注意力权重的数值范围、 以及四条性质校验(行和为 1 / 与循环实现等价 / 缩放系数的影响 / 因果掩码)。 """ import numpy as np rng = np.random.default_rng(0) def split_heads(t: np.ndarray, H: int) -> np.ndarray: """[B, N, H*D] -> [B, H, N, D]。 多头不是「复制 H 份再算」,而是把 d_model 这个维度切成 H 段, 每段独立算一次注意力,最后再拼回去。这一步只是 reshape + transpose, 本身不做任何算术,但它决定了后面所有矩阵乘的形状。 """ B, N, d_model = t.shape D = d_model // H return t.reshape(B, N, H, D).transpose(0, 2, 1, 3) def merge_heads(t: np.ndarray) -> np.ndarray: """[B, H, N, D] -> [B, N, H*D],split_heads 的逆操作。""" B, H, N, D = t.shape return t.transpose(0, 2, 1, 3).reshape(B, N, H * D) def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray: """数值稳定的 softmax:先减去最大值再取指数,避免 exp 溢出。""" x = x - x.max(axis=axis, keepdims=True) e = np.exp(x) return e / e.sum(axis=axis, keepdims=True) def mha(X: np.ndarray, W_q, W_k, W_v, W_o, H: int, causal: bool = False): """完整的多头自注意力前向。返回 (输出, 注意力权重 P)。""" B, N, d_model = X.shape D = d_model // H # ① 三个投影:每个 token 独立地把自己映射成 query / key / value Q = split_heads(X @ W_q, H) # [B, H, N, D] K = split_heads(X @ W_k, H) # [B, H, N, D] V = split_heads(X @ W_v, H) # [B, H, N, D] # ② 打分:每个 query 和所有 key 做内积,得到 N×N 的分数矩阵 S = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(D) # [B, H, N, N] # ③ 因果掩码:第 i 个 query 只能看 0..i 的 key if causal: mask = np.triu(np.ones((N, N), dtype=bool), k=1) S = np.where(mask, -np.inf, S) # ④ 归一化成权重,再对 V 做加权求和 P = softmax(S, axis=-1) # [B, H, N, N] ctx = P @ V # [B, H, N, D] # ⑤ 拼回 d_model,过输出投影 O = merge_heads(ctx) @ W_o # [B, N, d_model] return O, P def mha_naive_loop(X, W_q, W_k, W_v, W_o, H): """完全不用矩阵乘的「照着定义抄」版本:四重循环。 用来验证上面的向量化实现没有把下标搞反。逻辑等价但慢得多(O(B·H·N²·D))。 """ B, N, d_model = X.shape D = d_model // H Q = (X @ W_q).reshape(B, N, H, D).transpose(0, 2, 1, 3) K = (X @ W_k).reshape(B, N, H, D).transpose(0, 2, 1, 3) V = (X @ W_v).reshape(B, N, H, D).transpose(0, 2, 1, 3) out = np.zeros((B, H, N, D)) for b in range(B): for h in range(H): for i in range(N): # 先逐个算出这一行 N 个分数,再 softmax,再加权求和 s = np.array([float(Q[b, h, i] @ K[b, h, j]) / np.sqrt(D) for j in range(N)]) p = softmax(s) out[b, h, i] = p @ V[b, h] return merge_heads(out) @ W_o def main(): B, N, d_model, H = 2, 8, 32, 4 D = d_model // H print(f"配置: B={B} N={N} d_model={d_model} H={H} D={D}") X = rng.normal(size=(B, N, d_model)) * 0.5 W_q, W_k, W_v, W_o = (rng.normal(size=(d_model, d_model)) / np.sqrt(d_model) for _ in range(4)) # ── 1. 每一层的真实形状 ────────────────────────────────── Q = split_heads(X @ W_q, H) K = split_heads(X @ W_k, H) S = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(D) print("\n[1] 各张量形状") print(f" X {X.shape} 输入") print(f" Q/K/V {Q.shape} 切头后 [B, H, N, D]") print(f" S {S.shape} 分数矩阵,N×N 是显存爆炸的源头") print(f" O {mha(X, W_q, W_k, W_v, W_o, H)[0].shape} 输出,和输入同形") # ── 2. 分数的量级:为什么要除以 sqrt(D) ────────────────── raw = Q @ K.transpose(0, 1, 3, 2) # 不除以 sqrt(D) print("\n[2] 缩放系数 sqrt(D) 的作用") print(f" D = {D}, sqrt(D) = {np.sqrt(D):.4f}") print(f" 未缩放分数的标准差 : {raw.std():.4f} 分布范围 [{raw.min():.2f}, {raw.max():.2f}]") print(f" 缩放后分数的标准差 : {S.std():.4f} 分布范围 [{S.min():.2f}, {S.max():.2f}]") print(f" 未缩放 softmax 的最大权重: {softmax(raw).max():.4f}") print(f" 缩放后 softmax 的最大权重: {softmax(S).max():.4f}") # D 越大,不缩放的分数方差越大,softmax 越容易塌成 one-hot print(" 固定 scale=1,扫一遍 D 看 softmax 最大权重怎么变:") print(" D 分数标准差 最大权重(未缩放) 最大权重(缩放后)") for D_test in (8, 16, 32, 64, 128, 256): q = rng.normal(size=(256, D_test)) k = rng.normal(size=(256, D_test)) s_raw = q @ k.T s_scaled = s_raw / np.sqrt(D_test) print(f" {D_test:4d} {s_raw.std():9.4f} " f"{softmax(s_raw).max():14.4f} {softmax(s_scaled).max():15.4f}") # ── 3. 性质校验 ──────────────────────────────────────── print("\n[3] 性质校验") O, P = mha(X, W_q, W_k, W_v, W_o, H) print(f" (a) 权重行和为 1 : max|sum(P)-1| = " f"{np.abs(P.sum(axis=-1) - 1).max():.2e}") O_loop = mha_naive_loop(X, W_q, W_k, W_v, W_o, H) print(f" (b) 向量化 vs 四重循环 : max|O - O_loop| = " f"{np.abs(O - O_loop).max():.2e} allclose=" f"{np.allclose(O, O_loop)}") _, P_causal = mha(X, W_q, W_k, W_v, W_o, H, causal=True) upper = np.triu(P_causal, k=1) print(f" (c) 因果掩码上三角全 0 : max = {upper.max():.2e}") print(f" 且下三角行和仍为 1 : max|sum-1| = " f"{np.abs(P_causal.sum(axis=-1) - 1).max():.2e}") # 随机初始化的不同投影产生不同权重;这不是训练后头功能分化的证据 P_mean = P.mean(axis=1) # [B, N, N],把 H 个头的权重平均 diff = np.abs(P[0, 0] - P_mean[0]).mean() spread = np.abs(P[0] - P[0].mean(axis=0, keepdims=True)).mean() print(f" (d) 各头权重并不相同 : 头 0 与 H 头平均的平均绝对差 = {diff:.4f}") print(f" 头之间的平均离散度 : {spread:.4f}(0 表示所有头完全一样)") # ── 4. 不看自己的极端情形:注意力塌缩成「复制」 ────────── print("\n[4] 极端情形:把 K 设成和 Q 完全一样(自注意力且 W_q=W_k)") X2 = rng.normal(size=(1, 4, d_model)) Q2 = split_heads(X2 @ W_q, H) S2 = Q2 @ Q2.transpose(0, 1, 3, 2) / np.sqrt(D) P2 = softmax(S2) diag_share = np.trace(P2[0, 0]) / 4 print(f" 对角线权重占比 = {diag_share:.4f}(1.0 表示每个 token 只关注自己)") if __name__ == "__main__": main() attention_memory.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """自注意力的显存账本:把一层 attention 的每一笔开销都算成字节。 只依赖 numpy(其实只用到它做格式化,核心就是整数乘除)。 直接 `python attention_memory.py` 即可运行。 算的是「训练时一层 MHA 需要为反向传播留下来的激活」, 这是显式保留中间量的保守教学模型,并非框架实测或反向所需最小值。默认的模型尺寸 采用简化 Transformer:d_model=3072、H=24、D=128。 """ import numpy as np BYTES = {"fp32": 4, "fp16": 2, "bf16": 2, "fp8": 1} MIB = 1024 ** 2 GIB = 1024 ** 3 def fmt(nbytes: float) -> str: if nbytes >= GIB: return f"{nbytes / GIB:8.2f} GiB" if nbytes >= MIB: return f"{nbytes / MIB:8.2f} MiB" return f"{nbytes / 1024:8.2f} KiB" def ledger(B: int, N: int, d_model: int, H: int, dtype: str = "bf16", dropout: bool = True): """返回一层 MHA 的逐项激活显存(字节)。 N² 项有三个:分数矩阵 S、softmax 后的 P、以及与所选 dtype 同宽的浮点 dropout 乘子(非 bool mask)。 这三项就是 O(N²) 显存的真身——不是「注意力复杂度是 N²」这句话, 而是本教学模型假设同时保存的三个 B×H×N×N 张量;框架可复用或省略它们。 """ b = BYTES[dtype] linear = B * N * d_model * b # 每个 [B, N, d_model] 的张量 square = B * H * N * N * b # 每个 [B, H, N, N] 的张量 return { "输入 X": linear, "Q": linear, "K": linear, "V": linear, "分数矩阵 S": square, "softmax 权重 P": square, "浮点 dropout 乘子": square if dropout else 0, "加权和 ctx": linear, "输出 O": linear, } def main(): d_model, H, D = 3072, 24, 128 print(f"模型尺寸: d_model={d_model} H={H} D={D} (H*D = {H * D})") # ── 1. 单个 N 下的逐项账本 ────────────────────────────── B, N = 1, 4096 items = ledger(B, N, d_model, H, "bf16") total = sum(items.values()) print(f"\n[1] 逐项账本 B={B} N={N} bf16(2 字节)") print(f" {'项目':<16}{'形状':<22}{'字节数':>14} 占比") for name, nb in items.items(): shape = "[B, N, d_model]" if nb == B * N * d_model * 2 else "[B, H, N, N]" print(f" {name:<16}{shape:<22}{nb:>14,} {nb / total:6.1%}") print(f" {'合计':<16}{'':<22}{total:>14,} {fmt(total)}") quad = sum(v for k, v in items.items() if "S" in k or "P" in k or "乘子" in k) lin = total - quad print(f" → N² 项共 {fmt(quad)},线性项共 {fmt(lin)},N² 项占 {quad / total:.1%}") # ── 2. N 增长时账本怎么变 ─────────────────────────────── print("\n[2] N 增长时一层 MHA 的激活显存(B=1, bf16, 含 浮点 dropout 乘子)") print(f" {'N':>7}{'线性项':>12}{'N² 项':>12}{'合计':>12} " f"{'融合侧教学值':>15} 倍数") for N in (1024, 2048, 4096, 8192, 16384, 32768): it = ledger(1, N, d_model, H, "bf16") q = it["分数矩阵 S"] + it["softmax 权重 P"] + it["浮点 dropout 乘子"] l = sum(it.values()) - q # 与 FlashAttention 文统一:六份线性量的教学预算,忽略小的 LSE 与工作区 flash = 6 * 1 * N * d_model * 2 print(f" {N:>7}{fmt(l):>12}{fmt(q):>12}{fmt(l + q):>12} " f"{fmt(flash):>15} {(l + q) / flash:5.1f}x") # ── 3. 交叉点:从哪个 N 开始 N² 项压过线性项 ──────────── # 3·B·H·N² = 6·B·N·d → N* = 2d / H N_star = 2 * d_model / H print("\n[3] 交叉点") print(f" 3·B·H·N²·bytes = 6·B·N·d·bytes 解得 N* = 2d/H = {N_star:.0f}") print(f" 也就是说 N 超过 {N_star:.0f} 之后,一层 attention 的激活显存") print(f" 就由 N² 项主导;N=4096 时已经超出 {(4096 / N_star):.0f} 倍。") # ── 4. 乘上层数:为什么 L 层比 N 更狠 ──────────────────── print("\n[4] 乘上层数 L(无重算训练,按上述教学模型保留各层中间量)") print(f" {'L':>4}{'N=2048':>12}{'N=4096':>12}{'N=8192':>12} 说明") for L in (12, 24, 32, 48): row = [] for N in (2048, 4096, 8192): tot = sum(ledger(1, N, d_model, H, "bf16").values()) row.append(fmt(tot * L)) print(f" {L:>4}{row[0]:>12}{row[1]:>12}{row[2]:>12}") print(" 以上均未计入模型权重与优化器状态;梯度检查点减少内部保存量,但仍需保留边界激活,并非通用的 1/L。") # ── 5. 推理侧的另一种账本:KV cache ───────────────────── print("\n[5] 推理侧:KV cache 的账本(随 N 线性增长,但随并发数 B 线性增长)") print(" 公式: 2 · B · N · L · H · D · bytes") print(f" {'B':>4}{'N':>7}{'L=32':>12}{'L=80':>12}") for B in (1, 8, 32): for N in (4096, 32768): r = [] for L in (32, 80): nb = 2 * B * N * L * H * D * 2 r.append(fmt(nb)) print(f" {B:>4}{N:>7}{r[0]:>12}{r[1]:>12}") # ── 6. 权重和激活谁更大 ──────────────────────────────── print("\n[6] 权重 vs 激活:别只盯着模型大小") L = 32 # 一层 transformer 的权重:4 个 d×d 投影 + 2 个 FFN 的 4d×d w_per_layer = (4 * d_model * d_model + 2 * d_model * 4 * d_model) * 2 print(f" 一层权重(4 个 d×d + FFN 的 2×4d×d,bf16): {fmt(w_per_layer)}") print(f" L={L} 层总权重 : {fmt(w_per_layer * L)}") for N in (2048, 8192): act = sum(ledger(1, N, d_model, H, "bf16").values()) * L print(f" N={N} 时 L={L} 层激活 : {fmt(act)} " f"(是权重的 {act / (w_per_layer * L):.1f} 倍)") if __name__ == "__main__": main() flops_ledger.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """自注意力的算力账本:把一层 attention 拆成「线性项」和「二次项」两笔。 只依赖 numpy。直接 `python flops_ledger.py` 即可运行。 最后有一段 CPU 实测,耗时约 10 秒;不同机器数值会不同,看趋势即可。 FLOPs 口径:一次矩阵乘 [M,K] @ [K,N] 需要 M·K·N 次乘加(MAC), 一次乘加算 2 个浮点运算,所以 FLOPs = 2·M·K·N。全文统一用这个口径。 """ import time import numpy as np def layer_flops(N: int, d: int, with_ffn: bool = True): """一层 transformer 的前向 FLOPs,拆成线性项和二次项。""" # 4 个 d×d 投影:Q、K、V 三个输入投影 + 1 个输出投影 proj = 4 * 2 * N * d * d # FFN:升到 4d 再降回来,两个矩阵乘 ffn = 2 * 2 * N * d * 4 * d if with_ffn else 0.0 # 注意力内部两个 N×N 的矩阵乘:QK^T 与 P·V attn = 2 * 2 * N * N * d return proj + ffn, attn def fmt_flops(f: float) -> str: for unit, scale in (("P", 1e15), ("T", 1e12), ("G", 1e9), ("M", 1e6)): if f >= scale: return f"{f / scale:7.2f} {unit}FLOPs" return f"{f:7.2f} FLOPs" def main(): d, H = 3072, 24 print(f"模型尺寸: d_model={d} H={H}") # ── 1. 一层里两笔账各是多少 ───────────────────────────── print("\n[1] 一层 transformer 的前向 FLOPs(N=4096, 含 FFN)") linear, quad = layer_flops(4096, d) print(f" 线性项(QKV 投影 + O 投影 + FFN): {fmt_flops(linear)}") print(f" 二次项(QK^T 与 P·V) : {fmt_flops(quad)}") print(f" 二次项占比 : {quad / (linear + quad):.1%}") # 只看 attention 内部:4 个投影 vs 2 个 N×N 矩阵乘 l_attn_only, q_attn_only = layer_flops(4096, d, with_ffn=False) print(f" 只算 attention 模块本身(去掉 FFN): 二次项占 " f"{q_attn_only / (l_attn_only + q_attn_only):.1%}") # ── 2. 交叉点:二次项什么时候超过线性项 ────────────────── # 4N²d = 24Nd² → N* = 6d print("\n[2] 交叉点") print(f" 含 FFN : 4N²d = 24Nd² → N* = 6d = {6 * d}") print(f" 只含投影: 4N²d = 8Nd² → N* = 2d = {2 * d}") print(f" 结论:N 小于 {2 * d} 时,二次项的 FLOPs 少于 attention 投影,") print(f" 这只比较运算量;是否为耗时瓶颈仍取决于访存、内核与硬件。") # ── 3. 占比怎么随 N 变化 ─────────────────────────────── print("\n[3] 二次项占比随 N 变化(d=3072,含 FFN)") print(f" {'N':>7}{'线性项':>16}{'二次项':>16}{'二次项占比':>10}") for N in (512, 1024, 2048, 4096, 8192, 16384, 32768, 65536): l, q = layer_flops(N, d) print(f" {N:>7}{fmt_flops(l):>16}{fmt_flops(q):>16}{q / (l + q):>10.1%}") # ── 4. 训练总算力:前向的三倍 ─────────────────────────── print("\n[4] 训练一个 token 的 FLOPs(前向 + 反向 ≈ 3 倍前向)") for N in (2048, 4096, 8192): l, q = layer_flops(N, d) print(f" N={N:<6} 每 token 前向 {fmt_flops((l + q) / N)}" f" 训练 {fmt_flops(3 * (l + q) / N)}") # ── 5. CPU 实测:二次项的增长是不是真的更快 ────────────── print("\n[5] CPU 实测(本机一次运行,看趋势不看绝对值;numpy float32)") d_t, H_t = 1024, 16 rng = np.random.default_rng(0) X = rng.normal(size=(4096, d_t)).astype(np.float32) W = rng.normal(size=(d_t, d_t)).astype(np.float32) / np.sqrt(d_t) def bench(fn, repeat=3): best = float("inf") for _ in range(repeat): t0 = time.perf_counter() fn() best = min(best, time.perf_counter() - t0) return best * 1000.0 print(f" {'N':>6}{'投影 (ms)':>12}{'QK^T (ms)':>12}{'实测比':>9}" f"{'FLOPs 比':>10}") measured = {} for N in (256, 512, 1024, 2048, 4096): x = X[:N] t_proj = bench(lambda: x @ W) D_t = d_t // H_t Q = (x @ W)[:, :].reshape(N, H_t, D_t).transpose(1, 0, 2) t_qk = bench(lambda: Q @ Q.transpose(0, 2, 1)) # 理论 FLOPs 比:二次项 / 一个投影 ratio_flops = (2 * N * N * d_t) / (2 * N * d_t * d_t) measured[N] = t_qk / t_proj print(f" {N:>6}{t_proj:>12.2f}{t_qk:>12.2f}" f"{t_qk / t_proj:>9.2f}{ratio_flops:>10.2f}") print(f" 注意看 N=256 那一行:QK^T 的算术量只有投影的 0.25 倍,") print(f" 本次 QK^T / 投影耗时比为 {measured[256]:.2f}。耗时比不必等于 FLOPs 比。") # ── 6. 算术强度:FLOPs 回答不了「为什么慢」 ─────────────── print("\n[6] 算术强度 AI = FLOPs / 访存字节(float32,4 字节/元素)") print(" 估算 AI 较低提示应检查访存;是否实际受限需测硬件计数器。") print(f" {'算子':<22}{'N=256':>12}{'N=1024':>12}{'N=4096':>12}") rows = [] for N in (256, 1024, 4096): # 投影 X[N,d] @ W[d,d]:读 X 与 W,写输出 f_proj = 2 * N * d_t * d_t m_proj = (N * d_t + d_t * d_t + N * d_t) * 4 # QK^T:读 Q 与 K,写 N×N 的分数矩阵(H 份) f_qk = 2 * N * N * d_t m_qk = (2 * N * d_t + H_t * N * N) * 4 rows.append((N, f_proj / m_proj, f_qk / m_qk)) print(f" {'投影 X@W_q':<22}" + "".join(f"{r[1]:>12.1f}" for r in rows)) print(f" {'注意力 QK^T':<22}" + "".join(f"{r[2]:>12.1f}" for r in rows)) print(f" {'(投影 / QK^T)':<22}" + "".join(f"{r[1] / r[2]:>12.1f}" for r in rows)) print(" 投影的 AI 高出 4~7 倍:权重矩阵被整批 token 复用,读一次能做 N 次乘加。") print(" QK^T 的 AI 低:输出是 H·N²,写完就走,数据复用少。") print(" 该估算不能解释本机时延;FlashAttention 原论文通过分块减少 HBM IO,") print(" 避免显式落地完整 N×N 权重矩阵,但算术量仍是二次。") if __name__ == "__main__": main() make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「自注意力机制的计算与显存账本」的三张解释图。 数值全部来自同目录下的 attention_memory.py 与 flops_ledger.py, 改了那两个脚本的话这里要跟着重跑,避免图与正文数字不一致。 只依赖 numpy + matplotlib,直接 `python make_figures.py` 即可运行。 """ from pathlib import Path import matplotlib.pyplot as plt import numpy as np from matplotlib.patches import FancyBboxPatch ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, }) C_QUAD = "#e05263" # N² 项:红 C_LIN = "#3f7fbf" # 线性项:蓝 C_FLASH = "#2f9e6f" # FlashAttention:绿 C_GREY = "#94a3b8" INK = "#182238" MIB = 1024 ** 2 GIB = 1024 ** 3 def _style(ax, title, xlabel, ylabel): ax.set_title(title, fontsize=14, weight="bold", color=INK, pad=10) ax.set_xlabel(xlabel, fontsize=11, color="#475569") ax.set_ylabel(ylabel, fontsize=11, color="#475569") ax.tick_params(colors="#475569", labelsize=10) for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("left", "bottom"): ax.spines[s].set_color("#cbd5e1") ax.grid(axis="y", color="#eef2f7", linewidth=1) ax.set_axisbelow(True) # ──────────────────────────────────────────────────────────── # 图 1:显存账本的构成 # ──────────────────────────────────────────────────────────── def fig_memory_ledger(): d_model, H = 3072, 24 b = 2 # bf16 fig, axes = plt.subplots(1, 2, figsize=(13.2, 4.9), facecolor="#fbfcfe") # 左:N=4096 时的逐项占比(一根堆叠条) ax = axes[0] N = 4096 lin_item = 1 * N * d_model * b quad_item = 1 * H * N * N * b labels = ["输入 X", "Q", "K", "V", "分数矩阵 S", "softmax 权重 P", "浮点 dropout 乘子", "加权和 ctx", "输出 O"] sizes = [lin_item, lin_item, lin_item, lin_item, quad_item, quad_item, quad_item, lin_item, lin_item] colors = [C_LIN] * 4 + [C_QUAD] * 3 + [C_LIN] * 2 total = sum(sizes) left = 0.0 for lab, s, c in zip(labels, sizes, colors): ax.barh([0], [s / GIB], left=left / GIB, color=c, edgecolor="white", linewidth=1.2, height=0.55) if s / total > 0.05: ax.text(left / GIB + s / GIB / 2, 0, f"{lab}\n{s / total:.1%}", ha="center", va="center", fontsize=9.5, color="white", weight="bold") left += s ax.set_xlim(0, total / GIB) ax.set_yticks([]) for s in ("left", "top", "right"): ax.spines[s].set_visible(False) ax.set_xlabel("单层 MHA 的教学激活存储(GiB)", fontsize=11, color="#475569") ax.set_title(f"N={N} 时,三个 N² 项吃掉 94%", fontsize=14, weight="bold", color=INK, pad=10) ax.text(0, -0.42, f"合计 {total / GIB:.2f} GiB | 模型 d_model={d_model}, " f"H={H}, bf16", ha="left", fontsize=10, color="#64748b") # 右:随 N 增长,朴素实现 vs FlashAttention ax = axes[1] Ns = np.array([1024, 2048, 4096, 8192, 16384, 32768]) lin = 6 * Ns * d_model * b # 与逐项账本相同:6 个 [B,N,d] 张量 quad = 3 * H * Ns ** 2 * b # 3 个 [B,H,N,N] 张量 naive = (lin + quad) / GIB flash = (6 * Ns * d_model * b) / GIB # 统一教学线性预算,忽略小的 LSE ax.plot(Ns, naive, "o-", color=C_QUAD, linewidth=2.4, markersize=6, label="朴素实现(落 N² 到显存)") ax.plot(Ns, flash, "s-", color=C_FLASH, linewidth=2.4, markersize=6, label="融合侧教学值(6Nd)") ax.fill_between(Ns, flash, naive, color=C_QUAD, alpha=0.10) ax.set_yscale("log") ax.set_xscale("log") ax.set_xticks(Ns) ax.set_xticklabels([f"{n:,}" for n in Ns], fontsize=9) ax.legend(fontsize=10, frameon=False, loc="upper left") _style(ax, "两者比值随 N 近似线性增长", "序列长度 N(token 数)", "单层教学激活存储(GiB,对数轴)") ax.text(0.98, 0.06, f"N=32768 时相差 {naive[-1] / flash[-1]:.0f} 倍", transform=ax.transAxes, ha="right", fontsize=10, color=C_QUAD, weight="bold") fig.suptitle("自注意力的显存账本:钱花在哪一笔", fontsize=17, weight="bold", color="#0f172a", y=1.03) fig.tight_layout() fig.savefig(OUT / "memory_ledger.png", bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close(fig) # ──────────────────────────────────────────────────────────── # 图 2:算力账本里二次项的占比 # ──────────────────────────────────────────────────────────── def fig_flops_ratio(): d = 3072 Ns = np.logspace(9, 16.5, 300, base=2) # 512 ~ 92682 quad = 4 * Ns ** 2 * d lin_proj = 8 * Ns * d * d # 只含 4 个投影 lin_all = lin_proj + 16 * Ns * d * d # 再算上 FFN fig, ax = plt.subplots(figsize=(8.6, 5.0), facecolor="#fbfcfe") r_only = quad / (quad + lin_proj) r_all = quad / (quad + lin_all) ax.plot(Ns, r_only, "-", color=C_QUAD, linewidth=2.6, label="只算 attention 模块(4 个投影 vs 2 个 N×N 乘)") ax.plot(Ns, r_all, "-", color=C_LIN, linewidth=2.6, label="算上 FFN(整层的线性项)") ax.axhline(0.5, color=C_GREY, linestyle="--", linewidth=1.2) ax.text(Ns[0], 0.52, "50% 线", fontsize=10, color=C_GREY) for N_star, color, tag in ((2 * d, C_QUAD, "N* = 2d = 6,144"), (6 * d, C_LIN, "N* = 6d = 18,432")): ax.axvline(N_star, color=color, linestyle=":", linewidth=1.6) ax.text(N_star * 1.05, 0.06, tag, fontsize=10.5, color=color, weight="bold", rotation=90, va="bottom") ax.set_xscale("log", base=2) ax.set_xticks([512, 1024, 2048, 4096, 8192, 16384, 32768, 65536]) ax.set_xticklabels(["512", "1K", "2K", "4K", "8K", "16K", "32K", "64K"]) ax.set_ylim(0, 1) ax.legend(fontsize=10.5, frameon=False, loc="upper left") _style(ax, "二次项的 FLOPs 占比随 N 增长(不等于耗时占比)", "序列长度 N(token 数)", "二次项占该层前向 FLOPs 的比例") ax.text(0.98, 0.30, "常用的 N=4096:\n只算 attention 时二次项占 40%,\n算上 FFN 后只占 18%", transform=ax.transAxes, ha="right", fontsize=10, color="#475569", bbox=dict(boxstyle="round,pad=0.4", facecolor="#f1f5f9", edgecolor="#e2e8f0")) fig.suptitle("算力账本:二次项什么时候才是主角", fontsize=16, weight="bold", color="#0f172a", y=1.00) fig.tight_layout() fig.savefig(OUT / "flops_ratio.png", bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close(fig) # ──────────────────────────────────────────────────────────── # 图 3:多头到底切了什么 # ──────────────────────────────────────────────────────────── def fig_head_split(): fig, ax = plt.subplots(figsize=(11.4, 5.0), facecolor="#fbfcfe") ax.set_xlim(0, 1) ax.set_ylim(0, 1) ax.axis("off") def block(x, y, w, h, title, lines, color, fs=11): ax.add_patch(FancyBboxPatch( (x, y), w, h, boxstyle="round,pad=0.01,rounding_size=0.02", facecolor=color, edgecolor="white", linewidth=1.6)) ax.text(x + w / 2, y + h * 0.74, title, ha="center", va="center", fontsize=fs + 2, weight="bold", color=INK) ax.text(x + w / 2, y + h * 0.33, "\n".join(lines), ha="center", va="center", fontsize=fs, color="#334155", linespacing=1.5) def arrow(x1, y, x2, label): ax.annotate("", xy=(x2, y), xytext=(x1, y), arrowprops=dict(arrowstyle="-|>", color=C_GREY, linewidth=2.0, mutation_scale=16)) ax.text((x1 + x2) / 2, y + 0.045, label, ha="center", fontsize=11, color="#475569", weight="bold") # X block(0.02, 0.30, 0.16, 0.42, "X", ["[B, N, d_model]", "d_model = H · D"], "#dbeafe") arrow(0.185, 0.51, 0.245, "三个投影") # Q/K/V 未切头 block(0.25, 0.30, 0.17, 0.42, "Q, K, V", ["各 [B, N, H·D]", "仍是完整宽度"], "#e0e7ff") arrow(0.425, 0.51, 0.485, "切头") # 切头后 block(0.49, 0.30, 0.17, 0.42, "切头之后", ["[B, H, N, D]", "reshape + transpose"], "#ede9fe") arrow(0.665, 0.51, 0.725, r"$QK^\top/\sqrt{D}$") # 分数矩阵 block(0.73, 0.22, 0.25, 0.58, "分数矩阵 S", ["[B, H, N, N]", "H 个 N×N,不是 1 个", "← 显存就花在这里"], "#fecaca", fs=11) # 底部说明 ax.text(0.5, 0.10, "每个头只用自己的 D = d_model / H 维去做内积,H 个头各算各的 N×N;\n" "最后把 H 份 D 维结果拼回 d_model,再过一次输出投影 W_o。", ha="center", va="center", fontsize=11.5, color="#475569", linespacing=1.6) ax.text(0.5, 0.02, "关键:分头不改变总算术量(H · N² · D = N² · d_model),它改变的是每个子空间的表达能力。", ha="center", va="center", fontsize=11, color=C_QUAD, weight="bold") fig.suptitle("多头注意力切的是 d_model 这一维,不是复制 H 份", fontsize=16, weight="bold", color="#0f172a", y=0.99) fig.tight_layout() fig.savefig(OUT / "head_split.png", bbox_inches="tight", facecolor=fig.get_facecolor()) plt.close(fig) if __name__ == "__main__": fig_memory_ledger() fig_flops_ratio() fig_head_split() print(f"已生成 3 张图到 {OUT}") for p in sorted(OUT.glob("*.png")): print(f" {p.name} {p.stat().st_size / 1024:.0f} KiB")
2026年09月24日
6 阅读
0 评论
0 点赞
2026-09-24
AIGC 每日速读|2026-09-24|阿里Qwen3.8-Omni发布,视频推理输Gemini
今日 AIGC 论文速览 今日共 10 篇 · 全模态基座与语音智能体 2 篇 · 视频编辑与生成智能体 2 篇 · 生成推理加速与量化 2 篇 · 视觉Tokenizer与音视频生成 2 篇 · 奖励模型与生成安全评测 2 篇 重点论文标题列表 Qwen3.8-Omni(Qwen Team, Alibaba Group):原生全模态智能体,百万上下文 VideoX-Qwen(Nanjing University):120万对指令视频编辑数据 QuantWM(Harbin Institute of Technology (Shenzhen)):2比特KV量化压住时序闪烁 Flash-dLLM(VILA Lab, MBZUAI):扩散LLM提速五倍且不伤精度 VideoGen-Agent(Princeton University):RL教会智能体用工具做视频 今日论文速览 1. Qwen3.8-Omni:原生全模态智能体,百万上下文 Qwen3.8-Omni: Towards Native Omni-Modal Agents | Qwen Team, Alibaba Group | arXiv:2609.25611 关键词:全模态智能体,稀疏MoE,百万上下文,视频剪辑,插件框架 前序问题:全模态模型此前主要强调感知和交互,真正的长程智能体任务——调工具、跑多步、把一段素材剪辑成片——要靠外挂 harness 拼,而现有 agent harness 对音频和视频没有原生支持。缺一个在多模态上原生 co-training、又不丢文本能力的底座。 本文贡献:发布 Qwen3.8-Omni-Flash:Thinker 继承 Qwen3.8-Next 的稀疏 MoE 骨干(GDN 与交错注意力混合),上下文扩到 100 万 token;原生多模态 co-training 把文本侧的智能体能力迁移到音视频任务,再经多教师蒸馏合并领域专精策略;配套开源 Qwen-MM-Plugins 插件框架与 Qwen-Live-Harness 实时交互框架。 Qwen3.8-Omni-Flash is a unified end-to-end model capable of processing multiple modalities 实验效果:WildClawBench-MM 从上代 34.5 涨到 71.0(+36.5)、AgenticVBench 14.5→36.8、OmniGAIA 57.2→74.0;Video-MME-v2 47.9→65.0。开智能体模式后 LVOmniBench 63.3→73.6,OmniVideoBench 单次 query 消耗 token 从 14.6 万降到 7.9 万(约 -45.7%)。 The overview of Qwen3.8-Omni-Flash-Realtime and Qwen-Live-Harness 批判点评:摘要口气很大,但表格里对手没输干净:Gemini 3.8 Flash 在三个视频推理基准的静态设置上全面更高,开智能体后 OmniVideoBench(70.1 vs 67.8)和 Video-MME-v2 仍被压着;teaser 图的 OmniCap-IF 一栏 Flash 拿 20.2,落在上代 Plus 的 28.3 和 Muse Spark 1.2 的 26.8 之后。所谓「领先」主要立在自己参与的新评测上。 2. VideoX-Qwen:120万对指令视频编辑数据 VideoX-Qwen: Data-Centric Instruction-Based Video Editing | Nanjing University;Jiutian Research | arXiv:2609.26015 关键词:指令视频编辑,数据合成,Qwen3-VL,Wan,成对监督 前序问题:通用视频编辑缺的不是模型而是监督:编辑要在执行指令的同时保住无关主体、场景结构和时序连续性,而成对(源视频、指令、目标视频)数据的规模和任务覆盖都远远不够,指令视频编辑长期落后于图像编辑。 本文贡献:一条可扩产的数据流水线:Qwen3-VL-235B 做源视频筛选与目标解析,加/删走 SAM3 加 Minimax-Remover,替换/属性走 Qwen-Image-Edit 首帧编辑加 WanAnimate 传播,自动通过率 89%,产出 120 万条方向性编辑对(每组任务超 40 万条)。模型侧统一 Qwen-Wan 编辑器:多模态语义条件加稠密源视频 latent 引导,图像-视频三阶段渐进训练。 Paired-video construction pipeline 实验效果:100 例对比 UniVideo 和 Kling O1:11 项指标里 9 项均值最优,覆盖指令遵循、编辑质量、内容保持、结构与感知相似度和视频分布质量。 Qualitative comparisons on compound and attribute editing 批判点评:「9/11 最优」的另一面是 11 项里输了 2 项,且对比规模只有 100 例;加/删两条数据路线直接复用 Minimax-Remover 和 WanAnimate,编辑对质量上限被这些现成生成器封顶。论文把「逐条校验富化后的指令」画成 proposed extension(虚线框),等于承认 120 万对里指令与实际编辑的一致性还没有被逐一验证。 3. QuantWM:2比特KV量化压住时序闪烁 QuantWM: Temporally Consistent 2-Bit KV Cache Quantization for World Models and Video Generation | Harbin Institute of Technology (Shenzhen);National University of Singapore | arXiv:2609.26425 关键词:KV cache量化,视频生成加速,世界模型,注意力补偿,训练无关 前序问题:视频生成和世界模型的 KV cache 是部署瓶颈,已有 2-bit 量化方法在 VBench 这类指标上「几乎无损」,但作者实测发现:指标没掉,画面却在闪——逐帧指标天生看不见时序问题。 本文贡献:训练无关、严格因果的 2-bit KV 量化框架。诊断先行:Key 的量化重建误差比 Value 小,输出退化却更大,因为小扰动就能改变 QK^T 注意力 logits、挪走 Query 选中的时序-空间 token。对症两招:QSAC 聚类联合历史 Query 敏感度与残差范围选 INT2 友好的 Key 质心;PSAC 用低秩投影沿 Query 主子空间补偿残余 Key 误差。 Overview of QuantWM's inference system 实验效果:Causal-Forcing、LingBot-World-v2、HY-World 1.5、Matrix-Game-2、Longcat-Video 五个底模上,KV cache 最高压缩 6.20×,图像与视频质量指标全面超过现有 2-bit 方法,额外开销有限,且保持严格因果可流式推理。 Visual comparison of BF16, K16V2, K2V16 and K2V2 on world models 批判点评:这篇最值钱的是诊断而不是方法:它解释了为什么「VBench 几乎无损」的结论靠不住——逐帧指标对时序闪烁盲。但时序改善的验收仍以看图为主,论文没给出时序稳定性的量化指标;6.20× 的压缩上限是在 93 帧生成这一档测的,更长序列的收益未报告。 4. Flash-dLLM:扩散LLM提速五倍且不伤精度 Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs | VILA Lab, MBZUAI | arXiv:2609.26796 关键词:扩散语言模型,KV cache,并行解码,Triton,推理加速 前序问题:扩散 LLM 靠并行去噪生成文本,但双向注意力让自回归那套 KV cache 直接失效;已有加速工作把 KV cache 和并行解码分开研究,两者合用时 GPU 显存 I/O 反而成了主要瓶颈。 本文贡献:训练无关两件套。Flash-Cache:把 QKV 投影、RoPE、cache 写入、注意力计算融进一个 I/O 感知的融合 Triton 核,砍掉冗余访存;Flash-Verify:让 dLLM 自己既当 drafter 又当 verifier 做草稿-验证式并行解码,不挂辅助模型。 Overview of Flash-dLLM 实验效果:LLaDA-1.5 上对此前最强的 Elastic-Cache 在 GSM8K/HumanEval 分别提速 5.1 倍和 11.0 倍;Flash-Cache 加 Flash-Verify 达到精度 83.02%、吞吐 210.6 tokens/s(Elastic-Cache 为 82.79%、41.7);batch 16 显存约 26GB,对照 Fast-dLLM 的 50GB 降约 48%,后者在 batch 24 直接 OOM。 Throughput and peak memory versus batch size on GSM8K 批判点评:赢面都在 dLLM 阵营内部:可扩展性图里自回归的 Llama3 在 batch 32 的吞吐仍高于 Flash-dLLM,「扩散 LLM 加速后追平 AR」还没有发生。精度与吞吐是硬权衡(80.2% 到 83.2% 对应 278 降到 186 tokens/s),宣传里 81 倍是对无 cache 基线的倍数,直接引用容易被误读成对 SOTA 的提速。 5. VideoGen-Agent:RL教会智能体用工具做视频 VideoGen-Agent: Reinforcing Video Generation Agents | Princeton University;Stanford University;UC Davis;MMLab, CUHK;BenchFlow;GWU | arXiv:2609.24997 关键词:生成智能体,强化学习,工具调用,视频生成,VABench 前序问题:视频生成模型面对需要专业知识、特定身份、物理一致或时序事件的 prompt 经常直接翻车:太极招式做不对、指定角色画不像、物理参数乱来。单靠扩大底模收益有限,缺一个会检索、仿真、检测来补信息的生成智能体。 本文贡献:六类任务上训一个共享策略:先在教师轨迹上 SFT 建立工具使用习惯(检索文本/图像、仿真、目标检测、深度估计等),再用类别感知混合奖励(工具调用合法性、任务适配度、视频质量)做 RL。配套 VABench:600 条 held-out prompt,覆盖程序知识、单/多实体身份、物理一致、场景合成、多镜头时序。 Overview of VideoGen-Agent 实验效果:VABench 上从底模 56.5 提到 75.6(+19.1 分);把生成工具换成更强的 Seedance 2.0 后,不重训智能体就到 86.1;人评 84.3% 偏好升级配置而非最强独立基线。 Qualitative comparison on Procedural Knowledge generation 批判点评:86.1 这档数字藏着归因问题:从 75.6 到 86.1 的增量完全来自换生成工具,智能体一行代码没动——agent 框架的收益上限被底模和工具质量强绑定,这一段「白捡的分」其实是底模的分。混合奖励里「视频质量」一项依赖 VLM 裁判打分,裁判偏差会直接写进策略,论文未讨论防 reward hacking 的机制。 6. Vorch-Human:一个模型统一说话人与歌声视频 Vorch-Human: Unified Multi-Task Human-Centric Generation via Long-Horizon Continuation | Vorch Team;Harbin Institute of Technology, Shenzhen;Tongji University | arXiv:2609.26117 关键词:音视频统一生成,数字人,扩散transformer,长视频,身份保持 前序问题:以人为核心的音视频生成被拆成好几个模型:语音驱动数字人、参考音色合成说话视频、外观加声音参考生成场景,各自训练各自维护。它们本质上是同一组模态、只是条件不同的任务。 本文贡献:双流音视频扩散 transformer,在传统 noisy 音频/视频接口之外增加 clean 条件 token 组,用 per-token 任务嵌入、时间位置类型、条件掩码和共享多模态 prompt 编码器,把驱动语音、音色样本、首帧、主体图统一进一个模型;两级数据管线做说话人聚类与跨片段身份串联;长视频靠冻结 latent 前缀递归、每段只生成新后缀。 Role-aware unified human-centric architecture 实验效果:5 分钟长生成保持身份一致与音画同步;GSB 人评对 LongCat-Video-Avatar-1.5:动作/姿态/表情净偏好 +46.9%,口型准确度 +47.5%,单人音视频驱动整体净偏好 +27.1%。 Raw GSB breakdown for LongCat-Video-Avatar-1.5 批判点评:同一张 GSB 图的另一半:视觉质量净偏好 -13.1%,多人场景再输 -9.2%——赢的两项全在「跟音频对齐」上,纯画质反而输给 LongCat,摘要里 strong identity preservation 的措辞没覆盖这个短板。署名主体写作「Vorch Team」,机构披露不完整,技术报告性质明显。 7. StableVQ:三招治住VQ码本训练崩溃 StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training | Huazhong University of Science and Technology;KlingAI Research;South China Normal University | arXiv:2609.26774 关键词:VQ tokenizer,码本利用率,训练稳定,自回归生成,ImageNet 前序问题:共享投影 codebook 把 VQ tokenizer 的利用率推上去之后,训练稳定性成了新瓶颈:码本利用长期趴在低位、出现永久死码、甚至在 26 万步附近突然崩到 0。现有方法各打各的补丁,缺少对根因的解释。 本文贡献:把根因归结为 Encoder-Decoder 与 Codebook 的「纠缠训练」:两个子系统各自都无法独立完成任务,只能靠碰巧合作维持。三个零参数改动:Dynamic STE 按量化距离压低不可靠梯度;Region VQ Loss 让激活码向邻近未激活码按比例传播目标;Decoupled Schedule 给两个子系统各配独立学习率调度。 Three characteristic failure modes of VQ training 实验效果:ImageNet 256 重建:262k 码本 120 epoch rFID 0.92(对照 FVQ 1.29、SimVQ 3.16),16k 码本 1.13;鲁棒性指标 UR-AUC 60.59,比 FVQ 的 8.08、SimVQ 的 2.17 高一个量级;只用单个线性投影层就追平需要 ViTBlock 投影器的 FVQ。 Qualitative reconstruction comparison with a 16k codebook 批判点评:主张是「治崩溃」,但主表里 FVQ 和 SimVQ 在各自标准配置下利用率本来就是 100%,差距全在故意制造码本-token 分布错配的鲁棒性测试里拉开——这三招买的更像「最坏情况的保险」,常态训练收益主要是 rFID 那 0.2 到 0.4。下游生成只给「competitive」,FID 全表在附录。 8. Qwen-Audio 3.1:全双工误答率从73%压到13% Qwen-Audio-3.1-Realtime: Towards Reliable Agentic Voice Interaction | Alibaba Token Foundry, Alibaba Group | arXiv:2609.25176 关键词:全双工语音,语音智能体,GRPO,在线蒸馏,工具调用 前序问题:实时语音助手要同时干三件互相打架的事:理解还在演化的请求、执行动作、遵守对话规则(何时说、何时闭嘴、是否插话)。上一代在背景有人说话时 73% 的轮次会错误应答,全双工的「闭嘴」比「会说」更难。 本文贡献:Think/Act/Speak and Coordinate 三层:Core-Cocktail SFT 加 M2-OPD 多教师在线策略蒸馏迁移文本能力;自进化可执行环境加多粒度 rollout 的 GRPO 教工具调用与参数落地;全双工决策模型独立管「说不说」。另给 Voice Harness 前后台原型:对话在前台、任务在后台异步跑。 Qwen-Audio-3.1-Realtime model framework for continuous spoken interaction 实验效果:tau-Voice 半双工改编任务成功 78.4%→82.0%(Airline 64.0%→74.0%);Full-Duplex-Bench v1.5 背景语音应答率 73.0%→13.0%、恢复率 0.26→0.87;FLEURS 宏平均 WER 9.01→3.98;WebSearch1K 检索调用从 4.37 次降到 1.05 次(-76%)。 Qwen-Audio-3.1-Realtime post-training pipeline 批判点评:两个回撤写在正文里:EVA-A 的 Mean 分从 70.50 掉到 66.26,WebSearch1K 的 F1 从 60.87% 降到 58.61%——少说话、少调用是把双刃剑;Daily Chat 和 persona 交互的 Spoken 分也在微跌。摘要只报提升项,回撤要翻到结果表才看得见,读技术报告别只读开头。 9. RULER:六项Rubric奖励治SVG生成 RULER: Instance-aware Rubric Rewards for SVG Generation | Ant Group;The Hong Kong University of Science and Technology (Guangzhou);Independent Researcher;University of Oxford | arXiv:2609.25270 关键词:SVG生成,Rubric奖励,GRPO,reward hacking,视觉裁判 前序问题:自然语言生成 SVG 是没有绝对真值的开放任务:CLIP、Aesthetic 这些在自然图像上标定的标量指标迁到风格化矢量图上会误判,直接拿来当 RL 奖励还会触发 reward hacking——评测和优化同时失去可靠信号。 本文贡献:实例化 rubric 奖励:每条指令由前沿模型生成六项 rubric(语义、视觉、风格三轴),裁判 VLM 逐项给渲染 rollout 打分,加权满意度经 GRPO 优化。rubric 只从文本推导,不需要配对 SVG 真值、也不需要人类偏好标注。训练用 Qwen3-VL-8B 做裁判,与评测用的 GPT-5-mini 分开。 Qualitative comparison of scalable vector graphics generation between RULER and four baselines 实验效果:MMSVG-Illustration/Icon 上 rubric 分从 0.432/0.395 提到 0.693/0.683,超过专用 SVG 模型、追平大得多的 DeepSeek-V3;人评盲测对自家 Qwen3-8B 骨干胜率 89.7%、对 Qwen3-32B 66.1%;消融显示去掉语义轴掉 10.1%、换静态 rubric 掉 22.7%。 Ablation of rubric design on MMSVG-Illustration and MMSVG-Icon 批判点评:「追平 DeepSeek-V3」在可视化图里并非全赢:五个示例里同心圆靶一栏 DeepSeek-V3 的 0.94 高于 Ours 的 0.92;Ours 的 CLIP 分(0.28-0.35)与基线几乎无差,说明传统指标确实分不出好坏,但反过来也意味着 rubric 分数与人类直觉的对应完全押在裁判 VLM 上。人评规模只有 150 条 prompt。 10. Moderation Gap:三成有害视频骗过逐帧审核 The Temporal Moderation Gap: Text-to-Video Safety Filters Are Blind to Harm in Motion | National University of Singapore;University of New South Wales;Fuzhou University | arXiv:2609.26233 关键词:文生视频,安全审核,帧序,红队测试,评测方法 前序问题:T2V 服务的安全栈是从图像生成继承的:关键词过滤加随机抽 8 帧逐帧检查。这套栈对「伤害只存在于帧序」的视频天然失明——闯红灯、往转动的车床里伸手、误服药片,每一帧单看都无害。 本文贡献:理论上证明:任何忽略帧序的审核器只要放行某片段,就必然放行其良性重排。实测未修改的基准 prompt 就能让 32.7% 的 Sequential-Action 目标落入审核盲区,改写、拆场景、反馈式搜索都无显著改善(McNemar p≥0.12)——不需要任何 prompt 技巧就能触发。修复方向是读帧序:时序感知检测器 AUC 0.74,逐帧检查等于抛硬币。 The temporal moderation gap 实验效果:对 97 帧全部打分发现:约三分之一的交付片段只是把不安全帧藏起来,其余按帧序整体看仍有害,尽管每一帧都通过了审核;用户实验确认人能区分片段与其重排的安全性。 Three Sequential-Action gap clips 批判点评:论文顺带揭了一个测量陷阱:在搜索出的 prompt 上用它自己的渲染种子打分,会把 7.5% 的单次生成率虚标成 46.7%——这个坑对所有做安全评测的人都成立。但修复方案的 AUC 0.74 离可用尚远,且全部证据来自 Sequential-Action 一类动作伤害,跨帧累积的暗示性内容等其他类型未覆盖。 趋势观察 全模态模型开始从「能听会说」转向「能干活」 Qwen3.8-Omni 把上下文推到 100 万 token 并配好插件与实时 harness,Qwen-Audio-3.1-Realtime 则把「何时闭嘴」当成一等公民单独训练。两条线指向同一件事:多模态模型的竞争正在从感知指标挪向任务成功率,模型外面的 harness、插件和记忆管理层开始决定产品形态。 视频编辑与视频生成的瓶颈,先卡在数据而不是模型 VideoX-Qwen 用 120 万对合成编辑数据喂出一个统一编辑器,VideoGen-Agent 则用六类任务的教师轨迹加 RL 教模型调用工具。两者都在补同一个洞:指令视频任务缺少可靠的监督与反馈信号,谁的流水线能稳定产出高质量配对数据,谁就拿到上限。 加速这条线在「省显存」和「省访存」上分头推进,且都是训练无关 QuantWM 发现 2-bit KV 量化的真正受害者是时序稳定性而非单帧指标,用注意力补偿把闪烁压回去;Flash-dLLM 则把瓶颈定位到 GPU 显存 I/O,用融合核加自验证并行解码把吞吐推到 210.6 tokens/s。共同点是都不动训练,靠推理系统层的重新安排。 评测开始把「裁判本身」当成被审对象 RULER 证明静态标量指标在 SVG 上失灵,改用逐指令生成的 rubric 当奖励;Moderation Gap 则证明逐帧安全审核对帧序盲。两者都在说同一件事:在开放生成任务里,裁判的构造方式本身就是被评测对象,默认信任任何单一打分器都危险。 人工智能炼丹君 整理 | 2026-09-24 更多 AIGC 论文解读,关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注
2026年09月24日
16 阅读
0 评论
0 点赞
2026-09-23
AIGC 每日速读|2026-09-23|NVIDIA像素扩散FID 1.46,仍输latent
今日 AIGC 论文速览 今日共 10 篇 · 图像生成与像素扩散 2 篇 · 视频世界模型与 3D 一致 2 篇 · 视频生成加速与流式编辑 2 篇 · 动作与音频生成 2 篇 · 音视频因果与奖励评测 2 篇 重点论文标题列表 PixelDiT2(NVIDIA):像素扩散 FID 压到 1.46 WorldCrafter(ARC Lab, Tencent IEG):分钟级探索重访 PSNR 18.0 ⚡ GAE(HKUST):相机轨迹误差砍半 FVD 降 23% SparkDiffusion(Peking University, Melon Group):97% 稀疏下 265 倍单卡加速 MixiMotion(PTIT, Hanoi):一步动作对齐 0.835 逼近教师 今日论文速览 1. PixelDiT2:像素扩散 FID 压到 1.46 PixelDiT2: Representation-Grounded Pixel Diffusion Transformers | NVIDIA;University of Rochester | arXiv:2609.24919 关键词:像素空间扩散,表征先验,表示对齐,DINOv3,ImageNet ⚠️ 前序问题:像素空间扩散不用自编码器、直接在 RGB 上去噪,省掉了 latent 重建瓶颈,代价是收敛慢、最终画质长期落后 latent 扩散。作者的判断是缺一个显式表征先验:latent 扩散天然在紧凑结构化的隐空间里去噪,像素扩散却要一边学「好去噪的表征」一边学「怎么生成像素」,两件事挤在同一次训练里互相拖累。 本文贡献:提出 representation grounding:用一个冻结的预训练视觉基座模型在整个去噪过程里持续提供 per-patch 表征监督。带噪图像走像素去噪通路,同时冻结编码器产出 per-patch 特征,一个时间步条件化的 DiT 块 P_g 把它映射成 scaffold,再通过空间 AdaLN 调制扩散 transformer;REPA 只作为训练期辅助目标。另提出延迟 grounding dropout:训练初期保留 grounding,到 epoch 160 再开启 dropout。 PixelDiT2 at ImageNet-512x512: samples, architecture, and convergence. 实验效果:ImageNet-256×256 上 H/16 模型 600 epoch 达 FID 1.46、IS 301.6;512×512 上 680 epoch 达 FID 1.48、IS 295.7,两档都是像素侧最好(同表其余像素方法 1.61–3.94 / 1.78–3.95)。论文强调 epoch 预算:512 档 epoch 200 就超过 PixelDiT 850 epoch 的 FID,预算低 4.25 倍。 PixelDiT2 reaches FID 1.78 at epoch 200 and 1.48 at epoch 680. 批判点评:「补上表征先验」在像素侧成立,但把两张表横着读是另一面:256 上 latent 侧 RAE-XL FID 1.13、REPA-XL 1.29 都优于 1.46,512 上 RAE-XL 1.13 同样压过 1.48,IS 也不及 REPA-XL 306.3 与 JiT-G 306.8。也就是说「追平」只在像素扩散内部成立,跨阵营仍差约 0.3 FID,摘要里 narrowed the gap 的措辞容易被读成已经追平。另外延迟 dropout 的开启时机是在 512 档上调出来的,256 档用的是从头独立 dropout,两档训练协议并不一致,跨分辨率可比性打了折扣。 2. WorldCrafter:分钟级探索重访 PSNR 18.0 WorldCrafter: Consistent Video World Model with Implicit 3D-aware Memory | ARC Lab, Tencent IEG;Peking University | arXiv:2609.24984 关键词:视频世界模型,隐式3D记忆,长程一致性,相机可控,流式生成 ⚠️ 前序问题:视频世界模型做交互式探索时,长程一致性一直是硬伤:走出一段距离再回头,场景常常「认不出」是刚才那个地方。难点在于历史观测越来越多,而视频生成器的 token 预算是固定的,怎么把历史压进有限 token、又不依赖显式深度对应,是个两难——显式 3D 重建重且脆,简单堆历史帧则 token 直接爆掉。 本文贡献:核心洞察是「让被请求的视角来决定历史怎么压缩」。一个与视频生成器联合训练的 memory encoder 把历史观测编成紧凑的 3D-aware 表示,再由 pose-conditioned readout 按当前目标相机位姿读出固定数量的、目标视角专属 token,在去噪前注入,全程不需要显式深度对应。配合 max-coverage 历史检索、最近时间上下文和 few-step 蒸馏,实现从单张图或文本出发的流式分钟级探索。 Overview of the WorldCrafter pipeline. 实验效果:VBench 自建输入模式 725 段视频上 Overall 81.910 为全表第一(Alaya-EVOKE 81.406、SANA-WM 80.841、Echo-WM 80.211),主体一致性 82.695、背景一致性 90.589、运动平滑 98.787、整体一致性 26.745 均为第一。长程重访一致性上 MEt3R 0.166、PSNR 18.016、SSIM 0.517,相对次优 Lyra 2.0(0.334 / 14.050 / 0.390)近乎翻倍。 Qualitative long-horizon revisit comparison in a static scene. 批判点评:重访那张表真正的第一不是 WorldCrafter 而是它的加速版 WorldCrafter-fast(MEt3R 0.129、PSNR 20.868、SSIM 0.616)——少步蒸馏的变体反而在一致性上更好,这暗示「更多步数」本身可能带来累积漂移,论文没有就这条给出解释。视觉质量上也不是全胜:美学质量 AQ 61.432 只排第四,不及 SANA-WM 63.467;动态程度 DD 96.893 是全表最低(其余 97.087–100),说明换来的一致性里有一部分是靠画面动得更少买到的。评测均在自建输入模式下完成,未与公开榜单协议对齐。 3. GAE:相机轨迹误差砍半 FVD 降 23% GAE: Learning a Geometry-Native Latent Space for 3D-Consistent World Generation | The Hong Kong University of Science and Technology;ARC Lab, Tencent IEG;The University of Hong Kong;The University of Texas at Austin | arXiv:2609.24981 关键词:几何原生隐空间,3D一致性,自编码器表征,相机轨迹,感知生成统一 ⚠️ 前序问题:视觉生成器能出逼真画面,却保不住一个前后一致的 3D 场景。作者认为这不只是建模问题,更是表征问题:生成器演化的是外观为中心的隐空间,感知模型恢复几何用的是语义丰富、编码跨视角结构的空间,两者不在同一坐标系上。常见补丁是「再输出一个几何头」,等于把几何当副产品,治标不治本。 本文贡献:不把几何当额外输出,而是直接把一个几何基座模型的特征重参数化成紧凑、可扩散的隐空间——几何原生自编码器 GAE,其 latent 同时可解码到外观、深度、相机和点图。训练分两阶段:先训 codec 把冻结的多层几何特征压进 latent,再用条件流只在「待生成视角」的 latent 上做流匹配,参考视角 latent、相机射线和文本作为控制,让感知与生成共用同一个隐空间接口。 Overview of Geometry-Native Autoencoder (GAE). 实验效果:在固定生成器与训练协议的控制对比下:重建侧 GAE-128 用比原始几何特征少 24 倍的通道,PSNR 28.76 / LPIPS 0.036 / rFID 5.4(RealEstate10K)反而略优于原始 L0 特征;生成侧 GAE-64 最优,RealEstate10K FVD 225.7、PSNR 20.02、SSIM 0.711,DL3DV FVD 287.0、PSNR 18.00,相对最好的非 GAE 受控 latent 分别降 FVD 12.7% 与 23.1%;ATE 在 RealEstate10K 降 52.8%、DL3DV 降 23.3%,MEt3R 取到受控组最佳 0.1208 / 0.1347。 Multi-frame RGB, geometry, and camera-pose comparison across RealEstate10K scenes. 批判点评:最有价值的是控制变量设计,但同一张表也留了口子:受控组之外还列了 Gen3R 与 GLD 两个参考方法,其中 Gen3R 在 RealEstate10K 的 MEt3R 是 0.1157,优于 GAE-64 的 0.1208——「最佳」限于受控 latent 这个口径。DL3DV 的重投影误差上 GAE 只有 0.0028–0.0029,反而不如 WAN2.1 VAE 的 0.0019,说明几何原生 latent 在像素级回投影精度上并不占优。另外通道数选择本身就不一致:128 在重建上最好、64 在生成上最好,论文没给选择依据,实际部署还得自己再权衡一次。 4. SparkDiffusion:97% 稀疏下 265 倍单卡加速 SparkDiffusion: Mitigating the High-Sparsity Trap --- A Unified Framework for up to $265\times$ Single-GPU Acceleration of Visual Generation | Peking University, Melon Group;Tsinghua University;Alibaba Group;University of Electronic Science and Technology of China;Harbin Institute of Technology | arXiv:2609.23153 关键词:稀疏注意力,推理加速,少步蒸馏,FP8量化,视频扩散 ⚠️ 前序问题:视频扩散 transformer 贵在注意力要处理超长时空 token 序列,做稀疏注意力是自然的加速思路。但作者发现一个「高稀疏陷阱」:稀疏度推到极致时,逐步训练损失还在稳步下降,最终生成质量却停滞甚至倒退。诊断指出这是监督问题——主要误差来自高噪声的结构生成阶段,逐步局部训练修不动它,只有对齐终态的训练才修得动。 本文贡献:由此给出一条分阶段原则:先把稀疏架构适配成粗粒度先验,再校正终态分布。SparkDiffusion 按此串起三件事——短程稀疏 warm-up、少步 trajectory-mixed 蒸馏、FP8 量化配融合 kernel,并把三者做成可乘的加速因子。覆盖 Wan2.1/Wan2.2 骨干、T2V/I2V 任务、480P/720P 分辨率,以及 H100 与 RTX 5090 两种 GPU。 Overview of the SparkDiffusion framework. 实验效果:3 步无 CFG 推理下,Wan2.1-T2V-14B-720P 在单张 RTX 5090 上端到端 265× 加速(H100 上 220×);Wan2.1-T2V-1.3B-480P 端到端 1.3 秒出片。长序列 720P 上维持 97% 注意力稀疏,1.3B-480P 上 90%。VBench 总分从 dense 的 83.69 到 83.15,掉 0.54 分;同表延迟与显存项均为最好(18.0 / 8.0)。 End-to-end latency and speedup of SparkDiffusion over Full Attention on NVIDIA H100 and RTX 5090 GPUs. 批判点评:265× 是三个因子相乘的结果,其中大部分来自 3 步蒸馏而非稀疏本身,读标题很容易把功劳全记在稀疏上——论文在图表里做了分解,但正文强调的仍是合成后的倍数。质量代价是实打实的:83.15 低于 dense 83.69,且稀疏度从 90% 提到 97% 时分数从 83.42 掉到 83.15,说明高稀疏陷阱只是被缓解、没被消除。另外源码里仍能看到若干未清理的 TODO 占位注释(作者元数据、BF16 两列标注为占位),作为预印本部分数值还需等正式版复核。 5. MixiMotion:一步动作对齐 0.835 逼近教师 MixiMotion: One-Step Text-to-Motion Generation via Asymmetric Set Distillation | Posts and Telecommunications Institute of Technology, Hanoi, Vietnam | arXiv:2609.23010 关键词:文本到动作,一步生成,集合蒸馏,运动学监督,推理延迟 ⚠️ 前序问题:文本到动作生成质量上来了,但迭代采样要几十上百次网络评估,延迟高,交互式动画、游戏这类场景用不了。压缩成严格一步模型有两大障碍:一是文本条件下的动作天然多模态,同一提示对应多个合理动作,固定的教师-学生一一对应会把多种有效模式平均掉;二是只在归一化表征空间里对齐,解码后可能出现明显运动学瑕疵。 本文贡献:MixiMotion 用离线集合蒸馏:对每个提示让冻结的多步教师先生成 K=4 个动作存进离线教师库,学生训练与推理时都不再查教师;学生从独立噪声采样 M=8 个一步样本,用非对称双向集合匹配训练——教师到学生方向鼓励覆盖教师支持的多样动作,学生到教师方向抑制不被支持的生成,两侧权重 α=(1, 0.20) 故意不对称。再叠加解码空间的端点、轨迹与身体分组运动学监督。 MixiMotion: offline asymmetric set distillation for one-step text-to-motion generation. 实验效果:ViMoGen 上语义对齐 0.835,一步方法里第一(MotionLCM-1 0.825、naive 一步教师 0.711),逼近 50 步 HY-Motion-1.0-Lite 教师的 0.858;六个类别中有四类超过 MotionLCM-1。25 人盲评总体 4.33 对教师 4.50,同为一步/少步里最高(MotionHiFlow 4.20、MotionLCM-1 3.98)。单次网络评估延迟从 829.58 ms 降到 9.30 ms。 Qualitative comparisons between HY-Motion-1.0-Lite and MixiMotion. 批判点评:摘要里的 89.2× 是单次网络评估口径,正文自己给出的端到端加速是 6.75×——差了十几倍,文本编码与渲染等开销没算进去,读者很容易把 89.2× 当成真实提速。另外 0.835 其实仍略低于多步的 MLD 0.840(与 MDM 0.833 基本持平),「逼近教师」成立,但严格说还没超过最好的多步基线。消融里非对称权重只带来 0.829→0.835 的 0.006 增益,解码空间监督贡献 0.832→0.835,各组件边际收益都很小;学生从教师热启动、460M 参数与教师同量级,并没有变小。语义分数由 Qwen3-VL-30B 判 yes/no 问题自动给出,非人类标注。 6. SVEET:双向模型改流式 15 FPS Streaming Video Editing with Easy Adaptation | Shanghai Jiao Tong University | arXiv:2609.24788 关键词:流式视频编辑,双向转自回归,特征解耦,正交解耦训练,实时生成 ⚠️ 前序问题:预训练的双向视频扩散模型编辑质量好,但它需要看完整段视频,做不了逐帧自回归的流式编辑。直接改造会撞上两个矛盾:双向骨干的特征在时间上互相依赖,而流式推理要求条件帧彼此独立;可控性与因果性两个优化目标还会互相拉扯,联合训练通常顾此失彼。 本文贡献:先系统复盘已有 video-to-video 扩散方案,归纳出流式适配的两条原则:骨干特征解耦、条件帧独立。据此给预训练双向模型加一条辅助控制分支,用时间独立自注意力编码源视频,中间特征注入对应骨干块;为弥合双向与流式的特征空间差异,再提出解耦训练 ODT,显式约束「视频可控性」与「模型因果性」两个优化方向正交,使二者在推理时兼容,并实现跨异构骨干的零样本迁移。 Our proposed SVEET. 实验效果:视频风格迁移上 VLM 编辑准确率 7.4322 为全表第一(LiveEdit 5.8672、DayDream+CF 5.2321、SDV2 4.2297),CLIP 文本对齐 0.2287、主体一致性 0.9447、背景一致性 0.9298、运动平滑 0.9898、美学质量 0.5550 均为第一。消融中 ODT 把 2D 版从 7.1898 提到 7.4322、3D 版从 7.0653 提到 7.3315,推理期投影(6.4250)与两阶段 teacher forcing(7.0927)都不如它。单张 H100 上 15 FPS,未用任何额外加速手段。 Qualitative comparison results on stylization. 批判点评:唯一没拿第一的一列是整体一致性(0.1123 第二),而拿下这列的 Channel concat 编辑准确率只有 1.8902——几乎没在编辑,说明这个指标在这组对比里更接近「越不改越高分」,拿它当唯一失分项反而暴露了指标本身失效。实验只在风格迁移这一个任务上给了完整数字,其余编辑任务的优势主要靠定性图支撑。15 FPS 是在 H100 上测的,离真正的端侧实时还有距离,且论文没有给出不同分辨率或时长下的 FPS 曲线。 7. COT-TTS:听对话推断语气 MOS 4.15 COT-TTS: Audio Context-Aware Text-to-Speech with Chain-of-Thought Reasoning | The Hong Kong University of Science and Technology;Nanjing University;JIUTIAN Research, China Mobile;Peking University;China Mobile (Hong Kong) Innovation Research Institute | arXiv:2609.22697 关键词:上下文感知语音合成,链式推理,可编辑风格,级联对比,时长一致性 ⚠️ 前序问题:语音合成的说话方式现在基本靠用户显式下指令,但自然对话里该怎么说本应从上下文推断出来。现有系统要么让用户逐句指定,要么把 ASR、LLM、TTS 串成级联——参数动辄 30B 以上,而且级联转换会把原始语音里的副语言信息(非语言发声、节奏、情绪强度)丢掉。 本文贡献:把「上下文感知推理式 TTS」定义成新任务:给定历史对话音频、目标文本和参考语音,先理解上下文,再显式推理出一段中间 CoT,最后按指定音色合成目标语音,且这段 CoT 可被检查与编辑。配套建了大规模双语对话语音数据集(900 万训练样本,含 100 万高质量子集)和 800 条人工核验、源不重叠的评测基准,并给出 0.6B / 1.7B 端到端自回归模型,直接产出带情绪标注的转写、可编辑风格推理和语音 token。 Overview of the end-to-end CoT-guided autoregressive model. 实验效果:英文基准上 COT-TTS-0.6B 拿到 Human MOS 4.150 全表第一(级联 3.050–3.500),情绪一致性 0.954 第一,时长误差 1.155 秒远优于级联的 5.1–12.5 秒;中文基准上 0.6B Human MOS 4.200 同样第一。可编辑 CoT 变体在中文上拿到最高的 LLM 综合分 3.461 与历史理解 4.079。参数量只有级联系统(34–39B)的几十分之一。 Distributions of the normalized audio quality score, naturalness score, target-audio effective-speech ratio, and emotional expression intensity over 50K randomly sampled examples. 批判点评:「参数少几十倍、效果相当」要看是哪一列:UTMOSv2 2.943 低于 two-a3b-fish 的 3.138,DNSMOSPro 3.180 明显不及 two-qwen3omni-seedvc 的 4.240,WER 0.041 也高于 three-dia-a3b-fish 的 0.012;LLM 评测的 CoT 逻辑 4.558 不及 4.836,综合分 2.304 不及 3.601。真正领先的是人类整体评分、情绪一致性和时长误差三项。更值得注意的是两个反常:0.6B 的人类评分(4.150 / 4.200)高于 1.7B(4.050 / 4.150),规模变大反而变差;开启可编辑 CoT 后推理分数上去了,英文 Human MOS 却从 4.150 掉到 3.500——论文把这解释为改动把分布推离训练域,等于承认「可编辑」目前要付音质代价。 8. Common Cause:灰度视频让音频翻车 62% Common Cause, Not Cross-Attention: Blocking Visual Shortcuts in Audio-Video Generation | RIKEN iTHEMS;RIKEN AIP;South China University of Technology;Columbia University | arXiv:2609.22361 关键词:音视频生成,因果推断,反事实不变性,视觉捷径,共享表征失效 ⚠️ 前序问题:音视频联合生成器的训练数据里,「看起来是什么」和「听起来是什么」高度相关——某种材质、纹理或物体外观总与某种声音同时出现。这种相关往往是虚假的:模型可以只从外观预测声音,完全不碰「是什么事件」这个因果变量。一旦测试时外观与事件的关联被打破,模型就会直接合成错误事件的声音。 本文贡献:这是一篇受控的因果研究而非又一个模型。先建音视结构因果模型,让音频在构造上独立于视频的干扰外观,再系统检验流行解法——直接 cross-attention、共享 latent、bottleneck、无监督 shared/private 分解、忠实共享先验,证明这些「走公共因」的做法全都抓不住要害,因为共享 latent 不是干预,模型照样拿得到外观代理变量。真正堵住捷径需要对干扰变量做干预:在 SCM 与干预假设下,论文证明反事实不变性是识别因果预测器的充要条件。 The counterfactual intervention do(v:=v') made concrete (Colored-MNIST). 实验效果:在特征向量 SCM、程序化像素视频、真实图像配频谱音频、移动真实数字、条件生成器五个场景验证。反事实惩罚随因果线索变强从 263× 平滑坍缩到 2.7×,说明捷径只在因果线索难读时才被采用;Moving-MNIST 上直接模型 OOD 准确率 0.21 而本文方法 0.995。对真实预训练视频转音频模型 MMAudio 做输入干预:灰度化让 top-1 声音类别翻转 62%、色彩旋转 54%,最小编辑只翻 22%。 The failure, made audible. 批判点评:最尖锐的一段是「修不好」:作者用重着色一致性项把 MMAudio 的速度网络微调 1500 步,色彩旋转翻转率基本没动(62.5% → 63.9%),真正下降的是通用 OOD 敏感度(高斯噪声 83.3% → 70.8%)——这次微调买到的是整体更鲁棒,而不是论文想修的外观捷径,恰好反证诊断与修复之间还有距离。另外结论高度依赖 SCM 与干预假设成立,真实数据上的干预只能做「改写外观」这类近似;MMAudio 实验只有 24 条真实视频、其中 18 条用于微调,样本量偏小。 9. RewardVerse:先写准则再打分 PLCC 0.52 RewardVerse: Rubric-Guided Policy Optimization for Video Reward Modeling | Institute of Automation, Chinese Academy of Sciences;University of Chinese Academy of Sciences;The Hong Kong University of Science and Technology;Tencent | arXiv:2609.22947 关键词:视频奖励模型,评分准则中介,标量漂移,RGPO,强化学习奖励 ⚠️ 前序问题:用强化学习优化视频生成模型时,奖励模型是地基,但现有视频奖励模型直接把复杂主观质量映射成一个标量,没有显式评分准则,于是出现标量漂移:分数尺度在不同提示之间坍缩或整体平移,RL 拿这种奖励训练并不稳。 本文贡献:借鉴人工标注工程,在「评测请求」和「打分器」之间插入一个动态评分准则 rubric 作为中间表示:先显式生成评测标准,再按准则打分,给标量一个语义锚点。为训练这条协作流水线提出两阶段 RGPO:阶段一用自演化的种子 rubric 预热打分器(偏好奖励 + 格式奖励 + 校准损失),阶段二联合优化 rubric 生成器、产出随请求自适应的准则,同时持续把打分器对齐人类评分。 Overview of RewardVerse and its two-stage RGPO training. 实验效果:16 维 EvalVerse 基准上 Joint(Ours) 在 14 个维度取得最高 PLCC,宏平均 PLCC 0.520、SRCC 0.493 均为第一(次优 Q-Scorer 0.467 / 0.354);Logic 维度 PLCC 0.750 对次优 0.593,Action 0.566 对 0.390。漂移诊断上 RGPO 把均值偏差从 +0.610 降到 +0.164,PLCC/SRCC 从 0.240/0.222 提到 0.505/0.458,预测标准差从 0.457 扩到 1.332,被压缩的分数区间被重新拉开。换到闭源 Gemini-3.1-Pro 上,加 rubric 仍让 PLCC 从 0.490 提到 0.593。 Score distributions of six reward models on the shared 10-dimensional test subset. 批判点评:「14/16 维最高」是真的,但剩下两维恰好是最容易被低层画质解释的:Temporal 那列 Ours 0.538 不及 VideoScore2 的 0.697,SRCC 侧也有两列被 Q-Scorer 与 Raw-rubric 超过。也就是说在偏感知质量而非认知/时序的维度上,rubric 中介的优势会收窄——论文自己也点出 Q-Scorer 在光照、逻辑这类与低层画质相关的维度上仍有竞争力。Gemini 实验只对比了有/无 rubric 两个变体,且受限于 API 只能读自然语言浮点输出,无法与开源模型在同一训练协议下横比。此外 rubric 由模型自生成,其质量评估仍依赖同一套 VLM 裁判,存在自证循环的风险。 10. IMPLICIT-Bench:5493 条中性提示词测隐式偏见 IMPLICIT-Bench: Measuring Implicit Bias in Text-to-Image Models under Neutral Prompts | The University of Melbourne | arXiv:2609.24228 关键词:文生图偏见评测,隐式偏见,知识图谱三元组,中性提示,去偏见权衡 ⚠️ 前序问题:文生图模型的偏见评测大多用「一张 [职业] 的照片」这类槽位模板,只能孤立地探显性人口属性(性别、肤色)。它漏掉了更隐蔽的一类:在自然提示里,当与刻板印象相关的属性没有被指定时,模型仍然默认输出刻板结果——这种隐式偏见现有评测基本看不见。 本文贡献:用结构化知识图谱构造受控提示三元组:中性、刻板、反刻板三个变体只在单一偏见维度上不同,其余场景语义保持一致,从而能把偏见效应精确归因到模板类基准做不到的粒度。最终发布 5,493 条提示、覆盖 11 个偏见类别,并用多模型一致性、CLIP 验证和人工评测三重校验。基于它评测现有去偏见方法,揭示出偏见降低与语义保真之间的基本权衡。 Four-stage KG-grounded benchmark pipeline. 实验效果:过滤后保留的中性图显著偏向刻板端:VLM 均值从 1.75 升到 3.40(Qwen3-VL)、1.26 升到 2.65(Gemma-4),未参与筛选的 CLIP 验证器同向移动。最终基准上 CLIP 判定 67.7% 的 Δx 为正(Cohen's d = +0.39),Qwen3-VL 71.6%(d = +0.85),三个裁判在 Δx 层面正相关(Pearson r = 0.22–0.44)。12 名人工评审在 100 个样例上复现了同样排序:刻板 3.78 > 中性 3.17 > 反刻板 1.60,75.9% 的 KG 被认为确实反映社会刻板印象。 Cross-model qualitative comparison, set 1. 批判点评:三个裁判的一致性是这个方法最脆弱的一环:Gemma-4 的倾向比例只有 48.7%,几乎等于抛硬币,d 也只有 +0.29,与 CLIP 的 67.7%、Qwen3-VL 的 71.6% 不在一个量级;裁判间 Pearson r 仅 0.22–0.44,CLIP 与两个构建裁判的逐图相关性也只有 0.42 / 0.55。人工侧同样不完美:VLM 相对人类均值的 MAE 接近 0.97–0.99,且 12%–16% 的样例把刻板与反刻板的方向判反。样本分布也不均衡——physical-appearance 类只有 14 条,人工研究里 sexual-orientation 一行只落在单个样例上。作者自列的局限同样实在:KG schema 单轴、源数据带 WEIRD 偏斜、去偏见实验只在 Qwen-Image 上做过。 趋势观察 世界模型这一轮在比「记得住」,不是在比「画得好」 WorldCrafter 让目标相机位姿去决定历史怎么压成固定数量的 token,GAE 则干脆把几何基座的特征重参数化成生成用的隐空间。两者承认的是同一件事:分钟级的一致性不是把模型做大就能换来的,而是取决于历史和几何以什么形式存在。有意思的是 WorldCrafter 的加速版在重访一致性上反超了标准版,说明步数与一致性之间未必是正相关。 像素扩散与 latent 扩散的差距,正被「表征」而不是「架构」填平 PixelDiT2 不引入自编码器,而是让一个冻结的视觉编码器在整个去噪过程里持续提供 per-patch 表征,把表征学习和像素生成拆成两件事。不过两张主表横着读会发现,1.46 是像素侧最好、却仍不及 latent 侧 RAE-XL 的 1.13,IS 也不及 REPA-XL 的 306.3——差距在缩小,但没有消失。 加速方案的瓶颈已经从「算力」挪到了「监督」 SparkDiffusion 指出高稀疏下逐步损失在降、终态质量却不涨,根因是监督没有对齐终态;SVEET 则指出双向改流式的瓶颈不在算力而在两个优化方向互相干扰,于是显式约束它们正交。两者都不是靠新算子取胜,而是重新安排了训练时究竟该优化什么。 评测类工作开始把「裁判本身」当成被审对象 IMPLICIT-Bench 的三个裁判一致性只有 Pearson r = 0.22–0.44,其中一个几乎等于抛硬币;RewardVerse 则发现外部打分器在认知类维度上可以是负相关(VideoScore-v1.1 在 Logic 上 PLCC −0.245)。当奖励模型和偏见裁判都开始被自己的盲区反噬,评测可信度就成了需要单独论证的事,而不是默认前提。 人工智能炼丹君 整理 | 2026-09-23 更多 AIGC 论文解读,关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注
2026年09月23日
18 阅读
0 评论
0 点赞
2026-09-22
AIGC 每日速读|2026-09-22|阶跃星辰全双工98.9登顶,清华W4A4逼近FP16
今日 AIGC 论文速览 今日共 10 篇 · 语音与音频生成 2 篇 · 视频生成与物理一致 2 篇 · 图像编辑与三维生成 2 篇 · 生成理解一体化 2 篇 · 推理加速与评测 2 篇 重点论文标题列表 StepAudio 3(StepFun-Audio Team):全双工语音模型98.9分登顶 OmniVChat(CUHK):全模态对话视频2800题 ⚡ CompAdapt(HIT):物理视频生成PIS 0.862 KaiNinja(Alaya Lab):3D部件生成免分割器 4DGS-Fixer(Goertek Alpha Labs):稀疏四视角重建4D高斯 今日论文速览 1. StepAudio 3:全双工语音模型98.9分登顶 StepAudio 3 Realtime Technical Report | StepFun-Audio Team | arXiv:2609.14005 关键词:实时语音交互,全双工对话,边想边说,音频语言模型,MTP解码 ⚠️ 前序问题:实时语音交互要同时满足三件互相拉扯的事:听懂语音里丰富的声学线索(意图、情绪、环境),在对方还没说完时就能自然接话(停顿、附和、打断),以及先做足够深的推理再开口。传统做法要么把推理挪到对话之外、延迟跟着涨,要么为了低延迟直接砍掉推理深度——「想得深」和「答得快」长期被当成一对不可兼得的取舍。 本文贡献:StepAudio 3 Realtime 把整套系统组织成一个连续的「听—对话—思考—行动」循环。三个关键部件:Deep Perception 从原始音频里抽取丰富的声学线索来理解意图;Seamless Duplex 对同步的用户/模型双路音频流建模,处理停顿、附和与打断;核心创新是 Think-While-Speaking——让私有推理与语音输出并行执行,并用 Medusa 式的 MTP 解码加速思考,从而把「想」这一步塞进「说」的同一段时间窗里。训练侧配了一个 Adaptive Thinking 路由,先判断当前这一轮是否需要显式推理,若不需要就走 direct route、留一个空的 think block,避免为简单对话付出推理成本。此外集成了Voice Agent,在对话流不中断的前提下异步执行工具调用,把工具返回结果再融回对话。 实验效果:按官方技术报告的自测:推理模式下 StepAudioChat 宏平均 73.0,达到专用推理模型的水平;MMSU 90.6;Artificial Analysis Full-Duplex Bench Overall 98.9;τ-Voice 宏任务成功率 56.0%。同表里 ASR 侧 StepAudio 3 ASR Max 在 LibriSpeech clean 1.18、WenetSpeech meeting 4.35、AISHELL-1 0.49,三项均为全表最低错误率;音频理解侧 MMAR 86.5、AudioMultiChallenge 49.3;通用能力侧 HMMT 86.8、GPQA Diamond 83.0、MultiChallenge 59.7。 批判点评:优势几乎全部集中在「实时语音交互」这一列:Full-Duplex Bench 98.9 领先 Doubao 2.0 Lite 的95.3 与同表 Gemini 系列的 95.1~98.4,ASR 三项也确实拿到最低错误率。但同一张表暴露了明确的短板——AudioMultiChallenge 49.3,比同表最好成绩 67.0 落后近 18 分;GPQA Diamond 83.0 低于同表最高的 90.3;MultiChallenge 59.7 也落后最好成绩 68.1。τ-Voice 上 56.0 与 Doubao 2.0 Lite 的 56.5 基本持平甚至略低,说明「边想边说」这套机制在需要长链条决策的智能体任务上还没转化成稳定优势。此外全部数字均为团队自评、无第三方复现,报告也没有给出 Think-While-Speaking 与串行思考在同等算力预算下的严格对照——增益里究竟有多少来自 MTP 解码加速、多少来自推理与发声的并行重叠,需要拆开才能判断架构本身的净贡献。 2. OmniVChat:全模态对话视频2800题 OmniVChat: Synthesizing, Benchmarking, and Training for Native Audio-Visual Dialogue | The Chinese University of Hong Kong;Alibaba Token Hub;Shanghai Jiao Tong University;Shanghai Innovation Institute;Zhejiang University | arXiv:2609.21465 关键词:音视频对话,全模态模型,合成数据,强化学习奖励,基准评测 ⚠️ 前序问题:让全模态模型像人一样「看着你、听你说话就直接回应」,而不是先把语音转成文字、把视频转成描述再喂给语言模型。直接吃音视频能省掉转写带来的外部延迟与算力,也保住了表情、环境、手边物品这类感知线索。但这条路有两个卡点:一是人们用自己设备录的真实对话数据极其稀缺;二是同一个意图有无数种说法,好的回答往往要结合用户所处的环境与表情,用关键词匹配根本判不出回答质量。 本文贡献:论文把这条路线正式命名为 OmniVChat,并配齐了数据、评测、训练三件套。数据侧 OmniVChat-Studio 是一个多智能体合成引擎,用智能体协作加视频生成造出单轮与多轮的音视频对话;评测侧 OmniVChat-Bench 覆盖 5 大能力类、17 个子类、22 个场景域共 2800 条合成实例,另有人工实录的 OmniVChat-Bench-Human(360 条)作为真实分布对照;训练侧 OmniVChat-RL 给出一套把回答正确性、效率与风格联合起来的 rubric 奖励设计,直接在合成对话上训练 Qwen3-Omni-Instruct,验证奖励设计能否迁移到真实录制的对话上。 实验效果:RL 训练满 1000 步(记录 51 个 checkpoint)后,OmniVChat-Bench 与人工录制的 OmniVChat-Bench-Human 命中率同步上涨,且合成集与实录集的曲线走势一致,论文还逐条给出了 16 个子类的学习曲线,说明在合成数据上得到的增益确实能迁移到真实对话。对比表里 Qwen3-Omni-Instruct 基线的综合命中率 0.186,训练后与 Qwen3.5-Omni-Plus 的 0.361、Doubao Seed 2.0 Lite 的 0.330 进入同一档。 批判点评:最值得警惕的是奖励设计里的「效率」项。训练曲线显示 Reply Efficiency(每千词的 rubric 命中数)从约 5 涨到 20 的同时,平均回复长度从接近 100 词一路掉到 35 词左右——效率提升有相当一部分来自「少说几句」而不是「说得更准」,在固定 rubric 打分下这几乎必然发生,机制上和同期的 reference-free 语音指标那篇讲的 reward hacking 完全一致。第二,OmniVChat-Bench 全部是合成实例,而实录集只有 360 条,并且只覆盖十二个单轮子类中的一部分,所以「迁移到真实对话」的证据强度受限于实录集规模——论文自己也在多个子类上标注 Bench Only,意味着这些能力在真实录制条件下根本没有对照数据。第三,用合成对话同时做训练和评测,存在评测集与训练分布同源的风险,跨数据集的零样本能力没有被检验。 3. CompAdapt:物理视频生成PIS 0.862 CompAdapt: Adaptable Composite Motion Modeling for Physics-Consistent Text-to-Video Generation | Harbin Institute of Technology;Alibaba Taobao | arXiv:2609.21455 关键词:文生视频,物理一致性,复合运动建模,神经动力学,一次性适配 ⚠️ 前序问题:扩散式文生视频能出画质、也能出时序连贯,但常常违反基本物理:物体不受重力、碰撞后穿模、运动类型中途乱切。已有的物理约束方法大多只会处理单一运动类型,还要人工指定物理参数,换一个没见过的物理环境就失灵。问题在于真实场景里的运动几乎都是复合的——耦合、多阶段转换、多物体碰撞,单一类型的动力学先验覆盖不了。 本文贡献:CompAdapt 把神经动力学建模从单一运动类型扩展到复合物理行为。架构上有两个关键模块:Motion Type Composite(MTC)在潜空间里把多个运动类型按时间关系做算子级复合,Multi-Object Interaction(MOT)在检测到物体接触时切到显式的交互动力学。自然语言提示被解析成结构化的物理语义——运动类型、时序关系、初始物理参数——实现端到端指定,不再需要人工调参。跨新物理环境靠 dynamics-aware prior matching 做一次性适配,不必重训核心动力学模块;另有 physics-aware 潜特征融合模块负责高速与复杂运动下的画面保真。 实验效果:消融实验:去掉 MTC 后 PIS 掉到 0.615,去掉时序建模掉到 0.589,去掉双 LLM 语义解析掉到 0.724,完整版 0.862;FID 15.0 / FVD 91.4 / PSNR 17.2 / SSIM 0.61,相比 Wan-Move(16.3 / 93.5 / 16.5 / 0.58)与 Go-with-the-Flow(17.5 / 96.7 / 14.9 / 0.56)在四项上均有改善,生成开销 2.9 与 Wan-Move 持平。定性部分把 NewtonGen、PhysT2V、Wan2.2、CogVideoX-5B、Veo3、Sora 全部拉到同一组提示下逐帧对照,覆盖抛体旋转、阻尼振荡、均匀加速、形变碰撞等组合运动。 批判点评:三处需要打折。第一,主指标 PIS(物理一致性分数)由论文自行定义,没有第三方榜单或社区复现支撑,而它在消融中的跨度(0.589 → 0.862)恰恰是全文最大的增益来源,指标自证的风险较高。第二,所谓「复合运动」实际只覆盖 9 类情形,其中真正涉及物体接触的弹性碰撞仅 2 类,多物体场景仍是最薄的一环,所谓 multi-object collisions 的证据量偏小。第三,「一次性适配」并不等于「免数据」——流程本身仍需要目标环境的一段观测视频(one-shot reference)才能完成 prior matching,因此把它当作通用物理先验来用还有前置成本。论文 comments 自注为 NeurIPS 2026 投稿(23 页、4 图),尚未经同行评审。 4. KaiNinja:3D部件生成免分割器 KaiNinja: Extending Native 3D Generators to the Part Level | Alaya Lab;The University of Tokyo;UC Merced;Institute of Science Tokyo | arXiv:2609.15659 关键词:三维生成,部件级建模,双体积表示,稀疏体素,免分割 ⚠️ 前序问题:TRELLIS.2 这类原生 3D 生成器能把一张图变成高保真、带材质、非水密的网格,但输出是一个融合好的整体;而下游的编辑、绑定骨骼、物理仿真全都按部件操作。最直接的办法是在融合网格上再跑一个 3D 分割网络,可这样既慢、又被分割精度卡死。真正的技术障碍更基础:O-Voxel 网格每个体素只存一片表面,所以单个 volume 在任何分辨率下都无法表示两个部件贴合的那层界面。 本文贡献:论文提出双体积(dual-volume)表示来补上缺失的界面——同一位置存两套体积,让贴合面两侧各自有归属。KaiNinja 建立在 TRELLIS.2 的 O-Voxel 之上,分两阶段生成:Stage 1 的 Layout Flow 把 3D 噪声网格去噪成逐 volume 的 occupancy 布局,两条流各自有自注意力、再用零初始化的全局注意力块跨流通信;Stage 2 的 SLat Flow 复用 TRELLIS.2 权重,只在活跃体素上对稀疏噪声去噪,最后由 FVAE 解码成 O-Voxel 体积并装配成部件分离的网格。整条流水线里没有掩码、也没有分割器。训练数据来源混杂,包含 CAD 模型和由 LLM 驱动的智能体创作的资产,作者称这是首个用智能体创作的部件数据训练的 3D 生成模型。 实验效果:相对不同范式的部件生成流水线,整体 Chamfer distance 降低 40%、严格部件 F-score 提升 16%;在held-out 测试集与训练未出现过的源上,KaiNinja 产出的部件数量都最接近真值,而 TRELLIS.2+X-Part、Hunyuan3D-2.1+X-Part、OmniPart、PartPacker 这几个基线常常把部件融在一起或者碎开。作者还报告一个反直觉现象:同样的骨干、同样的数据,加上部件级监督之后整体几何保真度反而也提升了。 批判点评:「免分割器」是这篇最硬的卖点,而它恰好没有回答分割器方案真正的软肋——精度上限。论文比的是部件数量、Chamfer 距离与 F-score,没有给出「部件语义是否对得上真值」的证据:在训练未出现的源上,KaiNinja 的部件着色与真值明显对不齐(椅子、花瓶、水壶几行都能看出切分位置与真值不同),也就是说「切得开」和「切得对」之间仍有距离,而语义正确性恰恰是分割器方案此前唯一能保证的东西。其次,40% / 16% 相对的是各基线外挂 X-Part 分割器后的组合结果,比的是「你的流水线 vs 别人的流水线」,而不是「dual-volume vs legacy O-Voxel」,无法分离表示本身的净贡献。第三,所谓「首个用智能体创作数据训练」目前只是自述,论文没有给出这类数据在训练集中的占比,也没有做数据来源的消融,因此这个卖点暂时无法验证。 5. 4DGS-Fixer:稀疏四视角重建4D高斯 4DGS-Fixer: Generative Sparse-View 4D Gaussian Splatting with Iterative Refinement Guided by Video Diffusion Priors | Goertek Alpha Labs;Singapore Institute of Technology;The Hong Kong University of Science and Technology | arXiv:2609.21176 关键词:4D高斯泼溅,稀疏视角,视频扩散先验,迭代精修,动态场景重建 ⚠️ 前序问题:用稀疏视角视频重建动态场景是典型的病态问题:观测太少、信息缺失,再多优化技巧也补不出来。更具体的症结在初始化——只有几个输入视角时,COLMAP 只能重建出稀疏残缺的点云,大片区域没有足够的高斯基元支撑,后续优化根本无从下手。也就是说问题出在起点,不在优化器。 本文贡献:论文给出一个两阶段的迭代精修框架。Stage 1 初始化:先对多路同步视频估计多视角深度图,把它们融合成稠密点云,为动态 4DGS 提供更完整的几何起点,替代 COLMAP 的稀疏点云。Stage 2 精修:沿新的相机轨迹渲染出视频,送进一个预训练的视频修复/恢复模型得到「修复后的伪目标视频」,再把它当作伪监督去迭代更新 4DGS 表示,多轮往复。整条链路的思路是用视频扩散先验补足缺失的观测信息,而不是继续在几何约束上加码。 实验效果:在常用的稀疏视角动态基准上,相对此前最好的方法 PSNR 提升接近 2 dB,作者称 substantially outperforms 现有基线。定性对比给出 GT / 4C4D / Ours 三行并排结果,4DGS-Fixer 在人物细节与背景完整度上都比 4C4D 明显更干净、更少孔洞。论文 comments 标注已接收至 SIGGRAPH Asia TC。 批判点评:整条链路的关键一步是「拿视频修复模型的输出当监督」,这带来一个绕不开的疑问:修复模型的幻觉会被原样写进 4DGS。用修复后的序列做伪监督,本质上是把生成先验的偏好当成真值——一旦输入视角极少、修复模型开始「编」,4DGS 就会忠实地记住这些编出来的内容;而 PSNR 的提升恰恰可能来自「补全得更像训练分布」而不是「更接近真实场景」。论文没有做「关掉修复模型、只保留深度稠密初始化」的对照,因此也说不清 2 dB 里有多少来自几何初始化、多少来自扩散先验。其次是部署成本:方法需要四路同步相机加一个预训练视频修复模型,实时性与泛化到非受控场景的能力都未讨论。源码层面还有个小瑕疵——LaTeX 里仍留着旧名「STGFixer pipeline」的残句没清理,说明从 STGFixer 改名到 4DGS-Fixer 之后并未完全同步。 6. RefineEdit:九类编辑背景保持最优 Refinement Is Inherently Editable: Training-Free Prompt-to-Prompt Image Editing with Generative Refinement Network | City University of Hong Kong (Dongguan);City University of Hong Kong;MBZUAI | arXiv:2609.20633 关键词:图像编辑,免训练,生成式精修网络,二值图像码,背景保持 ⚠️ 前序问题:文本引导的图像编辑要在「改对」和「别乱改」之间走钢丝。扩散类编辑器靠空间控制(掩码、注意力图)来圈定改动范围,而掩码一旦不准,结果不是改不彻底就是动到无关区域;因果自回归类编辑器还有一层额外约束——解码顺序固定,前面已经决定的 token 后面没法回头修正,早期的一个坏决策会被一路带下去。 本文贡献:论文换了一个编辑载体:不去控制扩散过程,而是在 Generative Refinement Network 的二值图像码上做全局精修。核心机制是双分支耦合——编辑分支从源轨迹的某个中间状态初始化,复用已经成形的布局;两个分支对同一批源采样 bit 分别给出概率,用它们的带符号差值挑出「该改的位置和 bit」。被选中的 bit 进入编辑精修,其余 bit 复制源分支持续演化的状态,于是「定位」与「生成」被绑在同一个精修过程里,编辑证据可以随图像演化不断被重新评估。为跨精修步稳定,论文加了两道约束:adaptive spatial freezing 抑制掩码无谓扩张,finite bit locking 让最近选中的 bit 保持可编辑。整套方案免训练、不需要外部掩码,也不需要注意力控制。 实验效果:在 PIE-Bench 的九类编辑任务上,RefineEdit 在 PSNR / LPIPS / MSE / SSIM 四项背景保持指标上全部最好,同时整图 CLIP 与被编辑区域 CLIP 也最高——「改得准」与「背景不崩」这次没有被做成取舍。定性部分与 LEDits++、FlowEdit、ChordEdit、MasaCtrl 在替换、添加、擦除、属性/姿态/颜色/材质/背景修改与风格迁移九类上逐一并排对比。 批判点评:最大的边界条件藏在实验设置里:这套流程要求源图来自 GRN 自己的生成轨迹——编辑分支要「继承源状态」,因此它编辑不了任意真实照片。论文在 limitation 里自己承认了这一点,但对普通用户而言,「只能改模型自己生成的图」和「能改我相册里的照片」是两个完全不同的产品。其次是超参:switch step(从第几步开始分叉)以及空间/bit 两个阈值都需要逐类调整,而且最优值是在同一个 PIE-Bench 评测集上选出来的,相当于用测试集选超参;论文的 limitation 也承认这些类别特定设置是在该基准上挑的,跨数据集泛化没有被验证。第三,九类各自用自己那一档配置,但真实用户不会先告诉系统「这次是换材质还是换背景」,多类混合的编辑场景仍未被覆盖。 7. QuAKE:W4A4 文生图量化误差校正 Quantization-Aware Kalman Estimation for Diffusion Sampling | Tsinghua University | arXiv:2609.21407 关键词:扩散模型量化,推理加速,W4A4,卡尔曼滤波,采样校正 ⚠️ 前序问题:量化是扩散模型落地的现实路径,但 W4A4 这种激进压缩会让量化去噪器的输出明显偏离全精度版本,而误差还会沿时间步持续传播、越积越大。已有的采样期校正方法基本只看当前步的局部信息,没有利用整条轨迹的历史——而误差恰恰是跨时间步传播的,用局部信息去修一个全局累积的问题,方向就不对。 本文贡献:论文把「用量化去噪器采样」重新表述成一个在线估计问题:既然真实的全精度输出拿不到,那就用历史量化输出去把它估回来。QuAKE 是为此设计的量化感知 Kalman 估计器,用一个平滑轨迹先验叠加一个条件高斯观测模型,每一步以闭式解递归更新「输出窗口」的后验,再把后验均值喂给采样器。它不改动量化网络、天然支持任意高阶多步 ODE 采样器,是一个即插即用的轻量校正器。论文同时给出了量化噪声的实证统计(边缘分布接近高斯、联合分布呈明显相关结构、条件分布可线性拟合),用来说明为什么 Kalman 这一套假设在这类噪声上站得住。 实验效果:在多个 W4A4 量化的文生图模型上,QuAKE 在「与全精度采样的分布差异」这一指标上持续领先:FLUX.1-dev + MJHQ-30K 上 FID 7.02 / KID 0.158,优于朴素量化的 7.30 / 0.257、TAC-Diffusion 的 7.31 / 0.304、Q-Drift 的 7.16 / 0.209;Sana-0.6B 上 7.90 / 0.289;PixArt-α 上 13.71 / 1.987;PixArt-Σ 上 15.99 / 3.401。开销按论文自己的测量几乎免费:FLUX.1-dev 单步延迟 22.2 ms(BF16)降到 6.2 ms(INT4 与 QuAKE 相同),DiT 体积 160 GiB 降到 102 / 104 GiB,端到端 23.51 s 降到 2.20 / 2.24 s。 批判点评:有一个容易被标题盖住的细节:QuAKE 赢的是分布距离(FID / KID),但在感知质量列上并不总是赢朴素量化。以 FLUX.1-dev + MJHQ-30K 为例,IR 分数 QuAKE 0.8822 低于 Quantized 的 0.8887,CLIP-IQA 0.9042 也低于 0.9057——也就是说 Kalman 校正主要抹平的是低层分布偏差,并没有让图在感知偏好上更讨喜;如果优化目标是下游打分模型或人眼偏好,朴素 INT4 甚至可能更安全。这不是论文的错,但摘要里「持续优于已有方法」的表述容易被读成全面胜出。其次,效率那一组三个数字(6.2 vs 6.2 ms、102 vs 104 GiB、2.20 vs 2.24 s)方向一致地说明校正本身几乎免费,反过来也说明它的收益完全来自「估得更准」而不是「跑得更快」——真正的加速来自量化,QuAKE 只负责把量化弄坏的部分补回来。最后,全部实验集中在文生图,视频扩散这类时间步更多、误差累积更长的场景是否同样成立,尚未给出证据。 8. AURA:对话式音乐编辑只训91M AURA: Unified Multimodal Framework for Conversational Music Editing | Aalto University;Technical University of Denmark;University of South Dakota;OpenRB Lab | arXiv:2609.14344 关键词:音乐编辑,多轮对话,多模态大模型,LoRA,概念token ⚠️ 前序问题:现有的指令式音乐编辑器是「一问一答」的:每条指令独立处理,用户想「先弱化鼓、再压一下高频、最后把情绪调柔和」这种逐步打磨的工作流根本走不通,因为模型记不住上一轮改了什么。此外,要改音乐光靠文字往往不够——用户脑子里想的是「这段场景听起来该是什么样」,而不是一串精确的音频术语。 本文贡献:AURA 把音乐编辑变成多轮对话。用一个多模态大模型读完整的对话历史、可选的一张图以及参考音频,把编辑意图蒸馏成一组紧凑的 concept token;再由 concept-to-audio(C2A)模块把这些 token 与逐帧对齐的参考特征注入冻结的 MusicGen 骨干——内部通过交叉注意力与 BiFAM 做对齐,让编辑既精确又不碰无关内容。训练只更新 91M 参数(LoRA rank-64 / α=128 加两个投影器与 C2A),1.9B 骨干全程冻结;对话数据由 Slakh2100 合成,共 66,539 段多轮对话。 实验效果:在域内 Slakh2100 与域外 MoisesDB 上同时评测:Add 任务的域外 FAD 0.84 对 Instruct-MG 的 3.84,Remove 任务域外 FAD 0.72 对 3.55,均为 4~5 倍降低;Remove 任务域内 SI-SDR +11.32 dB、SDRi +4.89,而 Instruct-MG 是 −2.10 / −8.53。消融显示对话历史确实起作用:只把最后一条指令喂给模型,Remove 的 SI-SDR 会掉到 −5.2 dB,SSIM 从 0.84 掉到 0.54。 批判点评:摘要里「域外 4~5 倍 FAD 降低」是有选择性的——它只覆盖 Add 与 Remove 两个任务;在 Extract 上 AURA 是 4.24 对 Instruct-MG 的 5.20,只有约 1.2 倍,并没有复现那个量级。更要紧的是 Extract 的SI-SDR 虽然比基线好得多(−7.62 对 −15.18),但仍是负值,说明「从混音里把目标音轨抽出来」这个任务本身离可用还有距离。其次,Remove 的域内增益(+11.32 dB)看着最亮眼,但该任务的域内评测集就来自训练数据来源域(Slakh2100 合成),换到域外 MoisesDB 只剩 +2.54 dB,跌幅接近 4 倍,真正的跨域能力比表格第一屏呈现的要弱。第三,91M 可训练参数对 1.9B 冻结骨干确实是漂亮的效率数字,但它同时意味着能力上限被 MusicGen-medium 这个骨干锁死,而论文没有给出换更大骨干的对照。 9. Qwen-Audio-3.0-TTS:四代语音模型瓶颈迁移史 The Evolving Bottleneck in Speech Generation: Interface Co-design and Staged Alignment from CosyVoice to Qwen-Audio-3.0-TTS | Alibaba Token Foundry | arXiv:2609.16514 关键词:语音合成,接口设计,瓶颈迁移,流匹配,阶段性对齐 ⚠️ 前序问题:语音合成的发展通常被叙述成「模型更大、tokenizer 更好、数据更多」,但这类叙事解释不了为什么某些改动立竿见影、某些就收效甚微。真正值得追问的是:每一代系统当下的主导瓶颈在哪里——是表征不够内容对齐?是接口不能因果流式?是覆盖度与可学习性不足?还是模型容量与跨模块协调出了问题? 本文贡献:这是一篇技术回顾(technical retrospective),沿 CosyVoice → CosyVoice 2 → CosyVoice 3 → Qwen-Audio-3.0-TTS 这条线给出另一种解释:进步来自瓶颈的反复迁移。四代之间不变的是一条分解——自回归语言模型负责「规划」,flow-matching 模型负责「渲染」;变的是两者之间的接口契约。作者把契约形式化成四个维度:representation(表征)、ownership(归属)、availability(可用性)、gradient reach(梯度可达性),并明确区分「单篇论文内的证据」与「跨论文比较」。每一代针对一个不同的约束:CosyVoice 立下监督式语义 token 这个内容对齐接口;CosyVoice 2 让接口对因果流式可用、把句级说话人嵌入从语言模型里移出;CosyVoice 3 靠多任务监督、数据与模型扩容、以及可微奖励优化(DiffRO)提升接口的可学习性与覆盖;Qwen-Audio-3.0-TTS 把帧率从 25 Hz 降到 12.5 Hz、渲染器改为条件在连续隐状态而非离散 token 嵌入上,并用五阶段课程做联合分阶段对齐。 实验效果:文中给出的实证来自各代公开系统:语义 token 的引入把中文 CER 从 2.56 降到 1.45、英文 WER 从 3.81 降到 2.57(CosyVoice 2 的模块消融);数据从约 17 万小时扩到 100 万小时、语言模型从 0.5B 扩到 1.5B(CosyVoice 3);Qwen-Audio-3.0-TTS 把帧率减半并保留 K=6561 的码本,五阶段课程依次为独立预训练、联合训练与质量退火、语言模型强化学习、声学鲁棒性、flow-matching 强化学习。 批判点评:这篇最需要标注边界的地方,作者自己写在 Figure 1 的图注里:「箭头编码的是我们的回顾性解释,而非受控的跨版本实验」。也就是说「瓶颈迁移」是一个事后叙事,不是被控制变量检验过的因果结论——拿它指导新系统设计时,风险在于你可能把相关性当成了因果。另一个更具体的问题是证据来源:文中复现的 tokenizer 对比(例如把帧率减半同时保持 K=6561 会一并损伤内容一致性与相似度)引自 Xiang et al. (2026),属于「论文内对比」的二次引用,不是独立复现;Table 1 也是同样的性质。因此这篇更适合当作一份结构清晰的领域地图来读——四维接口那套分析框架确实好用,也确实解释了「为什么改 tokenizer 有效、改渲染器无效」这类问题——但不宜当成可直接复制的工程配方。论文 comments 自注为 technical retrospective(11 页),本身也不以实验贡献为目标。 10. Ref-Free Metrics:13类语音指标在干净音频失效 The Limits of Reference-Free Speech Quality Metrics as Evaluators and Rewards on Modern Text-to-Speech | AGIGO;ETH Zurich;Sapienza University of Rome;University of Zurich | arXiv:2609.13150 关键词:语音质量指标,无参考评测,奖励攻击,成对偏好,评测协议 ⚠️ 前序问题:UTMOS、DNSMOS、SCOREQ 这类无参考质量预测器,今天既是 TTS 的默认自动评测器,又越来越常被拿去做偏好优化的奖励信号。这两个角色都建立在同一个假设上:预测分数能跟上人类偏好。但这个假设在什么条件下成立、什么时候会崩,此前没有被系统检验过。 本文贡献:论文把假设直接摆到台面上测。做法是构造一个成对判定任务——给两个音频片段,问预测器打分更高的那个是不是人更偏好的那个——然后在 6 个人工评分语料上跑,覆盖从「瑕疵明显」到「几近无瑕」的完整质量区间。评测对象包括 13 个无参考预测器家族(33 个分数变体)、13 个 Praat 韵律描述符,以及一批经典信号处理特征,每项都给出相对「随机基线」与「人类天花板」的位置。第二部分把这些指标当奖励用,对比单分数奖励与复合奖励在策略优化下的行为差异,并给出数据效率曲线。 实验效果:结论相当刺。在瑕疵明显的 BVCC 上,各指标准确率能爬到人类天花板附近(天花板 0.887,标定的复合分数到 0.92);但一旦两边音频都干净(TTS-HP),没有任何单个预测器能可靠指认被偏好的样本,好几个的准确率甚至低于「直接挑时长更长的那个」这个傻瓜基线(argmax clip duration 在 BVCC 0.517、SOMOS 0.368、SingMOS 0.456、SpeechJudge 0.370、TTS-Arena 0.523、TTS-HP 0.524),复合分数也才 0.52,而人类天花板是 0.764。奖励侧更严重:优化单一分数会诱发 reward hacking——指标冲上去了,但独立留出评测器和人工听测同时变差;复合奖励能抵抗这个倾向。数据效率曲线还给了个实用兜底:每类只取 min(20%, 1000) 条域内成对数据拟合一个带正则的线性头,就能把准确率相对最好的单指标提升 +1.5%(BVCC)到 +6.6%(SOMOS)。 批判点评:这篇的价值是给出一把尺子,而不是给出一个能用的指标——它自己也承认在干净音频上没有任何固定公式可迁移。要留意三点。第一,复合分数需要通过域内偏好标注拟合(论文用了 min(20%, 1000) 条域内成对数据),相对增益只有 +1.5% 到 +6.6%,也就是说「改用复合分数」这个建议的前提是你已经有该领域的偏好数据,对做通用 Reward Model 的人是不小的门槛。第二,它测的是成对偏好准确率,不是绝对 MOS 拟合度,「指标失效」的结论不能直接推成「MOS 预测不准」——在需要绝对分数的场景里结论需要重新验证。第三,时长基线之所以能赢,是因为评测语料里存在系统性的时长偏差,这提示「最好用的信号」可能来自数据偏置而非感知质量;反过来说,这也意味着任何在干净音频上得到的排行榜,都可能只是某个数据偏置的代理指标。 趋势观察 「边想边说」正在把推理的延迟藏起来 StepAudio 3 让私有推理与语音输出并行执行,用 Medusa 式 MTP 解码把思考压进说话的同一段时间窗;QuAKE 则把量化误差的校正重新表述成一个在线估计问题,用轨迹历史去猜全精度输出本该是什么。两者都在回答同一个问题:当「想清楚」和「立刻给」不能同时满足时,能不能把其中一件事塞进另一件事的时间窗里,而不是二选一。 物理一致性开始从「单类型」走向「复合行为」 CompAdapt 把神经动力学从单一运动类型扩到耦合运动、多阶段转换与多物体碰撞,并用一次性适配去应对没见过的物理律;4DGS-Fixer 干脆不追求几何自洽,改用视频扩散先验把缺失的观测补出来当伪监督。一个在约束上加码,一个在数据上放开,方向相反,但都承认了同一件事:稀疏观测下的物理,靠纯几何推不出来。 「部件」正在成为 3D 生成的新交付单位 KaiNinja 指出 O-Voxel 每个体素只存一片表面、单个 volume 在任何分辨率下都表示不了两个部件贴合的那层界面,于是用双体积表示把界面补上,全程不需要掩码或分割器。下游的编辑、绑骨与仿真要的正是部件而不是一整块表面。这提示 3D 生成的下一步未必是把表面做得更细,而是把语义切分做得更明确。 质量指标同时当评测和奖励,正在被自己的盲区反噬 语音侧,Ref-Free Metrics 那篇的 13 个无参考预测器在干净音频上集体失效,好几个甚至不如「挑最长的那个」,而优化单一分数会直接诱发 reward hacking;对话视频侧,OmniVChat 的效率项让平均回复长度从约 100 词掉到 35 词,效率上涨里混着「少说几句」。当同一个分数既当尺子又当奖励,它没测到的那部分就一定会被优化出来。 人工智能炼丹君 整理 | 2026-09-22 更多 AIGC 论文解读,关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注
2026年09月22日
17 阅读
0 评论
0 点赞
2026-09-21
AIGC 基本功|音频质量评测:MOS、PESQ 与 FAD 各测什么-AudioEval
音频质量评测:MOS、PESQ 与 FAD 各测什么 所属方向:评测与指标 | 难度:工程实战 | 前置知识:无 关键词:音频质量、音频评测、MOS、PESQ、STOI、SI-SDR、FAD、Fréchet Audio Distance、感知评测、语音合成、音频codec、无参考评测 01. 为什么需要它 先看一个教学假设:团队比较两个音频 codec 版本,一个版本对齐更准,SI-SDR 更高,却同时新增了低能量的高频纯音。波形误差下降和听感变差可以同时发生。第 04 节用可复现合成信号展示这个指标盲区;本文没有受试者听测,因此不把「刺耳」或「无差别」写成测得的事实。7~7.5 kHz 的窄带成分可能可闻,其显著程度取决于响度、掩蔽、设备和听者;不能把这一段称为人耳最敏感频段。 这个问题不是「评测做少了」,而是指标回答的问题和团队想问的问题不是同一个。SI-SDR 回答的是「输出波形和参考波形逐采样点对得上吗」,团队想问的是「用户听着难受吗」。这两个问题之间的距离,就是这篇要讲的东西。 更麻烦的是,音频指标的「鄙视链」比图像指标更长、更碎:PESQ 面向语音质量,含窄带与宽带版本;STOI 面向语音可懂度,音乐生成用的是分布级的 FAD,而人工 MOS 需要组织听测,并报告受试者、协议与置信区间。没有一个指标是免费的午餐——先搞清楚每个指标在测什么,再决定用哪个,比记住任何一行数字都重要。 02. 最小可用理解 先把这篇要用的缩写认全——音频圈的缩写比图像圈更碎,不先对齐名词,后面每个数字都容易理解错: 缩写 全称 中文 一句话 MOS Mean Opinion Score 平均意见分 找一群人听、打分取平均。按指定听测协议汇总主观评分,有成本和统计不确定性 PESQ Perceptual Evaluation of Speech Quality 语音质量感知评估 语音有参考质量指标,P.862 为窄带、P.862.2 为宽带扩展 STOI Short-Time Objective Intelligibility 短时客观可懂度 只测「内容听不听得清」,不测好不好听 SI-SDR Scale-Invariant Signal-to-Distortion Ratio 尺度不变信号失真比 允许整体缩放后再比波形,衡量逐采样点对齐程度 FAD Fréchet Audio Distance Fréchet 音频距离 不需逐条配对参考,仍需参考音频集合,比较两堆特征分布,思想来自图像的 FID STFT Short-Time Fourier Transform 短时傅里叶变换 把波形切帧做 FFT,得到随时间变化的频谱 SSL Self-Supervised Learning 自监督学习 UTMOS 这类 MOS 预测器用来提特征的预训练模型 DNSMOS / UTMOS — — 用模型预测 MOS 的两种常见方案(前者微软、后者东京大学 SaruLab) 三句话讲完框架: 每个指标只回答一个具体的问题:「波形逐点对齐了吗」(SI-SDR/SNR)、「频谱形状像吗」(对数梅尔距离)、「这两堆音频听起来像同一个世界来的吗」(FAD)、「人听着舒服吗」(MOS,以及预测它的 UTMOS/DNSMOS)。问题不同,答案就不同——不存在「最好的指标」,只存在「问对了问题的指标」。 有参考和无参考是两大阵营。有参考指标(SI-SDR、PESQ、梅尔距离)需要对应的参考信号,适合 codec、增强、分离等任务;逐条无参考指标(UTMOS、DNSMOS)不需要对应参考;FAD 不需要逐条配对,但需要代表目标分布的参考集合。是否适合音乐或音效仍取决于训练域和特征器。 客观指标是 人耳的代理,代理就有失真。失真最大的地方往往是「在部分听测条件下不显著、指标却敏感」的维度(波形相位、时移、整体音量),和「人耳一耳朵就听出来、指标却毫无反应」的维度(高频毛刺、音色瑕疵)。第 04 节会用一张表把这两类失真都造出来给你看。 03. 数学推导 3.1 SI-SDR:为什么它对音量不敏感、对时移却如此脆弱 信号处理里的「失真」最朴素的定义是信噪比: $$\mathrm{SNR} = 10 \log_{10} \frac{\| s \|^{2}}{\| \hat{s} - s \|^{2}}$$ 其中 $s$ 是参考信号、$\hat{s}$ 是待测信号,范数就是逐采样点的能量。它有一个明显的不合理之处:如果输出和参考只差一个整体音量($\hat{s} = 0.95\,s$),按理说听感几乎没差别,但 SNR 会把它记成 $20\log_{10}(1/0.05) \approx 26$ dB 的「失真」。 SI-SDR 的解法是先做最优缩放,再算残差——这也是它名字里 Scale-Invariant 的来历: $$\alpha = \frac{\langle \hat{s}, s \rangle}{\langle s, s \rangle}$$ $$\mathrm{SI\text{-}SDR} = 10 \log_{10} \frac{\| \alpha s \|^{2}}{\| \hat{s} - \alpha s \|^{2}}$$ $\alpha$ 是把待测信号投影到参考方向上的最小二乘系数:分子是内积 $\langle \hat{s}, s \rangle$(两个波形逐点相乘再求和,衡量「同相程度」),分母是参考信号的能量,用来归一化。缩放完再算残差,整体音量的差异就被精确抵消了。 但这个设计同时埋了雷。SI-SDR 衡量的是逐采样点的对齐程度,而波形对齐对时间平移极端敏感:把信号往后挪 8 个采样点(16 kHz 采样率下只有 0.5 毫秒,单独播放时通常不易觉察),波形的峰谷就全部错开了,内积 $\langle \hat{s}, s \rangle$ 急剧下降、残差急剧上升,SI-SDR 可以从正几十分贝直接掉到负值——第 04 节的实验里,0.5 毫秒的时移让 SI-SDR 掉到了 $-6.25$ dB。 所以使用 SI-SDR 前要确认与参考的对齐约定(降噪、回声消除、codec 重建),并且要明白它测的是去掉全局增益后的波形误差,对齐只是影响因素之一。 顺带回答一个高频疑问:为什么叫 SDR 不叫 SNR? 因为这个指标不是从「信噪比」这棵树上长出来的,而是从盲源分离那套评价体系里长出来的。2006 年的 BSS_eval 工具箱(Vincent / Gribonval / Févotte,Performance measurement in blind audio source separation)把估计信号相对参考的总误差拆成四类: 残留干扰(interference):别的声源串进来多少 加性噪声(noise):环境噪声剩多少 算法伪影(artifacts):算法自己制造出来的人工痕迹 通道误差(channel errors):增益、滤波这类整体偏差 对应地给出 SIR / SAR / ISR 三个分项,以及一个总指标——SDR(Signal-to-Distortion Ratio)。 关键就在 Distortion 这个词:它涵盖上面四类,而 SNR 的「噪声」只对应其中一类。在盲源分离里,最伤听感的往往不是噪声,而是别的源串进来的干扰和算法自己的伪影,叫「信噪比」会名不副实。 后来 2019 年 Le Roux 等人在 SDR – half-baked or well done?(arXiv:1811.02508)里指出 BSS_eval 版 SDR 有两个致命问题:bss_eval_sources 允许用 512 阶滤波器去「修改参考信号」来拟合估计值,于是把某些频段直接置零也能拿到近乎无穷的 SDR;bss_eval_images 的 SDR 又退化成了普通 SNR,不允许全局缩放,反被一些算法无意间靠调音量刷了分。他们给出的修法就是本文用的 SI-SDR——误差项与参考信号严格正交,彻底去掉对幅度缩放的依赖,命名沿袭 SDR 保持与社区的连续性。 SDR 的 distortion 可以包括噪声、串扰和伪影等误差。SNR 与 SI-SDR 是否数值相等还取决于缩放、投影和预处理约定,不能仅凭「只有加性噪声」判定。工程上 SI-SNR 常用于同类投影指标,比较实现时须检查是否先去均值、如何处理静音。 3.2 对数梅尔距离:把「人耳怎么听」写进公式 梅尔刻度近似描述感知音高与频率的非线性关系,不能据此断言某两个高频纯音「人耳几乎分不清」。本文采用 HTK 风格的映射: $$m(f) = 2595 \log_{10}\left(1 + \frac{f}{700}\right)$$ 其中 $f$ 的单位为 Hz。短时傅里叶变换(STFT)得到幅度谱后,用三角滤波器聚合到梅尔频带,再取对数。本文的最小实现是梅尔幅度谱,不是常见的梅尔功率谱;两者必须区分。令 $M(k,t)$ 为第 $k$ 带、第 $t$ 帧的幅度,以每条音频的谱峰值 $M_{\max}$ 归一化,定义: $$L(k,t)=\ln\max\left(M(k,t)/M_{\max},10^{-4}\right)$$ $$D_{\mathrm{mel}} = \frac{1}{KT}\sum_{k,t}\left|L(k,t)-\hat L(k,t)\right|$$ $K,T$ 分别为梅尔带数和帧数。这里的单位是自然对数幅度差,不是 dB;换成 dB 要乘 $20/\ln 10$。独立归一化是为了在本实验里忽略整体增益,若任务关心响度,就不能先把这类差异消掉。幅度下限 $10^{-4}$ 对应峰值下 80 dB,防止对接近零的值取对数;若使用功率谱,同样的 80 dB 下限应是 $10^{-8}$。窗口、hop、滤波器、幅度/功率约定、floor 和归一化都会改变数值,需完整报告。 3.3 Fréchet Audio Distance:从「两条音频」到「两堆音频」 前面两个指标都需要一条参考音频。但音乐生成、音效生成没有「正确答案」——你没法说这段生成的钢琴曲应该长什么样。这类任务需要无参考指标,FAD 是其中最常用的一个。 FAD 的思路借自图像领域的 FID:不比较单条音频,而是比较两堆音频的分布。把每条音频过一个特征提取器(原论文用 VGGish,得到 128 维特征),每堆音频的特征就变成了高维空间里的一团点云,然后近似成高斯分布,问:这两个高斯离得多远? FAD 使用两个拟合高斯之间的 2-Wasserstein 距离平方(沿用 Fréchet distance 名称): $$\mathrm{FD} = \left\| \mu_{r} - \mu_{t} \right\|^{2} + \mathrm{Tr}\left( \Sigma_{r} + \Sigma_{t} - 2 \left( \Sigma_{r} \Sigma_{t} \right)^{1/2} \right)$$ $\mu_{r}, \mu_{t}$ 是参考集和测试集的特征均值(每维上特征的平均),$\Sigma_{r}, \Sigma_{t}$ 是协方差矩阵(特征之间的联合波动结构),$\mathrm{Tr}$ 是矩阵的迹(对角线元素之和)。直觉上:均值差衡量「两团云的中心差多远」,协方差项衡量「两团云的形状差多远」。 样本量会影响均值和协方差的估计。当样本数不超过特征维度时,中心化样本协方差必然秩亏;但秩亏不表示距离无法计算。更关键的是估计偏差和方差,其大小取决于特征谱、独立样本数和参考集,而非单由维度决定。第 04 节的 16 维合成例子中,同分布两组的距离从 8 个样本时的 13.1390 降到 300 个样本时的 0.1844;这是一组随机种子的结果,不是通用阈值。 实战里怎么检查样本量的影响 报告同分布基线及其波动。按原始音频为单位反复随机分组,让组大小与真正的比较一致,观察距离分布。不能拿一次基线的「2 倍」当显著性阈值;应结合效应量、重采样区间或预先设计的置换检验。 不要把切片数当独立样本数。0.96 秒窗口是 VGGish 常见的输入单位。帧级统计与 clip 级统计评价的对象不同,同一条音频的窗口又高度相关。切片可以提供局部信息,但不会凭空产生同等数量的独立录音;bootstrap 应按音频或更高层级成组抽样,不能把相关帧当成独立样本。 固定评测协议。报告参考集、生成集规模、总时长、窗口步长、特征器和归一化。更换这些条件后,旧 FAD 数字不再是同一把尺子。 原始 FAD 论文 与 后续 FAD 评测研究 讨论了特征器和样本量偏差。本次修订移除了缺少原始音频与评测脚本的历史「真实案例」图表;下面保留可由附录完整复现的合成实验,避免把不可复核的数值当作结论依据。 04. 代码实现 完整脚本在文末附录,这里只看核心片段与真实输出。数值脚本只用 numpy;配图脚本另需 matplotlib,不依赖音频模型权重。 第一段:构造三种不同机制的退化(metric_disagreement.py)。干净信号用「基频 + 6 个谐波 + 慢包络」合成;三种退化分别是音量 ×0.95、叠 7/7.5 kHz 纯音、循环时移 0.5 ms: def degrade(x, kind, sr=SR): if kind == "A": return 0.95 * x # 音量差 5% if kind == "B": # 7/7.5 kHz 纯音,均低于 16 kHz 采样率的 Nyquist 频率 t = np.arange(len(x)) / sr buzz = 0.02 * ( np.sin(2 * np.pi * 7000 * t) + np.sin(2 * np.pi * 7500 * t) ) return x + buzz if kind == "C": return np.roll(x, 8) # 时移 8 采样 = 0.5ms raise ValueError(kind) 第二段:SI-SDR 的最小实现(缩放抵消就三行): def si_sdr(ref, est): """Scale-Invariant SDR。 先把 est 投影到 ref 上(找一个最优缩放 alpha),剩下的残差才算失真。 所以「整体音量变了」在 SI-SDR 里被完全抵消——这是设计意图, 但也意味着它测不出增益错误。 """ if np.dot(ref, ref) <= 0 or np.dot(est, est) <= 0: raise ValueError("SI-SDR needs non-silent reference and estimate") alpha = np.dot(est, ref) / np.dot(ref, ref) e_target = alpha * ref e_res = est - e_target num = np.dot(e_target, e_target) den = np.dot(e_res, e_res) if den <= 1e-24 * num: return float("inf"), alpha return 10 * np.log10(num / den), alpha 真实运行输出(16 kHz、1 秒合成乐音): 操作 SI-SDR(dB) SNR(dB) 梅尔距离(ln幅度) 高频占比 高频谱平坦度 A 音量×0.95 inf 26.02 0.0000 3.808e-14 0.3256 B 7/7.5kHz纯音 26.58 26.58 0.3472 2.188e-03 近0 C 循环时移0.5ms -6.25 -0.51 0.0187 3.836e-14 0.3272 A 的全局增益在 SI-SDR 和独立归一化的梅尔谱里被消除;SNR 仍会计入它。$\alpha=0.95$ 表示待测信号为参考的 0.95 倍,并不是把 0.95「倒过来」。 B 的 SI-SDR 为 26.58 dB,但它不是主观音质的合格线。高频能量占比明显上升,说明多了高频成分;这些成分是否不悦耳,需要听测。 C 在波形上错位,SI-SDR 大幅下降,梅尔距离较小。这里 np.roll 是循环时移,会把尾部搬到开头;它只是对齐敏感性的演示,真实系统延迟应在共同有效区间评测。 谱平坦度描述所选频带的谱能量是否集中。纯音通常低、平坦噪声通常高,但音乐本来就含纯音,白噪声也可能是缺陷,因此不能用统一的 0.05 阈值判好坏。本实验干净信号的高频能量约为 $10^{-14}$ 量级,已几乎为空;此时平坦度主要由窗函数泄漏和数值 floor 决定,0.3256 不是「健康音频标准」。高频能量占比与平坦度在理想非零计算下都对全局增益不变,实际小能量频带则应先做绝对能量门控。 第三段:FAD 与样本量(fad_sample_size.py)。没有 VGGish,就用 16 维归一化对数梅尔幅度均值当特征——Fréchet 距离这一步的数学完全一样。核心实现采用对称半正定夹心矩阵。不能把一般非对称的 $\Sigma_r\Sigma_t$ 直接交给只处理对称矩阵的 eigh;通用矩阵平方根算法可以处理乘积,但不是这里的最小实现。附录检查非半正定输入,并仅把容差内的负舍入误差截为零。 真实运行输出(每行重新抽取样本,固定随机种子): 每堆样本数 同分布距离 带纯音距离 8 13.1390 47.5224 16 3.8499 49.1996 32 1.6452 46.6619 64 0.9992 46.3322 128 1.1044 45.4964 300 0.1844 45.4057 总体分布相同才有理论距离零,两个有限样本的估计值不必为零,也不必随样本量严格单调下降。这里是自定义 16 维梅尔特征的高斯距离演示,不是标准 VGGish FAD。 配图由 mel_spectrogram.py 重绘,单位为 dB,和表中自然对数单位相差固定倍数。B 的两条纯音均低于 Nyquist 频率;旧实现的 sin(2π·8000t) 在 16 kHz 采样时落在零点,不能用它模拟一个存在的 8 kHz 纯音。 图要看什么:B 在高频多出亮带;A 因独立归一化与参考重合;C 的谱形状变化较小。色谱只能证实频谱结构变化,不能代替主观听测。 05. 工业级实现对照 最小实现是为了讲清原理,生产里的实现多出的主要是「工程鲁棒性」: FAD:TorchEval 的 FrechetAudioDistance 的 FrechetAudioDistance(以 2026-09 的实现为准)把特征提取抽象成 preproc + model 两个可插拔组件,with_vggish() 一行构建官方配置;Microsoft 的 fadtk 则把 CLAP、MERT、HuBERT、Whisper 等十来种特征提取器都接了进来——FAD 的数值和特征提取器强绑定,换提取器等于换了一把尺子,不同论文的 FAD 数字不能直接比。另外它用流式更新(update/compute 两段式),避免一次把全部音频载入内存。 编解码器:facebookresearch/encodec/encodec/model.py 的 EncodecModel 提供 encode/decode 与可调目标码率的残差量化器,评测 codec 时要用它官方的 24 kHz / 48 kHz 配置,因为不同采样率下量化器的带宽分配不同,失真谱型也不同。 PESQ / STOI:两者来源不同。PESQ 对应 ITU-T P.862 系列,P.862.2 是宽带扩展;ITU 已于 2024-01-05 撤回该系列,并指向 P.863 系列。历史 benchmark 仍可报告 PESQ,但要写清版本、带宽和采样率。STOI 则来自 Taal 等人的可懂度研究,不是 P.862 标准;作者资料与代码 给出来源。常用实现为 pesq / pystoi,使用时仍须核对模式与预处理。 MOS 预测:UTMOS/DNSMOS 可以对单条样本输出预测分,适合筛选和回归测试。域内和域外的绝对值、排序都应在目标数据上验证,不能保证域内排序必然正确。人工听测也并非不可复现,可按 ITU-T P.800 等协议固定设计并报告评分分布和置信区间。 06. 代价与边界 指标 适合的问题 什么时候会骗你 SI-SDR 时间严格对齐的重建任务 对时移和相位差敏感;对单个全局增益不敏感 SNR 同上,且关心音量 对感知无感的相位问题同样过敏 对数梅尔距离 音色/频谱保真度 归一化会消除整体增益;逐帧比较仍对时序变化敏感,时间聚合又会丢失结构 PESQ / STOI 语音质量 / 语音可懂度 需按版本匹配带宽;不能直接推广为音乐或音效质量指标 FAD 生成系统的分布级评价 样本量不足时严重虚高;绑定特征提取器;对单条样本没有意义 MOS(人工) 指定协议下的主观评价 有采样不确定性;跨批次比较需要协议、样本与评分锚点控制 UTMOS / DNSMOS MOS 的廉价代理 排序与绝对值都需校准,域外泛化尤其容易失败 两条总原则: 任何单一指标都可能给出与听感相反的结论。开头的教学场景说明,波形误差变小并不排除感知缺陷。多指标交叉 + 小规模人工抽听,是底线配置。 指标回答「像不像」,不回答「好不好」。一个模型可以无限逼近参考(各种距离都趋近 0),同时毫无创造力;生成类任务的最终评价一定要留人工的位置。 (具体缺陷怎么查、该看哪个数字,见下面「诊断手册:按症状选指标」。) 07. 诊断手册:按症状选指标与读图 前面几节讲的是「每个指标测什么」,这一节反过来——从你听到或收到的问题出发,倒推该看哪个数字。(它和上面的边界表是一体两面:表说「指标会骗你」,这里说「那该看谁」。) 7.1 先区分配对参考、参考分布与无参考预测 有对应参考波形时,可算 SI-SDR、SNR、谱距离以及适用的语音指标。没有逐条参考时,MOS 预测器仍可逐条打分,文本条件模型也可以计算文本与音频相似度。经典 FAD 是集合统计,不直接定位单条缺陷;研究中的单样本扩展有另外的定义与边界。 比较两个生成器时,优先固定 prompt 分布、后处理、时长和测试规模。若确实有一一对应的 prompt 与样本,可以做配对分析;非配对设计同样可以成立,但需要匹配设计的统计方法。报告效应量和置信区间,必要时使用置换检验或按独立音频单位重采样;多指标检验还要考虑多重比较。p=0.02 不能因为只有 24 条样本就改称「勉强」,p=0.2 也不能证明没有差异。单个集合级 FAD 更不能直接塞进逐条配对 t 检验。 7.2 按症状安排检查 持续高音、电流声:先看谱图确认是否存在额外窄带峰,再看该频带能量与逐帧平坦度。嘶声可能是宽带噪声,平坦度反而较高;蜂鸣、鸟叫、乐器泛音也可能是正常窄带成分。阈值必须由正常/异常验证集校准。 发闷或带宽不足:检查绝对电平、高频能量、频谱质心和相对参考的带宽。贝斯、鼓等低频内容天然窄带,不能由质心低直接判缺陷。 金属声或音乐噪声:沿时间看瞬态谱峰、带宽和谱距离变化,避免全段平均掩盖短时问题。有参考时可结合分离任务的伪影指标,再人工确认。 相位或通道异常:明确是否允许补偿全局延迟,检查通道间相关和对齐后的波形误差。低 SI-SDR 加小谱距离只能提示对齐问题,不能唯一证明根因。 整体增益偏差:用 RMS、LUFS 或参考增益比。SI-SDR 对单个全局非零增益不变,对随时间变化的增益并不不变;响度泵动需要逐窗观察。 可懂度不足:语音任务可结合 STOI 与人工转写;WER/CER 还依赖 ASR 本身,不能把识别器错误全部算成生成缺陷。 内容跑偏或多样性不足:在一致协议下看 FAD、条件一致性和多样性统计。特征方差变大也可能只是噪声,不能以「方差越大越好」替代覆盖度分析。 7.3 排查顺序和判定边界 先检查格式、采样率、声道、时长、绝对电平,再检查相对谱、局部异常和整体统计。接近静音时,按自身峰值归一化会把噪声放大得像正常内容,平坦度也可能很高;此时应单独报告静音比例及条件质量,不能默默删掉静音样本来提高总分。 浮点信号超过 0 dBFS 说明超出常见满刻度约定,但尚不能证明已发生硬削波。真峰值与采样峰值也要区分,是否压限应依据目标交付规范;不能假设流媒体平台的响度归一化会自动修复过载。THD+N 通常需要规定的测试信号,不是对任意音乐直接测 50/60 Hz 哼声的万能工具,后者应在谱上检查基频、谐波与背景的相对强度。 最后把可疑样本与 prompt、参考和听测对照:可视化说明信号发生了什么,不能单独决定这种变化是否符合内容要求。附录实验能验证增益、时移和新增纯音如何影响指标;实际产品的合格线需要自己的数据与听测来确定。 08. 经典论文脉络 PESQ(ITU-T P.862,2001):通信时代的遗产。把心理声学掩蔽模型塞进了客观指标,第一次让「机器预测听感」在窄带语音上有了可信度——也把「只适用语音」的边界刻进了骨子里。 SDR – half-baked or well done? (arXiv:1811.02508)(ICASSP 2019):指出 BSS_eval 那版 SDR 在单通道场景下会被滥用(改参考信号、靠调音量刷分),给出本文用的 SI-SDR 定义——误差与参考正交、不依赖幅度缩放。读它是理解「为什么偏偏叫 SDR」的最短路径。 Fréchet Audio Distance: A Metric for Evaluating Music Enhancement Algorithms (arXiv:1812.08466):把 FID 的思想搬到音频,提出基于背景集特征统计的音乐质量指标,解决了生成式音乐「没有正确答案可比」的问题;留下的问题是 FAD 数值绑定特征提取器、且强依赖样本量。 High Fidelity Neural Audio Compression (arXiv:2210.13438)(EnCodec):神经音频 codec 的代表作,结合客观指标与主观听测评价低码率重建。不能由这篇论文推导出「SI-SDR + PESQ + 梅尔距离」是所有 codec 的固定验收标准。 UTMOS: UTokyo-SaruLab System for VoiceMOS Challenge 2022 (arXiv:2204.02152):用 SSL 特征 + 集成学习预测 MOS,拿了 VoiceMOS Challenge 2022 多项第一。它代表的方向是「用模型替代昂贵的人工听测」——可行,但域外泛化仍是软肋。 一条暗线贯穿始终:指标的演进一直在追赶生成能力的演进。信号指标管不了感知,感知指标管不了分布,分布指标管不了单条样本——每一代生成技术把旧指标逼失效一次,然后逼出下一代指标。 09. 常见误解 「SI-SDR 高就是音质好」。它衡量去掉全局增益后的波形误差,对齐、相位、噪声和伪影都会影响它,但影响权重并非人耳感知权重。 「FAD 低说明每条音频都好」。均值与协方差相近不保证所有样本合格,也不保证更高阶分布一致、文本匹配或内容多样性。 「FAD 数字可以跨论文直接比」。特征器、参考集、时长、分帧、样本量、预处理一致,比较才有明确含义。本例是自定义梅尔特征距离,不应混称标准分数。 「MOS 不可复现」。MOS 是按协议收集的主观统计量,可以复做实验估计误差;跨批次须控制听测设计,报告不确定性。VoiceMOS 也评价 MSE 等指标,并非只比相关系数而完全不看绝对误差。 「无参考模型可替代人工验收」。MOS 预测器有域和数据偏差,即使在训练域也可能错排;发布标准需要经过人工标注或听测校准。 10. 动手验证 先安装 numpy matplotlib,把文末三个脚本存到同一目录,分别运行。它们不下载音频或权重;可复现的结论来自数值输出,不包括主观听感结论。 把 metric_disagreement.py 的 7000/7500 Hz 改成 200/250 Hz,重跑。高频占比应基本不变,但低频梅尔谱会改变。记录实际距离,不能预设「一定更小」或「一定同样刺耳」。这说明频带指标有明确的观测范围。 把 np.roll(x, 8) 改成 np.roll(x, 1),比较 SI-SDR 与谱距离。变化取决于波形自相关,不存在通用于所有音频的固定分数;真实延迟应在共同有效区间对齐,避免循环边界干扰。 在 FAD 脚本的循环中加入样本量 4,并多换几个种子。16 维协方差的秩最多为 3,估计很不稳定,但仍可计算;单次结果也可能偶然小。观察分布而非某一次抽样,才是样本量实验的正确读法。 11. 延伸阅读 FID / CLIP Score 到底测了什么(image_metrics):FAD 的直系前身。FID 的样本量陷阱、特征提取器绑定问题在图像侧一模一样,先读它再看本文会有「同一个故事换了个领域」的感觉。 视频生成评测:VBench 与人工验收(video_metrics):评测的维度从音频扩展到视频,多维度分解 + 人工验收的思路一脉相承。 离散化表征:VQ-VAE 与 VQGAN(vq_tokenizer):音频 codec(EnCodec、SoundStream)就是 VQ-VAE 思想在波形上的应用,想理解被评测的对象,从这篇进。 官方工具:fadtk(多种特征提取器的 FAD 工具箱)、torcheval 的 FrechetAudioDistance、ITU-T P.862.2(历史宽带 PESQ 与撤回说明)。 附录:完整代码 09 节用到的脚本全文如下(metric_disagreement.py、fad_sample_size.py、mel_spectrogram.py)。复制到本地存成同名文件,按各脚本开头的依赖说明准备环境后即可运行。 metric_disagreement.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """同一个音频退化,四个指标给出四种排名。 这篇文章的核心论点就是这张表:指标不是「音质分数」,每个指标只回答一个 很具体的问题。用错了就会得到自欺欺人的结论。 三种退化分别是整体增益、窄带纯音和循环时移。 表中听感仅是待验证的假设,本文没有执行受试者听测。 干净的语音/音乐可以看成「基频 + 一串谐波」,这里就按这个造, 省得读者还要去找音频文件。只用 numpy。 """ import numpy as np SR = 16000 DUR = 1.0 N = int(SR * DUR) rng = np.random.default_rng(0) def make_clean(n=N, sr=SR): """基频 220Hz + 前 6 个谐波,加一点慢包络,听起来像有起伏的乐音。""" t = np.arange(n) / sr x = np.zeros(n) for k in range(1, 7): x += (1.0 / k) * np.sin(2 * np.pi * 220 * k * t) # 慢包络:2.5Hz 的轻微起伏,避免信号看起来像死板的稳态正弦 x *= 0.75 + 0.25 * np.sin(2 * np.pi * 2.5 * t) return x / np.abs(x).max() def degrade(x, kind, sr=SR): if kind == "A": return 0.95 * x # 音量差 5% if kind == "B": # 7/7.5 kHz 的窄带纯音:是否刺耳需听测,且两者均低于 Nyquist 频率 t = np.arange(len(x)) / sr buzz = 0.02 * ( np.sin(2 * np.pi * 7000 * t) + np.sin(2 * np.pi * 7500 * t) ) return x + buzz if kind == "C": return np.roll(x, 8) # 时移 8 采样 = 0.5ms raise ValueError(kind) def si_sdr(ref, est): """Scale-Invariant SDR。 先把 est 投影到 ref 上(找一个最优缩放 alpha),剩下的残差才算失真。 所以「整体音量变了」在 SI-SDR 里被完全抵消——这是设计意图, 但也意味着它测不出增益错误。 """ if np.dot(ref, ref) <= 0 or np.dot(est, est) <= 0: raise ValueError("SI-SDR needs non-silent reference and estimate") alpha = np.dot(est, ref) / np.dot(ref, ref) e_target = alpha * ref e_res = est - e_target num = np.dot(e_target, e_target) den = np.dot(e_res, e_res) if den <= 1e-24 * num: return float("inf"), alpha return 10 * np.log10(num / den), alpha def snr(ref, est): """不做缩放对齐的信噪比。音量差会被算成失真。""" res = est - ref den = np.dot(res, res) if den == 0: return float("inf") return 10 * np.log10(np.dot(ref, ref) / den) def stft_mag(x, n_fft=512, hop=128): win = np.hanning(n_fft) n_frames = 1 + (len(x) - n_fft) // hop frames = np.stack([x[i * hop:i * hop + n_fft] * win for i in range(n_frames)]) return np.abs(np.fft.rfft(frames, n_fft, axis=1)).T # [freq, time] def hz_to_mel(f): return 2595.0 * np.log10(1.0 + f / 700.0) def mel_to_hz(m): return 700.0 * (10 ** (m / 2595.0) - 1.0) def mel_filterbank(sr=SR, n_fft=512, n_mels=64, fmin=20.0, fmax=8000.0): n_freqs = n_fft // 2 + 1 hz = mel_to_hz(np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_mels + 2)) bins = np.floor(hz * n_fft / sr).astype(int) fb = np.zeros((n_mels, n_freqs)) for i in range(n_mels): left, center, right = bins[i], bins[i + 1], bins[i + 2] if center <= left: center = left + 1 if right <= center: right = center + 1 for k in range(left, min(center, n_freqs)): fb[i, k] = (k - left) / (center - left) for k in range(center, min(right, n_freqs)): fb[i, k] = (right - k) / (right - center) return fb def log_mel(x, sr=SR, n_fft=512, hop=128, n_mels=64): """归一化梅尔幅度谱的自然对数,单位为 ln 幅度(不是 dB)。 每条样本独立归一化是本实验为了忽略整体增益所作的选择, 不是所有音频评测的必需步骤。幅度 floor=1e-4 对应峰值下 80 dB。 """ fb = mel_filterbank(sr=sr, n_fft=n_fft, n_mels=n_mels) mel = fb @ stft_mag(x, n_fft=n_fft, hop=hop) mel = mel / (mel.max() + 1e-12) return np.log(np.maximum(mel, 1e-4)) # floor 在 log(1e-4) ≈ -9.2 def log_mel_distance(ref, est): """对数梅尔谱的平均绝对误差(单位为自然对数幅度差)。""" return float(np.mean(np.abs(log_mel(est) - log_mel(ref)))) def hf_energy_ratio(x, sr=SR, n_fft=512, hop=128, f_lo=6000.0): """6kHz 以上能量占总能量的比例——高频毛刺会把它顶起来。""" mag = stft_mag(x, n_fft=n_fft, hop=hop) freqs = np.fft.rfftfreq(n_fft, 1.0 / sr) power = (mag ** 2).mean(axis=1) return float(power[freqs >= f_lo].sum() / power.sum()) def spectral_flatness(x, sr=SR, n_fft=512, hop=128, f_lo=4000.0, f_hi=9000.0): """高频平均功率谱的几何均值/算术均值,范围 0~1。 小值表示谱能量集中,不等价于存在缺陷;纯乐音也会很低。 只在该频带有足够能量时解释,阈值必须按任务和频带校准。 与高频能量占比一样,理想计算对非零全局增益不变。 """ mag = stft_mag(x, n_fft=n_fft, hop=hop) freqs = np.fft.rfftfreq(n_fft, 1.0 / sr) band = mag[(freqs >= f_lo) & (freqs <= f_hi)] power = (band ** 2).mean(axis=1) + 1e-20 geo = np.exp(np.mean(np.log(power))) ari = np.mean(power) return float(geo / ari) def main(): ref = make_clean() print(f"参考信号: {len(ref)} 采样 = {len(ref) / SR:.2f}s @ {SR}Hz") print(f"参考信号高频能量占比: {hf_energy_ratio(ref):.3e}") print() header = (f"{'退化':<18}{'操作':<12}{'SI-SDR(dB)':>12}{'SNR(dB)':>10}" f"{'梅尔距离':>10}{'高频占比':>12}{'高频谱平坦度':>14}") print(header) print("-" * len(header)) rows = {} for kind, feel in (("A", "增益变化"), ("B", "新增纯音"), ("C", "循环时移")): est = degrade(ref, kind) s, alpha = si_sdr(ref, est) d = { "si_sdr": s, "snr": snr(ref, est), "mel": log_mel_distance(ref, est), "hf": hf_energy_ratio(est), "flat": spectral_flatness(est), "alpha": alpha, } rows[kind] = d print(f"{kind + ' ' + {'A': '音量×0.95', 'B': '高频毛刺', 'C': '时移0.5ms'}[kind]:<18}" f"{feel:<12}{s:>12.2f}{d['snr']:>10.2f}{d['mel']:>10.4f}" f"{d['hf']:>12.3e}{d['flat']:>14.4f}") print() print("读数(这才是重点):") print(f" A 音量小 5%:SI-SDR = {rows['A']['si_sdr']:.2f} dB(缩放被抵消,几乎满分)," f"SNR = {rows['A']['snr']:.2f} dB(老老实实记了一笔)") print(f" → 最优缩放 alpha = {rows['A']['alpha']:.4f},表示估计波形为参考的 0.95 倍") print(f" B 高频毛刺:SI-SDR = {rows['B']['si_sdr']:.2f} dB 看着还行," f"高频占比却从 {hf_energy_ratio(ref):.3e} 涨到 {rows['B']['hf']:.3e}") print(f" C 时移 0.5ms:SI-SDR 掉到 {rows['C']['si_sdr']:.2f} dB," f"但梅尔距离只有 {rows['C']['mel']:.4f} —— 谱几乎没动") print() print("结论:SI-SDR 衡量的是「波形对不对齐」,梅尔距离衡量的是「频谱像不像」,") print(" 高频占比描述频带能量分配,是否刺耳还需要听测。") if __name__ == "__main__": main() fad_sample_size.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """FAD 到底在比什么,以及为什么样本量不够时它根本不可信。 FAD(Fréchet Audio Distance)是 FID 的音频版:不问「这两条音频像不像」, 而是问「这两堆音频的分布像不像」。这个区别决定了它有两个反直觉的性质: 1. 它不配对。不需要参考音频和测试音频一一对应, 所以能评测「凭空生成」的音乐/音效——这是 PESQ 之类做不到的。 2. 它对样本量极其敏感。协方差是从样本里估的,样本不够时估出来的协方差 本身就是歪的,算出来的 FAD 跟真实值差得离谱。 这里用 16 维归一化对数梅尔幅度均值当特征(真正的 FAD 用 VGGish 的 128 维 embedding, 我们没有那个模型,但 Fréchet 距离这一步的数学完全一样)。 """ import numpy as np SR = 16000 N_FFT = 512 HOP = 128 N_MELS = 16 rng = np.random.default_rng(7) def stft_mag(x, n_fft=N_FFT, hop=HOP): win = np.hanning(n_fft) n_frames = 1 + (len(x) - n_fft) // hop frames = np.stack([x[i * hop:i * hop + n_fft] * win for i in range(n_frames)]) return np.abs(np.fft.rfft(frames, n_fft, axis=1)).T def hz_to_mel(f): return 2595.0 * np.log10(1.0 + f / 700.0) def mel_to_hz(m): return 700.0 * (10 ** (m / 2595.0) - 1.0) def mel_filterbank(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=8000.0): n_freqs = n_fft // 2 + 1 hz = mel_to_hz(np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_mels + 2)) bins = np.floor(hz * n_fft / sr).astype(int) fb = np.zeros((n_mels, n_freqs)) for i in range(n_mels): left, center, right = bins[i], bins[i + 1], bins[i + 2] if center <= left: center = left + 1 if right <= center: right = center + 1 for k in range(left, min(center, n_freqs)): fb[i, k] = (k - left) / (center - left) for k in range(center, min(right, n_freqs)): fb[i, k] = (right - k) / (right - center) return fb FB = mel_filterbank() def embed(x): """一段音频 → 16 维特征:各梅尔带的平均对数幅度。""" mel = FB @ stft_mag(x) mel = mel / (mel.max() + 1e-12) return np.log(np.maximum(mel, 1e-4)).mean(axis=1) def make_tone(f0, n=8000, buzz=False, snr_db=None): """一个乐音:基频 f0 + 谐波。buzz=True 时叠高频毛刺。""" t = np.arange(n) / SR x = np.zeros(n) for k in range(1, 6): x += (1.0 / k) * np.sin(2 * np.pi * f0 * k * t) if buzz: x += 0.05 * (np.sin(2 * np.pi * 7000 * t) + np.sin(2 * np.pi * 7500 * t)) if snr_db is not None: x = x + rng.normal(0, np.sqrt(np.mean(x ** 2) / (10 ** (snr_db / 10))), n) return x / (np.abs(x).max() + 1e-12) def sqrtm_psd(A): """对称半正定阵的平方根;仅裁掉舍入级负特征值。""" A = (A + A.T) / 2 w, V = np.linalg.eigh(A) if w.min() < -1e-10 * max(1.0, np.abs(w).max()): raise ValueError("matrix is not positive semidefinite") w = np.clip(w, 0.0, None) return (V * np.sqrt(w)) @ V.T def frechet_distance(X, Y): """两个特征集合的 Fréchet 距离。 把每堆特征当成一个高斯(只关心均值和协方差),再算这两个高斯之间的 2-Wasserstein 距离的平方。闭式解: FD = ||mu_x - mu_y||^2 + Tr(S_x + S_y - 2 (S_x S_y)^{1/2}) """ mu_x, mu_y = X.mean(axis=0), Y.mean(axis=0) S_x = np.cov(X, rowvar=False) S_y = np.cov(Y, rowvar=False) diff = mu_x - mu_y # 不可把一般非对称的 S_x @ S_y 直接交给 eigh。 # 使用对称半正定夹心矩阵,平方根的迹与原式相同。 s_x_sqrt = sqrtm_psd(S_x) covmean = sqrtm_psd(s_x_sqrt @ S_y @ s_x_sqrt) fd = diff @ diff + np.trace(S_x) + np.trace(S_y) - 2 * np.trace(covmean) tol = 1e-7 * max(1.0, np.trace(S_x) + np.trace(S_y)) if fd < -tol: raise ArithmeticError("negative distance exceeds numerical tolerance") return float(max(fd, 0.0)) # 仅容差内的负舍入误差截到 0 def main(): # 参考集:干净乐音,基频在 200~400Hz 之间随机 # 测试集 A:同分布(只是另一批采样)→ FAD 应该接近 0 # 测试集 B:带高频毛刺 → FAD 应该明显大于 0 def make_set(n, buzz=False): return np.stack([ embed(make_tone(rng.uniform(200, 400), buzz=buzz)) for _ in range(n) ]) print("特征维度: 16(各梅尔带的平均对数幅度)") print() print("样本量对 FAD 的影响(同一个分布对,只是抽的样本数不同):") print(f"{'每堆样本数':>10}{'同分布 FAD':>14}{'带毛刺 FAD':>14}{'相对基线':>10}") print("-" * 50) big_ref = make_set(600) for n in (8, 16, 32, 64, 128, 300): ref = big_ref[:n] same = np.stack([ embed(make_tone(rng.uniform(200, 400))) for _ in range(n) ]) buzz = make_set(n, buzz=True) fd_same = frechet_distance(ref, same) fd_buzz = frechet_distance(ref, buzz) ok = f"{fd_buzz / max(fd_same, 1e-6):.1f}x" print(f"{n:>10}{fd_same:>14.4f}{fd_buzz:>14.4f}{ok:>10}") print() print("注意最上面几行:样本只有 8 或 16 个时,「同分布」的 FAD 就已经不是 0 了,") print("这是均值与协方差估计误差;16 维、16 个样本的中心化协方差必然奇异。") print("真实 VGGish 特征为 128 维,但所需样本数还取决于谱结构、独立性和精度——") print("这就是为什么论文里报的 FAD 必须同时报样本数,否则数字不可比。") print() print("另一面:") print(f" 300 个样本时,同分布 FAD 与带毛刺 FAD 差了 " f"{frechet_distance(big_ref[:300], make_set(300, buzz=True)) / max(frechet_distance(big_ref[:300], np.stack([embed(make_tone(rng.uniform(200, 400))) for _ in range(300)])), 1e-9):.1f} 倍") print(" → 倍数只是描述量,不是统计检验;需多次抽样或按音频做 bootstrap。") if __name__ == "__main__": main() mel_spectrogram.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """把「指标打架」画出来看:三种退化的梅尔谱。 波形指标(SI-SDR)分不清的三件事,在谱上一眼就能看出来: 音量 ×0.95 —— 本图每条单独归一化,整体增益被抵消 高频毛刺 —— 谱的中上部多出一条亮带(听感需另行验证) 时移 0.5ms —— 几乎看不出差别(但 SI-SDR 可能明显下降) 产出 figures/mel_compare.png,文章里用  引用。 """ import os import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib import font_manager # matplotlib 默认字体没有中文,标题会渲染成一排豆腐块。 # 按 macOS / Linux 常见字体逐个试,找一个系统里真有的。 _AVAILABLE = {f.name for f in font_manager.fontManager.ttflist} for _cand in ["PingFang SC", "Heiti SC", "Arial Unicode MS", "Songti SC", "STHeiti", "Noto Sans CJK SC", "WenQuanYi Zen Hei"]: if _cand in _AVAILABLE: plt.rcParams["font.sans-serif"] = [_cand] break plt.rcParams["axes.unicode_minus"] = False # 负号不要用中文缺字替代 SR = 16000 N_FFT = 512 HOP = 128 N_MELS = 64 OUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "figures") def make_clean(n=16000, sr=SR): t = np.arange(n) / sr x = np.zeros(n) for k in range(1, 7): x += (1.0 / k) * np.sin(2 * np.pi * 220 * k * t) x *= 0.75 + 0.25 * np.sin(2 * np.pi * 2.5 * t) return x / np.abs(x).max() def degrade(x, kind, sr=SR): t = np.arange(len(x)) / sr if kind == "A": return 0.95 * x if kind == "B": return x + 0.02 * (np.sin(2 * np.pi * 7000 * t) + np.sin(2 * np.pi * 7500 * t)) if kind == "C": return np.roll(x, 8) raise ValueError(kind) def stft_mag(x, n_fft=N_FFT, hop=HOP): win = np.hanning(n_fft) n_frames = 1 + (len(x) - n_fft) // hop frames = np.stack([x[i * hop:i * hop + n_fft] * win for i in range(n_frames)]) return np.abs(np.fft.rfft(frames, n_fft, axis=1)).T def hz_to_mel(f): return 2595.0 * np.log10(1.0 + f / 700.0) def mel_to_hz(m): return 700.0 * (10 ** (m / 2595.0) - 1.0) def mel_filterbank(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=8000.0): n_freqs = n_fft // 2 + 1 hz = mel_to_hz(np.linspace(hz_to_mel(fmin), hz_to_mel(fmax), n_mels + 2)) bins = np.floor(hz * n_fft / sr).astype(int) fb = np.zeros((n_mels, n_freqs)) for i in range(n_mels): left, center, right = bins[i], bins[i + 1], bins[i + 2] if center <= left: center = left + 1 if right <= center: right = center + 1 for k in range(left, min(center, n_freqs)): fb[i, k] = (k - left) / (center - left) for k in range(center, min(right, n_freqs)): fb[i, k] = (right - k) / (right - center) return fb FB = mel_filterbank() def log_mel_db(x): mel = FB @ stft_mag(x) mel = mel / (mel.max() + 1e-12) return 20 * np.log10(np.maximum(mel, 1e-4)) # 归一化后取 dB,floor 在 -80dB def main(): ref = make_clean() cases = [ ("参考(干净)", ref), ("A 音量 ×0.95", degrade(ref, "A")), ("B 高频毛刺", degrade(ref, "B")), ("C 时移 0.5ms", degrade(ref, "C")), ] specs = [log_mel_db(x) for _, x in cases] # 每条谱已独立按峰值归一化,共用色标只用于比较相对谱形状。 vmin = min(s.min() for s in specs) vmax = max(s.max() for s in specs) fig, axes = plt.subplots(1, 4, figsize=(15, 3.6), constrained_layout=True) for ax, (name, _), s in zip(axes, cases, specs): im = ax.imshow(s, origin="lower", aspect="auto", vmin=vmin, vmax=vmax, cmap="magma") ax.set_title(name, fontsize=12) ax.set_xlabel("帧") ax.set_ylabel("梅尔带") fig.colorbar(im, ax=axes, shrink=0.8, label="dB(相对峰值)") fig.suptitle("同一个音频的四种形态:梅尔谱", fontsize=13) os.makedirs(OUT_DIR, exist_ok=True) out = os.path.join(OUT_DIR, "mel_compare.png") fig.savefig(out, dpi=150, bbox_inches="tight", facecolor="white") plt.close(fig) print(f"图已保存: {out}") print() print("每张谱的动态范围与「和参考谱的平均绝对差」:") print(f"{'形态':<16}{'最小值(dB)':>12}{'最大值(dB)':>12}{'与参考的平均差(dB)':>20}") print("-" * 62) base = specs[0] for (name, _), s in zip(cases, specs): print(f"{name:<16}{s.min():>12.2f}{s.max():>12.2f}" f"{np.mean(np.abs(s - base)):>20.3f}") if __name__ == "__main__": main()
2026年09月21日
17 阅读
0 评论
2 点赞
2026-09-21
AIGC 每日速读|2026-09-21|伯克利线性注意力让H3视频快14.5倍-VDN
今日 AIGC 论文速览 今日共 10 篇 重点论文标题列表 Video DeltaNet(UC Berkeley × Impossible Inc. × UT Austin):线性注意力让H3快14.5倍 dQwen3.5(University of Texas at Austin):混合骨干省一半适配token Edit-VAR(中山大学 × 华南理工大学 × 清华大学 × 山东大学 × 南开大学 × 湖南大学):免训练免反演改视频 Paint-Anything(ByteDance Seed × 浙江大学 × 南京大学):写个十六进制就能上色 StepAudio 3 Music(StepFun(阶跃星辰)× ACE × 香港中文大学 × UC San Diego):先写ABC谱再生成5分半歌 今日论文速览 1. Video DeltaNet:线性注意力让H3快14.5倍 Video DeltaNet: A Video-Native Hybrid Attention for Livestream Video Generation | UC Berkeley × Impossible Inc. × UT Austin | arXiv:2609.20744 关键词:视频扩散模型,线性注意力,混合注意力,推理加速 前序问题:视频扩散模型每一步去噪都要把又长又宽的时空 token 序列过一遍注意力,算力全压在这里。线性注意力在大语言模型里已经被验证好用,但直接搬到视频模型上会丢掉细粒度交互,画面质量跟着掉——于是「省算力」和「保质量」一直是二选一。 本文贡献:论文提出 Video DeltaNet(VDN),不做纯替换而是双分支并存:局部 Softmax 注意力管近处的细粒度交互,双向线性记忆管长程视频上下文。线性分支的核心是 Video Delta Attention(VDA),它不按 token 逐个更新记忆,而是以「帧」为单位、把该帧的全部空间 token 一起吸收进去,一帧只更新一次。两个分支各有独立的输出投影和可学习门控来校准配比;适配则走分阶段的教师对齐配方(Stage A1 先单独对齐每个线性分支,A2 再端到端训练组装后的混合层),把新通路渐进地塞进已预训练的模型里。作者把 VDN 落在 MiniMax H3 上,只对 video-to-video 这类交互用混合注意力,凡是牵涉文本或音频的交互仍然保留 Softmax。 实验效果:配合八步蒸馏和优化过的 SGLang 服务栈,VDN-H3 在 8 张 NVIDIA B200 上生成 14.3 秒 768p 视频,DiT 去噪只用 6.70 秒,相对同样 8 卡、50 步的 dense H3 基线提速 14.5 倍。质量方面:50 步时 Stage B 后的 VDN-H3 与蒸馏前的 Dense H3 基本持平(美学、技术、学习质量、FIRM 各项接近,仅两个端点保真度指标略降);八步时在五项无参考质量指标上全部最高,VQA_T 比 FastH3 v2 高 10.78 分;四步时同样在五项上超过 Dense H3 加 Larry 适配器与 FastH3 v1,VQA_T 领先 FastH3 v1 达 12.64 分。 批判点评:这篇的价值在于把「线性注意力用于视频」从整体替换改成了按交互类型分流:视频与视频之间用线性,涉及文本、音频的跨模态交互留给 Softmax,承认了不同交互对精度的敏感度本就不一样,比一刀切务实得多。以帧为单位更新记忆也和视频 tokenizer 的时间结构对齐(H3 的 VAE 按五个连续潜帧为一个时间块解码,窗口顺着它走以免 Softmax 边界切断 tokenizer 的自然时间单元)。要留意的是 14.5 倍里混进了八步蒸馏和服务栈优化的贡献,架构本身的净收益要看 50 步那组对照;论文也明确说八步那组的显著性标记是各变体与 Dense 50 比,并非 VDN-H3 直接对 FastH3。另外全部结论都长在 MiniMax H3 这一个骨干上,换模型是否还成立尚无证据。 2. dQwen3.5:混合骨干省一半适配token dQwen3.5: Hybrid-Attention Diffusion Language Models | University of Texas at Austin | arXiv:2609.20751 关键词:扩散语言模型,混合注意力,模型适配,并行解码 前序问题:把预训练好的自回归模型改造成扩散语言模型(DLM)是条省钱路线,但几乎所有现成改造都从全注意力 Transformer 出发。问题是自回归这边早就转向了注意力与 RNN 层交错的混合架构,而 RNN 结构上就是因果的,要把它双向化并不容易——骨干和目标范式对不上。 本文贡献:作者索性直接试:以 Qwen3.5 为起点,在 0.8B、2B、4B、9B 四个规模上做自回归到扩散的适配,得到 dQwen3.5 系列,并拿一个全注意力模型做对照,专门回答两个问题——混合骨干是否是高效的适配起点,适配出来的模型还保不保留定义 DLM 的并行与任意序解码行为。评测统一用同一套解码方案(默认画布 1024)而非各家自带的推理配方,并用 trunk(去掉 embedding 与输出头)参数量作为容量口径对齐比较组。 实验效果:混合骨干确实更省:达到同一训练 loss 所需的 token 量大约只要全注意力对照的一半,配图上逐个 loss 档位统计的中位数是 2.21 倍 token 效率差。跨规模看,dQwen3.5 在任意序解码行为上与全注意力 DLM 相似,并行解码下表现强劲——HumanEval 与 MBPP 上 1× 到 8× 加速区间基本压住全注意力对照。 批判点评:结论有用且反直觉:结构上「不适合」双向化的混合骨干,反而是更划算的 DLM 起点,这对想复用现成混合架构权重的人是直接的好消息。但配图暴露了摘要没提的代价——在 MATH500 上,100B token 档的混合版明显落后全注意力对照(1× 时约 9.3% 对 15.8%),数学推理这一块的损失不小,只有在 4× 以上高加速区间才靠后者衰减更快而追平。作者自己在注释里也怀疑退化可能来自训练数据混合不如 Qwen3/Qwen3.5 原始配方,这等于承认「骨干差异」和「数据差异」还没被干净地分离。另外全部比较都在基座 checkpoint 上做,后训练被明确排除在外,而后训练对下游表现影响很大,所以这套结论离「能直接上线」还有距离。 3. Edit-VAR:免训练免反演改视频 Edit-VAR: Taming Visual Autoregressive Model for Precise Video Editing | 中山大学 × 华南理工大学 × 清华大学 × 山东大学 × 南开大学 × 湖南大学 | arXiv:2609.21268 关键词:视频编辑,视觉自回归,免训练,token替换 前序问题:文本引导的视频编辑要改对目标内容,同时保住没被编辑区域的外观和时间连贯。训练型方法控制力强但吃数据吃算力;免训练方法分两条路,免反演的不用恢复轨迹,可它那套保源引导会限制编辑强度、让语义改动做不彻底;反演型先恢复潜轨迹再重生成,近似误差会累积,导致源内容漂移和时间不一致。 本文贡献:Edit-VAR 是第一个基于预训练视觉自回归视频模型、既免训练又免反演的文本引导视频编辑框架。它把源视频直接编码成多尺度离散 token,用概率引导的条件 token 替换来保源;再用注意力引导的 token 级与尺度感知调制,只在与编辑相关的位置和生成阶段选择性放松源约束。Scale-Decoupled Generation 以「晚期尺度解除约束」的形式实现,负责重新生成运动一致的细节、减少纹理碎裂。另有残差引导的 token 剪枝,利用最后两个高分辨率尺度上的冗余来压推理开销。 实验效果:大量实验加一项盲测用户研究显示,Edit-VAR 在编辑保真度、源内容保持、时间连贯性和推理效率上整体优于现有免训练视频编辑方法。定性对比里,同样把「red boat hull」改成「dark blue boat hull」、把「tree frog」改成「poison dart frog」,VACE 会连船型和背景一起改,RAVE 整片色调偏移、青蛙变成纯绿,Edit-VAR 则只换掉目标物的颜色与纹理,港口的其他船、水面倒影和叶片细节都留住了。 批判点评:把编辑搬到离散多尺度 token 上做,绕开了连续扩散反演的误差累积,这个切换是这篇最实在的地方——问题定位在「反演」而不是泛泛的「一致性」,配套的概率引导替换也把「保源还是接受编辑」变成了可按 token 判定的显式比较。但摘要通篇只给「整体优于」的定性结论,没有一个具体数字,编辑保真与源保持之间的权衡到底移动了多少无从判断;残差剪枝说是降推理成本,省了多少、掉不掉质量也没披露。更根本的限制是它绑在视觉自回归视频模型上,而这类骨干的开源生态和画质上限目前都不如主流扩散视频模型,方法再好也受骨干天花板约束。 4. Paint-Anything:写个十六进制就能上色 Paint-Anything: Unified Any-Color Control for Image Generation and Editing | ByteDance Seed × 浙江大学 × 南京大学 | arXiv:2609.20816 关键词:图像生成,图像编辑,颜色控制,数据管线 前序问题:专业设计需要的是任意颜色控制:用任一 24 位十六进制值指定某个物体的目标颜色,生成和编辑都得听话。以往工作在颜色生成、编辑、上色上各做一块,却往往依赖专门的颜色表示或特制的推理流程,不通用。而大语言模型这边提供了一个更简单的起点——连小模型都已经能把十六进制值和颜色语义对上。 本文贡献:Paint-Anything 直接把十六进制写进文本提示(包在 color 标签里)由文本编码器编码,通过物体级颜色监督学出一套生成与编辑共享的 hex-prompt 接口。配套数据管线从真实图像构建 Paint-500K:经 VLM 定位、SAM3 掩码、CIE 空间颜色提取、编辑对合成、过滤与重新描述而成。关键设计是:真实图像有阴影,标签只能算近似颜色,所以再补一批纯色锚点——它们的像素与配对的十六进制严格一致;而这些锚点只在高噪声时间步参与训练,低噪声阶段仍交给自然图像,避免把纯色的平坦质感带进成品。论文还提出 Any Color Benchmark(ACBench),含 ACBench-T2I 与 ACBench-Edit 两部分,专门量物体级十六进制颜色保真度。 实验效果:在 FLUX.2-4B 上,Paint-Anything 把 ACBench-T2I 和 ACBench-Edit 分数相对基座分别提升 85.3% 和 28.3%,并在参与比较的方法中取得最高的平均 CompColor 分。消融实验支持这套训练配方。定性对比里基座 FLUX.2-4B 常给出「语义上说得通但数值上离要求很远」的颜色——要 #FF5634 的车身给成偏红、要 #000080 的沙漠绿洲给成亮蓝,Paint-Anything 则能贴住指定色值。 批判点评:「纯色锚点只在高噪声时段用」这个细节是全文最值得抄的一笔:它承认了监督信号在不同噪声尺度上的可信度不同,用时间步做分工而不是简单混数据,比多数「加一路干净监督」的做法更讲究。把十六进制塞进文本提示、不引入专门颜色编码器,也让它能直接挂在现成模型上。但 85.3% 是相对基座的提升,基座本身在这项上分数低,相对增幅容易放大;编辑侧只有 28.3%,两者差距说明编辑任务里颜色约束和内容保持的冲突还没解决好。更要紧的是 ACBench 是作者自建的评测,「在自家基准上提升最多」这件事需要外部复现才能定性。另外全部结果都在 FLUX.2-4B 上,跨骨干的可迁移性未验证。 5. StepAudio 3 Music:先写ABC谱再生成5分半歌 StepAudio 3 Music Technical Report | StepFun(阶跃星辰)× ACE × 香港中文大学 × UC San Diego | arXiv:2609.16034 关键词:音乐生成,音频tokenizer,MoE,流匹配DiT 前序问题:长篇音乐生成要同时满足两件事:既能按开放域文本指令走,又得让和声、节奏、旋律结构立得住。以重建为导向的音频 tokenizer 会把音乐结构和精细声学细节混在一起,产出高熵 token 不好建模;而模型若只是端到端硬生,编曲结构就成了隐式的、不可控的副产物。 本文贡献:StepAudio 3 Music 走离散与连续混合的设计。StepAudio Music Tokenizer 把音频表示成 50 Hz 的单码本流(码本 65536 条),靠语义引导的自监督与多任务训练同时保住音乐结构和重建所需信息;论文用受控对比(统一 25 Hz 帧率)比较了单码本 VQ、Semantic RVQ、Acoustic RVQ 三种离散瓶颈,以及不同 DiT 配置,才定下这个方案。渲染侧由流匹配 DiT 预测连续的 StepAudio VAE 潜变量,再由 VAE 解码器还原 48 kHz 音频。显式规划则交给一个 MoE 自回归模型:它先用 ABC 记谱法产出中间编曲方案(ABC-CoT),再去预测音乐 token,让和声、节奏、旋律结构成为生成上下文的一部分。渐进式训练课程加监督微调支撑歌曲与纯器乐生成、由干声生成伴奏、翻唱合成,最长 5 分 30 秒。 实验效果:经 DPO 强化学习后,模型在 AudioBox 的 Content Enjoyment、Content Usefulness、Production Quality 三项以及 MuQ-MuLan 相似度上取得所评系统中的最高分,SongBench 结果具备竞争力。在 Artificial Analysis Music Arena Vocals 初步榜上 Quality Elo 为 1105。 批判点评:「先写 ABC 谱再生成音频」把编曲从隐式副产物提成了显式可读的中间态,这既是可控性抓手也是可调试性抓手,是这篇最有辨识度的设计;单码本 50 Hz 加流匹配 DiT 的分工,也把「离散好建模」和「连续好还原」各自的长处分开用了。但榜单要看细账:Vocals 榜上它实际是第 4 名(前面是 Suno V5.5 的 1170、Mureka V9 的 1159、Mureka V8 的 1140),摘要里「仅次于 Suno V5.5 和 Mureka」的说法把 Mureka 两个版本折叠了,读起来比实际名次靠前。更关键的是它的投票样本只有 2119,而同榜 Suno、MiniMax 等都在 1 万以上,95% 置信区间也宽到 ±14,和第 5 名 Suno V5(1097,±6)的差距在统计上并不稳。论文自己也说明这些评测衡量的是生成质量,并不直接度量对单个 ABC 音符、和弦、小节的遵循程度——也就是说 ABC-CoT 到底被执行得多准,目前没有直接证据。 6. PhysStream:边生成边推物理的流式视频 PhysStream: Streaming Physics-Grounded Video Generation with Structured Scene Memory and Fine-Grained Motion Control | University of Pennsylvania × Snap Inc. × KAUST | arXiv:2609.17521 关键词:视频生成,物理一致性,自回归,交互控制 前序问题:视频生成的交互控制正从粗糙提示走向对动态场景做细粒度、物理上有意义的操纵。但现有可控方法要么要求生成开始前就给出完整控制时间表,要么用像素空间信号去规定物体位置——那是在指派位置,而不是在指定物理动力学。 本文贡献:PhysStream 是面向物理接地的图生视频的自回归模型,带两样东西。一是结构化场景记忆:位置图与物体跟踪图都从已生成的帧里在线导出,边生成边更新。二是用稀疏的速度增量信号做细粒度运动控制,这些信号编码的是物理量,让模型去学底层动力学而不是背位置。训练分两阶段:先用运动控制条件微调一个双向模型,再训练一个因果自回归模型并额外接入结构化场景记忆,进一步提升物理一致性。 实验效果:PhysStream 支持对多物体桌面刚体场景做交互式的、生成中途的控制——这是此前方法不具备的能力。合成基准上运动分布距离(FVMD)比最强基线降低 33%,轨迹误差降低 12%;野外对比中超过 85% 的情况被人类评估者更偏好。长时程上,配图里一只玉色茶杯在桌面随机走动,从 t=0 一路走到 t=180 帧,远超 49 帧的训练窗口,仍持续跟随随机注入的速度增量交互并保持外观不崩。 批判点评:「控制信号编码物理量而非像素位置」这个取舍是关键:它把可控性问题从「让模型听话地摆物体」改成「让模型学会动力学然后自己推演」,这也是它能支持生成中途插入交互、而不需要预先排好完整控制表的根因。结构化场景记忆从已生成帧在线导出,避免了额外的外部状态追踪,工程上干净。但作者自己单列了一张图说明一致性指标的局限——FlashMotion 的一致性分数和本方法相当,却有明显伪影和幻觉物体,DragStream 生成近乎静止的画面也能拿高分,这等于承认 33% 和 12% 这两个提升所依赖的指标体系本身可信度有限,真正支撑结论的是那 85% 的人类偏好。适用范围也偏窄:主战场是多物体桌面刚体,非刚体只给了弹跳球和布料两个微调后的演示,离通用物理场景还远。 7. DeepSeek-V4.1-Flash:每token只留890字节KV DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression | DeepSeek-AI | arXiv:2609.19969 关键词:KV cache压缩,MoE,长上下文,推理成本 前序问题:长时程智能体大规模铺开后,模型负载越来越偏输入重。尽管此前工作已大幅降低长上下文计算成本,prefill 仍然很贵,而庞大的 KV cache 继续挤占 HBM 与 SSD 容量、吃掉数据传输带宽。算力、存储、带宽这三份需求合起来,构成了进一步压低部署成本的主要瓶颈。 本文贡献:DeepSeek-V4.1-Flash 是一个多模态 MoE 模型,骨干 552B 参数,支持最长 100 万 token 上下文。它采用因果编码器-解码器(CED)架构:解码时每 token 激活 16B 参数,prefill 时只激活 8B,这对智能体这类输入重负载显著改善成本效率。为把 KV cache 压缩推到极限,模型把 Compressed Sparse Attention 2(CSA2)里的跨层 KV cache 复用与 FP4 KV 缓存结合起来,使常驻 HBM 的全局 KV cache 降到每 token 890 字节,约为 DeepSeek-V4-Flash 的四分之一。再通过名为 SWA Bounded Replay 的部署优化,把常驻 SSD 或主机内存的持久 KV cache 压到约八分之一。论文另外精简了 V4 架构并引入若干高效扩展,在 45T token 的多模态语料上预训练并做了完整后训练。 实验效果:KV cache 占用大幅缩小的同时,性能明显好于基线。单 token decode FLOPs 随上下文增长的曲线上,V4.1-Flash 到 1024K 仍基本持平在 25 GFLOPs 上下,而 V4-Flash 已升到约 55、V3.2 约 600、V1 超过 3000。预训练阶段在自建留出集上,V4.1-Flash-Base 的 BPB 在所有任务上均低于 V4-Flash-Base 与 V4-Pro-Base。代码智能体基准上性能随 RL 训练规模持续提升,把最大上下文进一步扩到 1M token 后,在极长时程任务上还能继续涨。模型权重已在 HuggingFace 放出。 批判点评:890 字节每 token 这个数字值得记住:它把长上下文的讨论从「模型记不记得住」彻底拽到了「cache 放不放得下、搬不搬得动」,而 prefill 8B、decode 16B 的非对称激活更是直接冲着智能体负载「输入远大于输出」的真实形状去的,属于按负载画像改架构。但 FLOPs 曲线也显示,在约 100K 以下的上下文里 V4.1-Flash 反而略高于 V4-Flash(约 23 对 19 GFLOPs),这套设计是拿短上下文的一点效率换长上下文的大幅收敛,主要跑短请求的场景未必划算。另外 FP4 KV 缓存在什么任务上会先露出精度问题、SWA Bounded Replay 对首 token 延迟的影响如何,摘要都没交代;593 位作者的技术报告里工程优化与架构贡献高度耦合,外部想复现出同样的 890 字节与同等质量,难度不低。 8. The Weight Is Over:12GB显卡跑亚秒出图 The Weight Is Over - Interactive Diffusion on Consumer GPUs | Adobe × NVIDIA | arXiv:2609.21849 关键词:端侧推理,扩散加速,文本编码器蒸馏,显存预算 前序问题:端侧推理正在爆发,但势头几乎全在语言模型那边。扩散管线吃显存、对延迟敏感,还要编排文本编码器、Transformer、解码器以及后处理若干环节,这套流程远没有大模型推理循环那么标准化,落到消费级设备上格外难办。 本文贡献:论文在性能、质量、模型体积这个三角里找可落地的点,给出三项贡献:一个嵌入翻译器,把小文本编码器映射到大编码器的表示空间,从而砍掉权重与延迟;一套可复现的扫参配方,用来在扩散管线的速度/质量/显存三角里做导航;以及一个端侧交互式图像生成编辑器,在较新的 GPU 上实现亚秒级首图时间(TTFI)。 实验效果:跨设备扫参结果最能说明问题。RTX 4070 Ti(12 GB,可用约 11.2 GB)上,只要常驻权重预算没超出可用显存,每步延迟稳定在数百毫秒量级;一旦越过 11.2 GB 就触发 host-paging、必须从主机流式取权重,每步延迟直接跳到 10^4 毫秒级,相当于一个断崖。工作站级 RTX PRO 6000 Blackwell(96 GB)上「什么都放得下」,此时主导排序的就变成精度——NVFP4 最快,FP8 次之,FP16/BF16 落后。嵌入翻译器的容量门槛也量出来了:187M 的翻译器能紧跟 4B 参考编码器,29M 就会漂到「看着合理但内容不对」(要狐狸给猫、要拉面给黄色玩偶),3.5M 则塌缩成一种同质的暗黑森林风格。 批判点评:这篇的实用价值不在某个新模块,而在把「显存预算」立成了一等约束并画出了那条断崖:越过可用显存的惩罚不是线性变慢而是一到两个数量级,这意味着端侧调优的第一目标应该是「压进预算」而不是「提高吞吐」,顺序搞反了后面全是白做。把翻译器容量与语义保真的关系摸到 187M/29M/3.5M 三个档,也给了很实际的选型下界。但工作偏工程报告:只有两位作者、两台设备,没有跨更多消费级显卡的统计;质量评价主要靠定性图,缺少 FID 之类的量化指标,「187M 紧跟 4B」到底差多少无法量化。亚秒 TTFI 也限定在「较新的 GPU」上,对更老的消费级卡是什么表现没有交代。 9. MuSeC:语义声学分流的音乐codec Rethinking Music Tokenization: A Semantic Codec toward High-Fidelity LLM Music Generation | 西北工业大学 ASLP × 吉利汽车研究院(宁波) | arXiv:2609.21240 关键词:音乐tokenizer,语义codec,残差量化,音乐信息检索 前序问题:离散音频 tokenization 已成为原始波形与自回归建模之间的关键接口,所以音乐 tokenizer 得同时做到两件互相拉扯的事:支持高保真重建,又要产出适合语言建模的离散序列。以重建为导向的 tokenizer 常把音乐结构和精细声学细节混在一起,产出高熵 token 难建模;而语义引导的替代方案本是为语音设计的,套到音乐上不合身,往往还伤重建质量。 本文贡献:论文先把「音乐语义内容」重新定义成一个可度量的概念——以下游音乐信息检索(MIR)任务表现为锚。循着这个定义提出 MuSeC:一个音乐语义 codec,直接从混合信号里把语义与声学内容解耦,不需要做源分离。结构上,冻结的 SSL 编码器加 k-means 提供语义 token,另一路声学 RVQ 流捕捉细粒度声学信息,两路拼接后送解码器做高保真重建,训练时还有判别器提供对抗损失。 实验效果:MuSeC 在保住高保真重建所需信息的同时,产出对语言模型更友好的离散单元,重建质量有提升,token 序列也更可预测,为高保真 LLM 音乐生成提供了实际基础。MIR 任务上的自消融显示,完整 MuSeC(S1A16)在 GuitarSet/Chords(0.3648 对 0.1713)、EMO/R2A(0.6520 对 0.6110)、EMO/R2V(0.3696 对 0.3578)等任务上明显好于纯 k-means 量化版本,也普遍好于只有声学的 A16 变体。 批判点评:把「语义」从形容词变成以 MIR 任务表现为锚的可测量量,是这篇方法论上最扎实的一步——有了口径,语义与声学的分流才有得优化,而不是凭感觉设计码本。不做源分离直接从混合信号解耦,也省掉了一个容易引入误差的前置环节。但消融图显示优势并不一致:GTZAN/Genre 上 MuSeC-S1A16 是 0.5345,反而低于 k-means 的 0.6034,Lyrics 上两者基本持平(0.4903 对 0.4993);更重要的是全部设置都明显低于连续 LeVo BEST-RQ 这条语义上限(多数任务差 0.1 以上),说明离散化的语义损失远未补齐。另外论文标题指向「高保真 LLM 音乐生成」,但实际只做到 codec 这一层,接上语言模型后端到端能好多少并没有验证,「更可预测的 token 序列」目前也缺少熵或困惑度之类的直接数字支撑。 10. OmniVBench:1.2万条清单查参考跑没跑偏 OmniVBench: A Benchmark and Large-Scale Dataset for Omni Reference-to-Video Generation | 腾讯 × 香港科技大学(广州) | arXiv:2609.22069 关键词:视频生成评测,参考到视频,数据集,因子级评估 前序问题:参考到视频(R2V)生成正朝着越来越通用、越来越多样的参考控制演进,催生了所谓全模态 R2V 这一范式。但现有基准跟不上:测试用例覆盖的参考类型与组合都有限,评估协议大多只看整体参考一致性,忽略了参考因子是否被正确保留、解耦并路由到对应目标。同时全模态 R2V 训练数据构建成本高,合适的训练资源很稀缺。 本文贡献:论文同时给出 OmniVBench 与 Omni-R2V Dataset,一个管评测一个管训练。OmniVBench 把 R2V 评测扩展到更广的参考类型、更细的控制任务和更丰富的参考组合,覆盖 7 个任务族、18 个细粒度任务,横跨内容、运动、风格、结构、叙事与多参考设定,共 813 个评测用例。核心是因子级评估:用 12172 条针对具体用例的检查清单条目,逐条判断意图中的参考因子是否被忠实保留、是否被正确解耦并绑定到目标、是否按指令实现。Omni-R2V Dataset 主要取材于大规模专业视频素材库,包含 34 万条处理好的训练样本,覆盖多种参考类型与多参考组合,并给出各任务专用的参考-目标配对构建管线。 实验效果:对先进的开源与闭源 R2V 模型做的大范围评测显示,各任务族和各评估维度上都存在明显性能差距,暴露了当前 R2V 模型的残留局限。配图给出的多参考用例里,RF(参考保真)/IR(指令遵循)两项分数分化极大:UniVideo 只有 33.2/14.8,LoomVideo 43.1/36.7,OmniWeaving 48.0/56.0,Kling 3.0 Omni 62.9/100,Bernini 85.2/69.7,Gemini Omni 91.2/100,而 MiniMax H3、Seedance 2.0 与 Seedance 2.5 在该用例上拿到 100/100。 批判点评:把「参考一致性」拆成保留、解耦、路由三件事,是这套基准最有价值的判断——多参考场景真正容易错的不是「像不像」,而是把 A 参考的属性贴到了 B 目标上,而整体一致性分数恰恰看不出这类串台。12172 条逐用例清单让这种错误变得可指认。数据侧 34 万条专业素材样本对社区也是实打实的补给。但代价是清单法对判定器的依赖很重:谁来回答这 1.2 万条问题、判定器自身的偏差和一致性如何,摘要没有交代,而这直接决定分数可信度。813 个评测用例摊到 18 个细粒度任务上平均每任务仅四十余例,统计功效对细分结论偏弱;配图那个用例里三个模型同时打满 100/100,也提示清单在强模型区间容易饱和、区分度下降。另外基准与数据集同源、且出自同一团队,用该数据训练的模型在该基准上的成绩需要额外小心解读。 趋势观察 注意力的省法,从「换掉它」变成「按交互类型分流」 Video DeltaNet 只在视频与视频之间用线性注意力,凡涉及文本或音频的交互仍留给 Softmax;dQwen3.5 则把注意力与 RNN 交错的混合骨干当成扩散语言模型的适配起点。两篇的共同前提是承认不同交互对精度的敏感度本来就不同,一刀切替换必然在某处付出代价。 长上下文的成本,已经从算力转到显存和带宽 DeepSeek-V4.1-Flash 把常驻 HBM 的全局 KV 压到每 token 890 字节、持久部分再压到八分之一,并用 prefill 8B、decode 16B 的非对称激活去贴合智能体负载的真实形状;The Weight Is Over 则量出消费级显卡越过可用显存后延迟跳升两个数量级的断崖。两者都在说:放不放得下、搬不搬得动,已经比算得快不快更决定部署成本。 生成过程正在被拆出显式的中间态 StepAudio 3 Music 先用 ABC 记谱产出编曲方案再预测音乐 token,PhysStream 把位置图与跟踪图作为结构化场景记忆在线维护,Edit-VAR 缓存源概率与交叉注意力图再逐 token 决定保源还是改。把过去隐含在权重里的东西显式化,换来的是可控性与可调试性。 评测开始追问「参考有没有被放对位置」 OmniVBench 把参考一致性拆成保留、解耦、路由三件事,用 1.2 万条逐用例清单指认把 A 的属性贴到 B 上的串台;Paint-Anything 自建 ACBench 专量物体级十六进制色值保真;MuSeC 干脆把「音乐语义」定义成下游 MIR 任务表现。三者都不再满足于一个整体相似度分数。 人工智能炼丹君 整理 | 2026-09-21 更多 AIGC 论文解读,关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注
2026年09月21日
6 阅读
0 评论
0 点赞
2026-09-20
AIGC 基本功|视频 VAE 的常见 loss 组合-VAELoss
视频 VAE 的常见 loss 组合 所属方向:表征与压缩 | 难度:工程实战 | 前置知识:视频 VAE 的时空压缩结构(本文第 02 节另补最小背景) 关键词:重建损失、KL 损失、LPIPS、感知损失、GAN 损失、判别器、loss 权重 01. 为什么需要它 训练一个视频 VAE,最朴素的想法是「编码器压缩、解码器还原,加个 L1 或 KL 就够了」。不同目标可能暴露不同问题,效果取决于瓶颈、数据和训练配方: 只用 L1/L2 + KL:画面是稳的,但像蒙了一层猪油 —— 头发丝、文字、树叶这种高频纹理全被抹掉,人脸带着一层「塑料感」。在存在重建不确定性时,这与逐像素回归的性质有关,但不能断言 L1/L2 + KL 必然不可用(L2 的最优解是条件期望,天然倾向平均、发糊)。 加上 GAN 想救清晰度:判别器一旦过早发力或过强,解码器立刻走样 —— 出现棋盘格、闪烁、甚至凭空捏造纹理;训练损失曲线看起来在降,重建却越来越假。 把图像那套搬到视频上:每一帧单独看 PSNR 都不错,连起来播放却疯狂闪烁,相邻帧的纹理在抖动。单帧指标根本测不出时间维的失真。 这不是三个孤立的 bug,而是三类不同的失真,需要三种不同的尺子去量。视频 VAE 的训练目标因此通常是四项的组合:像素重建、KL 正则、感知损失、对抗损失,有时再加一项时间一致性约束。玄学的地方从来不是「用哪几项」—— 这几项高度趋同 —— 而是权重配比:为什么 HunyuanVideo 的 KL 权重能小到 $10^{-6}$,而感知项是 $0.1$、对抗项是 $0.05$?这些数字不是拍脑袋,是被各项的量纲、归一化口径和训练阶段逼出来的。 本文把每一项「在管什么失真、数学上长什么样、量级有多大、什么时候帮倒忙」逐一拆开,所有关键数字都来自随文可跑的脚本(只依赖 numpy)。 02. 最小可用理解 最小背景(前置 VideoVAE):视频 VAE 用 3D 因果卷积把一段视频 $x\in\mathbb{R}^{B\times C\times T\times H\times W}$ 压成低维潜变量 $z$(典型压缩率:时间 $4\times$、空间 $8\times$),解码器再从 $z$ 重建出 $r$。编码器输出一个高斯后验 $q(z\mid x)=\mathcal{N}(\mu,\sigma^2)$,训练目标要同时满足三件事:重建要像、潜变量分布要规整(好让后面的扩散模型去建模)、压缩率要够高。loss 就是在这三者之间做权衡的旋钮。 三句话讲清四项损失: 像素损失(L1/L2)和 KL 是「保正确」的:前者逐像素对齐内容与结构,后者把潜变量摁在标准正态附近、防止它为了重建而无限膨胀。 LPIPS 感知损失和 GAN 对抗损失是「保好看」的:LPIPS 在预训练深度网络的特征空间里比较,惩罚人眼在意的语义 / 纹理差异;GAN 让一个判别器去挑刺,逼着解码器还原高频细节。 这四项量纲完全不同,必须加权配平,GAN 常用「晚启动 + 自适应权重」来控制训练平衡:先让重建项把解码器教到成形,再让判别器入场,并按两项在最后一层上的梯度范数之比动态调权。视频还要额外盯时间维(时空判别器或帧间一致性项)。 03. 数学推导 3.1 从 ELBO 到「重建 + KL」两项 VAE 最大化证据下界(ELBO)。写成「最小化负 ELBO」,目标天然裂成两项: $$\mathcal{L} = \underbrace{-\mathbb{E}_{q_\phi(z\mid x)}[\log p_\theta(x\mid z)]}_{\text{重建项}} + \underbrace{D_{\mathrm{KL}}\big(q_\phi(z\mid x)\,\|\,p(z)\big)}_{\text{KL 正则项}}$$ 逐项解释符号:$\phi$ 是编码器参数、$\theta$ 是解码器参数;$q_\phi(z\mid x)$ 是编码器给出的后验分布;$p_\theta(x\mid z)$ 是解码器的似然;$p(z)=\mathcal{N}(0,I)$ 是标准正态先验。重建项要求「从采样出的 $z$ 能还原 $x$」,KL 项要求「后验别离先验太远」—— 它是潜变量的正则项,鼓励与先验对齐;它不保证感知空间的平滑或连通,也不是所有 latent 扩散能够学习的必要条件。 当后验取对角高斯 $q_\phi(z\mid x)=\mathcal{N}(\mu_\phi(x),\,\mathrm{diag}(\sigma_\phi(x)^2))$、先验取标准正态时,KL 有解析解(不用采样、不用估计)。令 $\ell_j=\log\sigma_j^2$(工程上网络直接预测 $\ell$,数值更稳): $$D_{\mathrm{KL}} = -\frac{1}{2}\sum_{j=1}^{d_z}\Big(1+\ell_j-\mu_j^2-\exp(\ell_j)\Big)$$ 这里求和下标 $j$ 跑遍所有潜变量维度。直觉:$\tfrac12\mu_j^2$ 惩罚均值偏离 0,$-(1+\ell_j-e^{\ell_j})$ 惩罚方差偏离 1($\ell=0$ 即 $\sigma^2=1$ 时该项为 0)。 第一个容易被忽略的坑是「求和 vs 平均」。教科书公式对 $d_z$ 个维度求和;而像素损失通常对全部像素求平均。两者元素个数差着几个数量级,直接相加 KL 会凭「项数多」碾压重建。随文脚本 vaeloss_terms.py 在一个合成视频上实测: # 解析 KL,logvar 是网络直接预测的 log(sigma^2) kl_per_element = -0.5 * (1 + logvar - mu ** 2 - np.exp(logvar)) kl_sum, kl_mean = kl_per_element.sum(), kl_per_element.mean() l1 = np.mean(np.abs(x - r)) print("L1 = %.5f" % l1) # 0.03158 print("KL(sum) = %.3f KL(mean) = %.4f" % (kl_sum, kl_mean)) print("KL(sum)/L1 = %.0f 倍" % (kl_sum / l1)) 输出(视频 $2\times3\times8\times32\times32$,latent $2\times4\times2\times4\times4$,像素 49152 个、全 batch 的 latent 共 256 个数值(每样本 128 个)): L1 重建 (mean) 0.03158 KL (全 batch sum,教学口径) 57.66490 KL (mean,换口径) 0.22525 KL(sum)/L1 的量级倍数 : 1826.0x 归约口径是开源视频 VAE 的 KL 权重有时出现 $10^{-6}$ 这种「看着离谱」的数字 —— 的原因之一,但不能由这个 toy 直接反推实际系数。通常应先对每样本 latent 求和,再对 batch 平均;这里为演示求和效应,57.66 是整个 batch 的和,batch=2 时常见口径为 28.83。抄权重之前先对齐归一化口径,这句话后面还会强调。 这张图要看什么:左图是四项损失在同一个合成视频上的原始数值(对数轴)——KL(sum) 一根柱子顶到 57.66,是 L1 的 1826 倍;右图套上 HunyuanVideo 的配方权重之后,各项贡献变为不同数量级,不能称为完全配平,最矮的 KL 贡献只有 $5.8\times10^{-5}$。权重同时反映口径、梯度和重建目标,不能仅按标量 loss 大小配平(图由 code/make_figures.py 生成,下同)。 3.2 像素损失:为什么多用 L1 而不是 L2 $$\mathcal{L}_{\mathrm{pix}} = \frac{1}{N}\sum_{n=1}^{N}\big\|x_n-r_n\big\|_1$$ L2(MSE)对大误差平方加权,其逐点最优解是给定 $z$ 下所有可能输出的条件均值—— 面对「这块纹理可能是 A 也可能是 B」的不确定性,它选择把 A、B 平均掉,结果就是模糊。L1 在非零残差处的梯度为 $\pm1$,其逐点最优解是条件中位数,不会因为误差大就给出更强的「往平均靠」的驱动力,对异常值通常更稳健,但条件中位数也可能模糊,不能保证所有纹理更锐利。代价是 L1 在零点不可导、对小误差的梯度恒定,容易留下颗粒感 —— 这正是要靠感知 / 对抗项补的地方。也有工作(如 LTX-Video)在像素项里混用 MSE 与小波域 L1(Video-DWT),在多尺度上约束。 3.3 感知损失 LPIPS:换一把「人眼的尺子」 像素距离有个致命问题:同样大小的像素误差,人眼感受天差地别。一个全局亮度偏移,MSE 不小但人眼几乎无感;一团等量的逐点噪声,MSE 相同却把纹理毁了。LPIPS(Learned Perceptual Image Patch Similarity)改用在 ImageNet 上预训练的分类网络(VGG/Alex,参数冻结)提特征,再在特征空间量距离: $$\mathcal{L}_{\mathrm{LPIPS}}(x,r) = \sum_{k}\frac{1}{H_kW_k}\Big\|\,w_k\odot\big(\hat{y}_x^{k}-\hat{y}_r^{k}\big)\Big\|_2^2$$ 符号:$y_x^k$、$y_r^k$ 是第 $k$ 层(LPIPS 用 VGG 的 relu1_2 到 relu5_3 共 5 层)对真实图和重建图提的特征图;$\hat{y}$ 表示沿通道做了归一化(除以通道维 L2 范数);公式中的 $w_k^2$ 对应实现里作用在平方特征差上的 $1\times1$ 非负通道权重;$w_k$ 可理解为它的平方根,而非直接把卷积权重再平方。该权重在人类主观偏好数据集 BAPPS 上学出来、之后冻结;最后空间平均、跨层求和。两个细节缺一不可:逐通道归一化让比较不被某些高幅值通道主导,学习权重 $w_k$ 让「哪些层的差异人眼更在意」由数据决定。它天然偏向语义 / 结构 / 纹理,而对整体明暗、轻微色偏不敏感。 pixel_vs_perceptual.py 用一个免权重的多尺度高通特征(高斯差分 DoG)复现 LPIPS 的结构,对比三种退化: # A:全局亮度偏移;B:同等 L2 能量的逐点噪声;C:高斯模糊 rA = x + delta # delta = 0.12 rB = x + rng.normal(0, delta, x.shape) # 标准差同为 0.12 # 像素 MSE:A 约等于 B;特征距离:B 远大于 A 真实输出: 重建方式 像素MSE 感知距离 感知/像素 A 亮度偏移(整体+0.12) 0.01440 0.00008 0.01 B 逐点噪声(σ=0.12) 0.01426 0.64084 44.93 C 高斯模糊 0.00508 0.50163 98.69 A、B 的像素 MSE 只差 0.95%,但特征距离里 B 是 A 的约 $8\times10^3$ 倍;模糊 C 的像素误差最小,感知距离却很高。高通教学代理天然抑制直流亮度,因此本例差距尤其大;真实 LPIPS 使用学习特征和权重,不能据此断言它对亮度或颜色变化不敏感。 3.4 对抗损失:雇一个判别器专挑高频毛病 GAN 引入一个判别器 $D$,训练它区分真实帧与重建帧;解码器(生成器)则努力骗过它。视频 / 图像重建里几乎都用 PatchGAN 式判别器:不输出整图真假,而是对每个空间 patch 打分,主要约束其感受野内的统计,常改善局部纹理,但并非数学上只看高频。最常用的 hinge 形式: $$\mathcal{L}_{D} = \tfrac{1}{2}\Big(\mathbb{E}_{\text{real}}[\max(0,\,1-D(x))]+\mathbb{E}_{\text{recon}}[\max(0,\,1+D(r))]\Big)$$ $$\mathcal{L}_{G}^{\mathrm{adv}} = -\mathbb{E}_{r}[D(r)]$$ 判别器希望真帧打分大于 1、重建打分小于 -1;生成器希望重建打分尽量大(为正)。gan_schedule.py 用合成 logits 算了判别器在三种强弱下的损失: 起步:判别器分不清 D(hinge)=1.0126 G(hinge)=-0.0312 real≈ 0.01 fake≈ 0.03 健康:适度拉开 D(hinge)=0.0767 G(hinge)= 1.2067 real≈ 1.21 fake≈-1.21 过强:margin 已饱和 D(hinge)=0.0000 G(hinge)= 5.9713 real≈ 6.02 fake≈-5.97 注意判别器 hinge loss 为 0 只说明这些样本的 margin 已满足,此时判别器对应损失的梯度为 0;不意味着生成器梯度消失。这里 $\mathcal L_G=-\mathbb E[D(r)]$ 对 fake logit 的导数仍是 −1,传回解码器的梯度还取决于判别器对输入的导数。仅看 +6/−6 logits 无法判断梯度是否健康。 其一,GAN 晚启动(warm-up)。重建项还没把解码器教出基本形状时,判别器挑的「毛病」没有意义,甚至会把训练带偏。taming 的做法是一个开关 adopt_weight,在第 $t_{\mathrm{start}}$ 步前把对抗权重置 0: $$\delta(t)=\begin{cases}0,&t<t_{\mathrm{start}}\\ \lambda_{\mathrm{adv}},&t\ge t_{\mathrm{start}}\end{cases}$$ 脚本实测 disc_start=2000 时,step 0 和 1999 的权重为 0,step 2000 起跳到 0.05。 其二,自适应对抗权重。固定 $\lambda_{\mathrm{adv}}$ 的两难:训练早期重建梯度很大、GAN 抢不过;后期重建收敛、梯度变小,同样的 GAN 梯度又会相对越来越强甚至压过重建。VQGAN 的解法是让两项在解码器最后一层权重 $w_L$ 上的梯度范数之比来决定权重: $$\lambda_{\mathrm{adv}}(t)=\mathrm{clip}\left(\frac{\|\nabla_{w_L}\mathcal{L}_{\mathrm{rec}}\|}{\|\nabla_{w_L}\mathcal{L}_{G}^{\mathrm{adv}}\|+\epsilon},\ 0,\ 10^4\right)\cdot\delta(t)$$ 直觉:它动态地让「对抗项在最后一层上产生的梯度量级」与「重建项的梯度量级」匹配,重建没收敛时比值大、收敛后比值自动变小。adaptive_weight.py 用一个可解析求导的迷你线性解码器精确计算两个梯度范数,并用中心差分验证(解析 0.012371 对差分 0.012371;0.255913 对 0.255913): 训练早期(未收敛) L_rec=0.73740 ||∇rec||=0.4071 ||∇gan||=4.2840 d_weight=0.0950 训练后期(近收敛) L_rec=0.00071 ||∇rec||=0.0127 ||∇gan||=4.2840 d_weight=0.0030 重建收敛后自适应权重从 0.095 掉到 0.003(约 32 倍),GAN 项被自动调小 —— 这正是固定权重给不了的能力。 这张图要看什么:横轴是合成出来的训练进程(从「解码器还没学会」到「重建已收敛」),三条曲线分别是重建损失、重建项在最后一层上的梯度范数、以及据此算出的 $\lambda_{\mathrm{adv}}$。要点是 $\lambda_{\mathrm{adv}}$ 不是人为排的衰减计划,而是被 $\|\nabla\mathcal{L}_{\mathrm{rec}}\|$ 拖着走的:早期 0.095、后期 0.003。对照上面的 warm-up 开关看更清楚——$\delta(t)$ 负责「第 2000 步之前完全不启用」,这条曲线负责「启用之后给多大」,两者相乘才是最终权重。 3.5 总目标,以及视频多出的时间维 把四项合起来,连续潜变量视频 VAE 的生成器目标是: $$\mathcal{L}_{G} = \lambda_{\mathrm{pix}}\mathcal{L}_{\mathrm{pix}}+\lambda_{\mathrm{p}}\mathcal{L}_{\mathrm{LPIPS}}+\lambda_{\mathrm{kl}}\mathcal{L}_{\mathrm{KL}}+\lambda_{\mathrm{adv}}(t)\,\mathcal{L}_{G}^{\mathrm{adv}}$$ (VQGAN 是离散码本,没有 KL,对应位置换成 codebook/commitment 损失;连续 KL-VAE 才是上面这版。)判别器另用 $\mathcal{L}_D$ 单独更新,二者交替。 视频比图像多一维,单帧损失管不到帧间。temporal_consistency.py 构造了一段运动视频,给两种重建:A 加逐帧独立噪声(播放时闪烁),B 加跨帧恒定的退化(不闪),两者单帧空间 L1 几乎相同: 重建 单帧空间L1 时序差分L1 时序差分MSE A 逐帧独立噪声(闪) 0.06393 0.09018 0.01275 B 跨帧恒定退化(稳) 0.06340 0.00455 0.00003 单帧 L1 几乎相等(0.0639 对 0.0634),时序差分 MSE 却差了约 393 倍。 这张图要看什么:上半部是同一个像素点随帧走的亮度轨迹。A(逐帧独立噪声)和 B(跨帧恒定退化)的逐帧平均误差几乎一样,但 A 的轨迹在真值附近高频锯齿抖动,B 的轨迹只是整体平移——前者播放起来就是闪烁,后者只是画质差一点。下半部把这件事量化:单帧空间 L1 两根柱子几乎齐平(0.0639 / 0.0634),时序差分 L1 差 20 倍,时序差分 MSE 差到 393 倍。所以「视频 VAE 的重建指标好看但看着闪」,根因是指标选错了:空间指标对时间频率不敏感。 补救有两条路:一是把判别器从 2D 扩到时空(3D/PatchGAN、混合外观 - 运动判别器),让它直接看片段;二是显式加时间一致性损失,用光流把相邻帧 warp 过来再比(静止区域退化成帧差): $$\mathcal{L}_{\mathrm{temp}}=\frac{1}{T-1}\sum_{t=1}^{T-1}\big\|r_t-\mathcal{W}(r_{t-1},\,F_{t\to t-1})\big\|_1$$ 其中 $\mathcal{W}$ 是 warp 算子、$F$ 是估计的光流;没有光流时可用相邻帧差分近似 $\mathcal{L}_{\Delta}=\frac{1}{T-1}\sum_t\|(r_t-r_{t-1})-(x_t-x_{t-1})\|_1$。 04. 代码实现 随文 5 个脚本只依赖 numpy,python 脚本名.py 即可运行,完整版在文末附录;另有 make_figures.py 需要 matplotlib,负责本文三张配图。这里串起主干。 (1)四项损失与量级对照(vaeloss_terms.py):在合成视频上算 L1、解析 KL(sum/mean 两种口径)、感知代理、GAN,核心是 KL 解析式: kl_per_element = -0.5 * (1 + logvar - mu ** 2 - np.exp(logvar)) kl_sum = kl_per_element.sum() # 教科书口径:对 latent 维求和 l1 = np.mean(np.abs(x - r)) # 工程口径:对像素求平均 print("KL(sum)/L1 = %.0f 倍" % (kl_sum / l1)) # 1826 倍 (2)像素 vs 感知(pixel_vs_perceptual.py):多尺度高斯差分特征加逐通道归一化,复现 LPIPS「在特征空间量距离」的结构。再次强调这是教学代理:本体用的是在 BAPPS 上学过权重的 VGG(见第 05 节真实代码)。 (3)GAN 损失与 warm-up(gan_schedule.py): def hinge_d(real, fake): return 0.5 * (np.mean(np.maximum(0.0, 1.0 - real)) + np.mean(np.maximum(0.0, 1.0 + fake))) def adopt_weight(weight, step, threshold=0, value=0.0): return value if step < threshold else weight (4)自适应权重(adaptive_weight.py):对迷你线性解码器用解析梯度(并用中心差分校验)算范数比: d_weight = np.clip(np.linalg.norm(g_rec) / (np.linalg.norm(g_gan) + 1e-4), 0.0, 1e4) print("早期 %.3f -> 后期 %.4f" % (w_early, w_late)) # 0.095 -> 0.003 (5)时间一致性(temporal_consistency.py):比较单帧空间 L1 与相邻帧差分误差,证明只有后者能抓到闪烁。 这些脚本的真实输出已散落在第 03 节,文末附录给出可直接运行的完整源码。 05. 工业级实现对照 最小实现是为了讲清原理,生产代码有几处关键的工程化。最权威的参照是 VQGAN 的损失模块(知识树锚点): CompVis/taming-transformers/taming/modules/losses/vqperceptual.py → VQLPIPSWithDiscriminator.forward(以 2026-09 的实现为准) 对照本文的四项,它的生成器一步几乎是公式的逐行翻译: rec_loss = torch.abs(inputs - reconstructions) # L1 像素 if self.perceptual_weight > 0: p_loss = self.perceptual_loss(inputs, reconstructions) # LPIPS rec_loss = rec_loss + self.perceptual_weight * p_loss nll_loss = torch.mean(rec_loss) logits_fake = self.discriminator(reconstructions) g_loss = -torch.mean(logits_fake) # hinge 生成损失 d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer) loss = nll_loss + d_weight * disc_factor * g_loss \ + self.codebook_weight * codebook_loss.mean() # VQ:codebook;连续 VAE 换成 KL 与最小实现的差异,每一处都有来由: 重建与感知合并成 nll_loss 一起算梯度:自适应权重需要对「合并后的重建项」和「GAN 项」分别在最后一层求梯度(torch.autograd.grad(..., retain_graph=True)),所以 LPIPS 不是独立加权、而是并进重建项。 双优化器、两次前向:optimizer_idx==0 更新生成器(含重建、LPIPS、codebook、GAN),==1 才更新判别器;判别器那一路对输入 .detach(),不让梯度流回解码器。 LPIPS 本体(taming/modules/losses/lpips.py):先过一个 ScalingLayer(把通常约定为 $[-1,1]$ 的输入变到 VGG 的归一化统计,shift/scale 是写死的常数),再取 VGG16 的 5 个 relu 特征,逐通道归一化、过学来的 $1\times1$ 卷积、空间平均、跨层求和 —— 正是公式 3.3 的完整实现,权重从 BAPPS 预训练 ckpt 加载且全程冻结。 连续视频 VAE 用 KL 约束替代离散码本损失,但具体实现还可能改变重建项归约、学习似然尺度、判别器和时序结构,不能机械替换一行就认为目标完全相同。 当代视频 VAE 的公开配方(权重都来自各自技术报告,不要跨项目照抄,口径不同): 模型 损失组合与权重(论文原文) HunyuanVideo (2412.03603) $\mathrm{L_1}+0.1\,\mathrm{L_{lpips}}+0.05\,\mathrm{L_{adv}}+10^{-6}\,\mathrm{L_{kl}}$;判别器做随机缩放加时间维扩展,视频与图像从零联合训练 LTX-Video (2501.00103) 像素 MSE 加 Video-DWT(小波域 L1)加 LPIPS 加 Reconstruction-GAN;并讨论 causal /non-causal VAE 的取舍 Seedance 1.0 (2506.09113) L1 加 KL 加 LPIPS 加对抗损失;用类 PatchGAN 的混合判别器同时约束外观与运动 H3AE (2504.10567) 反方证据:判别类损失收益小却显著拖慢训练,主张先用 L1+KL 收敛、再用潜空间一致性损失微调 表中只有 HunyuanVideo 给出了这里列出的明确系数,其他项目采用不同目标或归一化,不能据此称权重“高度趋同”。toy 的数值只解释口径为何重要,不能从 1826 倍损失比推导出真实训练必需的 $10^{-6}$ 权重。 06. 代价与边界 每一项都在解决一类失真,也都引入新的代价: L1/L2:稳、好训,但 L2 糊、L1 颗粒重;它们只能保证「像素对」,保证不了「看着真」。 KL:去掉它可能使后验更确定、尺度约束变弱,但没有 KL 的自编码器或 VQ latent 也可以训练扩散模型;但权重过大、瓶颈太紧,重建细节会被牺牲。$10^{-6}$ 这种小权重是「弱正则」,依赖特定归一化口径,换套实现可能就要重新定标。 LPIPS:贴人眼,但它的「审美」被冻结在 VGG 的自然图像特征里 —— 对医学影像、动画、线稿等域外数据可能偏置;它偏纹理,有时会鼓励「看起来有细节」的伪纹理。 GAN:是清晰度和真实感的主要来源,也几乎是所有训练不稳的来源:需要 warm-up、谱归一化、限制判别器更新次数、自适应权重;并且提升感知质量往往以 PSNR 下降为代价(感知 — 失真权衡,perception–distortion tradeoff),这不是没训好,而是规律。 视频时间项 / 3D 判别器:能压闪烁,但显著增算力、增训练时长;光流估计本身在遮挡和大运动处会出错,warp 损失可能误伤真实运动。 边界也要讲清:H3AE 等近期工作指出,在高压缩 VAE 上判别类损失的边际收益可能撑不起它的训练成本,先用重建加 KL、后期再针对性微调是更划算的路线。四项全开不是政治正确—— 数据域、压缩率、训练阶段不同,最优组合也不同,应当用消融实验决定。 07. 经典论文脉络 Auto-Encoding Variational Bayes(VAE,arXiv:1312.6114,2013):提出 ELBO、重参数化与连续高斯潜变量,KL 正则的源头。 Neural Discrete Representation Learning(VQ-VAE,arXiv:1711.00937,2017):改走离散码本,用 codebook/commitment 损失替代 KL,是 VQGAN 的前身。 Image-to-Image Translation with Conditional Adversarial Networks(pix2pix / PatchGAN,arXiv:1611.07004,2017):把判别器做成局部 patch 判定器,确立了「局部对抗项与像素重建互补」的分工。 The Unreasonable Effectiveness of Deep Features as a Perceptual Metric(LPIPS,arXiv:1801.03924,CVPR 2018):用 BAPPS 人类偏好数据证明深度特征距离远胜 PSNR/SSIM,并学出逐层权重。 Taming Transformers for High-Resolution Image Synthesis(VQGAN,arXiv:2012.09841,2021):L1、LPIPS、PatchGAN、codebook 四件套定型,配套自适应 GAN 权重与 warm-up,是本文工业对照的母本。 CogVideoX(arXiv:2408.06072)与 HunyuanVideo(arXiv:2412.03603):把这套组合搬到 3D 因果视频 VAE,后者给出明确的四项权重和时空判别器设计。 H3AE(arXiv:2504.10567,2025):对「判别损失是否值得」提出反方证据,代表这条线仍在演进。 08. 常见误解 「KL 权重抄 HunyuanVideo 的 $10^{-6}$ 就行」:错。权重取决于你的 KL 是 sum 还是 mean、latent 与像素各有多少元素、有没有做 loss balancing。本文实测同一组数据 sum 口径 KL 是 mean 口径 L1 的 1826 倍;换个压缩率或归一化,$10^{-6}$ 可能过大或过小。先统一口径,再谈数字。 「LPIPS 就是拿 VGG 特征算 L2」:漏了两个关键件 —— 沿通道的归一化和在 BAPPS 上学出来的 $1\times1$ 权重;输入还要先过 scaling layer 换到 VGG 的数值域。少了归一化,距离会被少数高响应通道主导。 「生成器 GAN 损失算出来是负数,训练崩了」:hinge 生成器目标下 $\mathcal{L}_G=-\mathbb{E}[D(r)]$(logistic non-saturating 常写成 $\mathbb E[\mathrm{softplus}(-D(r))]$,是另一种公式),为负恰恰说明重建帧已经把判别器骗到打正分,是预期现象;该盯的是梯度和平衡,不是损失正负。 「判别器越强,重建越清晰」:hinge margin 饱和会让判别器损失梯度为 0,但生成器目标对 fake logit 的导数仍为 −1,不能把 D loss=0 当成生成器梯度消失的证据。正确姿势是晚启动、谱归一化、限制判别器更新次数、用自适应权重。 「L1/L2 越低画面越好」:L2 的最优解是条件均值,越低往往越糊;清晰度是用 GAN / 感知项换来的,并伴随 PSNR 下降。要同时看像素指标和感知 / 对抗指标。 「视频 VAE 把图像四项 loss 直接套到 3D 卷积上就行」:单帧损失测不出帧间闪烁(实测单帧 L1 几乎相同、时序误差差 393 倍)。应额外评估时序误差,按消融结果决定是否需要时空判别器、Video-DWT 或显式帧间约束;时间网络结构本身也能学习一致性。 「四项应该从头一起训」:主流做法是重建加 KL(有时加 LPIPS)先预热,第 disc_start 步才开 GAN;H3AE 甚至主张慎用判别损失。晚启动是一种稳定训练的办法,是否必需以及何时启动需依据具体实验。 09. 动手验证 5 个脚本都在本文附录,仅需 numpy(pip install numpy);第 6 个 make_figures.py 额外需要 matplotlib,用来重画本文配图。预期关键结果如下(随机种子已固定): 运行 vaeloss_terms.py:看到 L1 约 0.0316、KL (sum) 约 57.7、KL (mean) 约 0.225,以及「KL (sum)/L1 约 1826 倍」和 HunyuanVideo 权重配方下各项贡献 —— 建立「权重是在配平量纲」的直觉。 运行 pixel_vs_perceptual.py:亮度偏移与同能量噪声的像素 MSE 几乎相等(差小于 1%),但感知距离相差约三个数量级(噪声约为偏移的 $8\times10^3$ 倍)。 运行 gan_schedule.py:warm-up 在 step 2000 起跳;判别器「过强」时 D (hinge)=0.0000 而 G 约 5.97,同时检查脚本新增的 logit 导数:D margin 梯度归零,G 对 fake logit 的导数仍非零。 运行 adaptive_weight.py:先看有限差分与解析梯度完全一致,再看 d_weight 从训练早期 0.095 降到近收敛 0.003(约 32 倍)。 运行 temporal_consistency.py:闪烁与稳定两种重建的单帧 L1 约 0.063 几乎相同,时序 MSE 却差约 393 倍 —— 理解视频为何要单独约束时间维。 运行 make_figures.py(这个需要额外装 pip install matplotlib):重新生成本文三张配图 —— 量级账本(加权前 vs 加权后)、闪烁轨迹与指标对比、自适应权重随收敛下降的曲线。把 vaeloss_terms.py 里的 KL 权重从 $10^{-6}$ 调大两个数量级再看图 1,KL 项由约 $5.8\times10^{-5}$ 升到 $5.8\times10^{-3}$,仍低于这里约 0.0316 的 L1,不能声称已经压过所有项。 建议进一步动手:把 vaeloss_terms.py 里的 KL 从 sum() 改成 mean(),观察 HunyuanVideo 配方里 KL 项的相对权重需要相应放大多少倍,就能切身体会「权重不能跨口径照抄」。 10. 延伸阅读 读这篇之前建议先看已发布的前置: 变分下界与重参数化:KL 解析式与重参数化技巧的完整推导。 VAE 结构与训练目标:编码器 / 解码器、scaling factor、后验坍缩。 视频 VAE 的时空压缩结构:3D 因果卷积、时间 / 空间压缩比与分块解码。 读完这篇可以继续看: 视频生成评测:VBench 与人工验收(还没写):学会在分维度指标上看出「总分没变、时序一致性掉了」这类压缩副作用。 FID / CLIP Score 到底测了什么:感知与分布距离指标的适用边界,与本文 LPIPS/GAN 的质量观互补。 离散化表征:VQ-VAE 与 VQGAN(还没写):本文连续 KL-VAE 的「离散码本」对照版本,codebook 损失替代 KL。 附录:完整代码 09 节用到的脚本全文如下(vaeloss_terms.py、make_figures.py、pixel_vs_perceptual.py、gan_schedule.py、adaptive_weight.py、temporal_consistency.py)。复制到本地存成同名文件,按各脚本开头的依赖说明准备环境后即可运行。 vaeloss_terms.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """视频 VAE 四类损失的最小可运行对照:L1 / KL / 感知 / GAN。 只依赖 numpy,直接 `python vaeloss_terms.py` 复现。 这一步不训练任何网络,只把四项损失放到同一个合成视频张量上算一遍, 回答一个最容易把人带沟里的问题:它们的「数值量级」根本不在一个频道上, 为什么工业配方里 KL 的权重会小到 1e-6。 """ import numpy as np rng = np.random.default_rng(0) # ---------------------------------------------------------------------- # 一个迷你的多尺度「特征提取器」,用来演示感知损失的*结构*。 # 注意:它不是 LPIPS 本体——LPIPS 用的是在 BAPPS 上学过权重的 VGG/Alex # 特征(见工业对照一节的 taming/lpips.py)。这里用可分离高斯 + 差分(DoG) # 构造确定性的多尺度边缘特征,目的是让「在特征空间而非像素空间比较」 # 这件事可以离线、免权重地跑起来。 # ---------------------------------------------------------------------- def _gauss_kernel(size=5, sigma=1.0): ax = np.arange(size) - (size - 1) / 2.0 k = np.exp(-(ax ** 2) / (2 * sigma ** 2)) return k / k.sum() def _conv_separable(img, k): """对最后两轴(H, W)做可分离卷积;img 形状 [..., H, W],边界 reflect。""" pad = len(k) // 2 x = np.pad(img, ((0, 0),) * (img.ndim - 2) + ((pad, pad), (0, 0)), mode="reflect") acc = np.zeros_like(img) for i, w in enumerate(k): acc += w * x[..., i:i + img.shape[-2], :] x = np.pad(acc, ((0, 0),) * (img.ndim - 2) + ((0, 0), (pad, pad)), mode="reflect") acc = np.zeros_like(img) for j, w in enumerate(k): acc += w * x[..., :, j:j + img.shape[-1]] return acc def _normalize_channels(feat, eps=1e-10): # 对应 LPIPS 的 normalize_tensor:沿通道维归一化 norm = np.sqrt(np.sum(feat ** 2, axis=1, keepdims=True)) return feat / (norm + eps) def feature_stack(x): """x: [B, C, T, H, W] -> 多尺度边缘特征列表(先把 T 折叠进 batch)。""" B, C, T, H, W = x.shape f = x.transpose(0, 2, 1, 3, 4).reshape(B * T, C, H, W) k = _gauss_kernel(5, 1.0) b1 = _conv_separable(f, k) b2 = _conv_separable(_conv_separable(b1, k), k) dog1 = f - b1 # 高频细节 dog2 = b1 - b2 # 中频边缘 feats = [_normalize_channels(f), _normalize_channels(dog1), _normalize_channels(dog2)] return feats def perceptual_proxy(x, y): """结构对齐 LPIPS:逐层归一化特征差的平方,空间平均后跨层求和。""" total = 0.0 for fx, fy in zip(feature_stack(x), feature_stack(y)): total += np.mean((fx - fy) ** 2) return total def make_video(B=2, T=8, H=32, W=32): """合成一段有运动内容的视频:移动的亮圆 + 网格背景,取值[0,1]。""" x = np.zeros((B, 3, T, H, W), dtype=np.float64) yy, xx = np.mgrid[0:H, 0:W] for b in range(B): for t in range(T): cx = W * (0.3 + 0.5 * t / (T - 1)) cy = H * (0.3 + 0.15 * np.sin(2 * np.pi * t / T)) circle = ((xx - cx) ** 2 + (yy - cy) ** 2) < (H * 0.12) ** 2 grid = ((xx // 8 + yy // 8) % 2) * 0.15 frame = 0.2 + grid frame[circle] = 0.95 x[b, 0, t] = frame x[b, 1, t] = frame * 0.9 x[b, 2, t] = frame * 0.7 return x def degrade(x): """重建结果:轻微模糊 + 小噪声 + 偏色,模拟一个训练中段的解码器。""" B, C, T, H, W = x.shape flat = x.transpose(0, 2, 1, 3, 4).reshape(B * T, C, H, W) k = _gauss_kernel(3, 0.8) out = _conv_separable(flat, k) out = out + rng.normal(0, 0.02, out.shape) out[:, 0] += 0.03 # 轻微红色偏置 return np.clip(out, 0, 1).reshape(B, T, C, H, W).transpose(0, 2, 1, 3, 4) def main(): x = make_video() r = degrade(x) print("视频张量 x / r :", x.shape, " 取值范围 %.2f~%.2f" % (x.min(), x.max())) # 编码器输出的后验 q(z|x):latent 在时间压缩4x、空间压缩8x B, C, T, H, W = x.shape Cz, Tz, Hz, Wz = 4, T // 4, H // 8, W // 8 mu = 0.3 * rng.standard_normal((B, Cz, Tz, Hz, Wz)) logvar = -1.0 + 0.2 * rng.standard_normal((B, Cz, Tz, Hz, Wz)) n_pix = B * C * T * H * W n_lat = mu.size print("latent 形状 :", mu.shape, " 像素数=%d latent数=%d\n" % (n_pix, n_lat)) # ① 像素重建 L1(对全部元素求平均,量纲与像素一致,O(0.01~0.1)) l1 = np.mean(np.abs(x - r)) # ② KL 解析式(标准正态先验)。注意它天然是「求和」式 kl_per_element = -0.5 * (1 + logvar - mu ** 2 - np.exp(logvar)) kl_sum = kl_per_element.sum() kl_mean = kl_per_element.mean() # ③ 感知损失(多尺度特征距离,量级被归一化压在 O(0.01)) l_p = perceptual_proxy(x, r) # ④ GAN 生成器损失。这里直接喂一组判别器对假图的打分 logits # hinge 生成损失 = -mean(logits_fake);先假设判别器刚起步、打分偏正 logits_fake = rng.normal(0.3, 0.5, size=64) g_loss = -np.mean(logits_fake) print("%-28s %10s" % ("损失项", "原始数值")) print("-" * 40) print("%-28s %10.5f" % ("L1 重建 (mean)", l1)) print("%-28s %10.5f" % ("KL (sum,常见写法)", kl_sum)) print("%-28s %10.5f" % ("KL (mean,换口径)", kl_mean)) print("%-28s %10.5f" % ("感知 proxy", l_p)) print("%-28s %10.5f" % ("GAN g=-mean(D(r))", g_loss)) print("\n--- 为什么 KL 权重能小到 1e-6 ---") # 若直接把 sum 口径的 KL 与 mean 口径的 L1 相加,KL 会凭元素数量碾压: print("KL(sum)/L1 的量级倍数 : %.1fx" % (kl_sum / l1)) print("HunyuanVideo 配方 L1 + 0.1*LPIPS + 0.05*GAN + 1e-6*KL :") total = l1 + 0.1 * l_p + 0.05 * g_loss + 1e-6 * kl_sum print(" 各项贡献: L1=%.5f LPIPS=%.5f GAN=%.5f KL=%.6f" % (l1, 0.1 * l_p, 0.05 * g_loss, 1e-6 * kl_sum)) print(" 合计 = %.5f" % total) print("\n结论:权重不是玄学,是在给「不同口径、不同元素数、不同量纲」的项找平。") if __name__ == "__main__": main() make_figures.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """生成「视频 VAE 的常见 loss 组合」的三张解释图。 数值一律从同目录的三个脚本里取(vaeloss_terms / temporal_consistency / adaptive_weight),这里只负责画——改了那边这里要重跑,免得图和正文数字打架。 三张图分别回答: 1. 四项损失的原始量级差着几个数量级,工业配方加权后为什么能凑到一起 2. 单帧指标完全分不开的两种重建,时间维损失一眼看穿 3. 自适应 GAN 权重为什么随训练自动变小 只依赖 numpy + matplotlib。跑法:python make_figures.py """ import textwrap from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import adaptive_weight as AW import temporal_consistency as TC import vaeloss_terms as VT ROOT = Path(__file__).resolve().parents[1] OUT = ROOT / "figures" OUT.mkdir(parents=True, exist_ok=True) plt.rcParams.update({ "font.sans-serif": ["PingFang SC", "Arial Unicode MS", "DejaVu Sans"], "axes.unicode_minus": False, "figure.dpi": 160, "savefig.bbox": "tight", }) INK = "#1f2937" MUTE = "#6b7280" C_PIX = "#2f6fb0" # 像素项:蓝 C_PERC = "#8b5cf6" # 感知项:紫 C_GAN = "#e0a03c" # 对抗项:橙 C_KL = "#d1495b" # KL:红 C_A = "#d1495b" # 闪烁重建 C_B = "#2f9e6f" # 稳定重建 C_GT = "#1f2937" # 真值 def style(ax, title, xlabel=None, ylabel=None): ax.set_title(title, fontsize=11.5, color=INK, pad=10, loc="left") if xlabel: ax.set_xlabel(xlabel, fontsize=10, color=MUTE) if ylabel: ax.set_ylabel(ylabel, fontsize=10, color=MUTE) ax.tick_params(colors=MUTE, labelsize=9) for s in ("top", "right"): ax.spines[s].set_visible(False) for s in ("left", "bottom"): ax.spines[s].set_color("#d1d5db") ax.grid(alpha=0.25, linewidth=0.6, axis="y") ax.set_axisbelow(True) def footer(fig, text, width=118): """把「这张图要看什么」放到坐标轴下方。 必须放在 y<0 的位置:bbox_inches="tight" 会把负坐标的 artist 一起收进来, 放在 0~0.05 之间的话会和 x 轴标签叠在一起。 """ wrapped = "\n".join(textwrap.wrap(text, width=width)) fig.text(0.012, -0.13, wrapped, fontsize=8.5, color=MUTE, va="top", ha="left", linespacing=1.6) # ────────────────────────────────────────────────────────────────────── # 图 1:四项损失的量级账本 # ────────────────────────────────────────────────────────────────────── def fig_magnitudes(): x, r = VT.make_video(), VT.degrade(VT.make_video()) # 与 vaeloss_terms.main 完全同口径重算一遍(rng 序列一致) B, C, T, H, W = x.shape Cz, Tz, Hz, Wz = 4, T // 4, H // 8, W // 8 mu = 0.3 * VT.rng.standard_normal((B, Cz, Tz, Hz, Wz)) logvar = -1.0 + 0.2 * VT.rng.standard_normal((B, Cz, Tz, Hz, Wz)) l1 = float(np.mean(np.abs(x - r))) kl_el = -0.5 * (1 + logvar - mu ** 2 - np.exp(logvar)) kl_sum, kl_mean = float(kl_el.sum()), float(kl_el.mean()) lp = float(VT.perceptual_proxy(x, r)) gan = abs(float(-np.mean(VT.rng.normal(0.3, 0.5, size=64)))) fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.4)) ax = axes[0] names = ["L1 重建", "感知 proxy", "KL(mean)", "KL(sum)", "|GAN|"] vals = [l1, lp, kl_mean, kl_sum, gan] cols = [C_PIX, C_PERC, C_KL, C_KL, C_GAN] bars = ax.bar(names, vals, color=cols, width=0.62) ax.set_yscale("log") ax.set_ylim(1e-3, 3e2) for b, v in zip(bars, vals): ax.text(b.get_x() + b.get_width() / 2, v * 1.25, f"{v:.4g}", ha="center", fontsize=8.5, color=INK) style(ax, "加权前:同一视频上四项损失的原始数值", None, "损失值(对数轴)") ax = axes[1] wnames = ["1.0 × L1", "0.1 × 感知", "0.05 × GAN", "1e-6 × KL(sum)"] wvals = [l1, 0.1 * lp, 0.05 * gan, 1e-6 * kl_sum] wcols = [C_PIX, C_PERC, C_GAN, C_KL] bars = ax.bar(wnames, wvals, color=wcols, width=0.62) ax.set_yscale("log") ax.set_ylim(1e-6, 1e-1) for b, v in zip(bars, wvals): ax.text(b.get_x() + b.get_width() / 2, v * 1.6, f"{v:.4g}", ha="center", fontsize=8.5, color=INK) ax.axhline(l1, color=MUTE, ls=":", lw=1.0) ax.text(2.6, l1 * 1.5, "L1 的量级", fontsize=8, color=MUTE) style(ax, "加权后:HunyuanVideo 配方下各项的真实贡献", None, "对总损失的贡献(对数轴)") fig.suptitle("图 1 损失归约口径与加权贡献:标量大小不等于梯度大小", fontsize=12, color=INK, x=0.012, ha="left", y=1.02) footer(fig, "要看什么:左图是四项损失在同一个合成视频上的原始数值——KL 按教科书口径对 latent 维" "求和,一上来就是 L1 的 1826 倍;右图是套上工业权重之后各项的真实贡献," "贡献仍不同量级(KL 约5.8e-5,而L1约0.0316);这些权重不能由toy损失比唯一推导。" "KL 权重小到 1e-6 不是不重要,是在补偿「求和口径 vs 平均口径」的元素数之差。") fig.savefig(OUT / "loss_magnitudes.png") plt.close(fig) print(f"[图1] L1={l1:.5f} KL_sum={kl_sum:.5f} 倍数={kl_sum / l1:.0f}x;" f"加权后贡献 L1={l1:.5f} KL={1e-6 * kl_sum:.2e}") # ────────────────────────────────────────────────────────────────────── # 图 2:单帧指标看不见的闪烁 # ────────────────────────────────────────────────────────────────────── def fig_temporal(): x = TC.make_sequence() # [T,1,H,W] fixed = TC.SIGMA * TC.rng.standard_normal(x.shape[1:])[None] rA = x + TC.SIGMA * TC.rng.standard_normal(x.shape) rB = x + fixed + 0.05 * TC.SIGMA * TC.rng.standard_normal(x.shape) # 取一条穿过运动边缘的水平线上的一个像素,看它随帧的取值 t_axis = np.arange(TC.T) y0, x0 = TC.H // 2, 30 gt = x[:, 0, y0, x0] a = rA[:, 0, y0, x0] b = rB[:, 0, y0, x0] fig, axes = plt.subplots(1, 2, figsize=(11.6, 4.3)) ax = axes[0] ax.plot(t_axis, gt, lw=2.6, color=C_GT, label="真值", zorder=3) ax.plot(t_axis, a, lw=1.2, color=C_A, alpha=0.9, label="A 逐帧独立噪声(闪)") ax.plot(t_axis, b, lw=1.2, color=C_B, alpha=0.9, label="B 跨帧恒定退化(稳)") ax.set_ylim(gt.min() - 4 * TC.SIGMA, gt.max() + 4 * TC.SIGMA) ax.legend(fontsize=8.5, frameon=False, loc="upper left") style(ax, "同一个像素随帧的变化:A 在真值附近抖,B 整体平移但不抖", "帧 t", "像素值") ax = axes[1] mets = ["单帧空间 L1", "时序差分 L1", "时序差分 MSE"] va = [TC.spatial_l1(rA, x), TC.temporal_error(rA, x), TC.temporal_mse(rA, x)] vb = [TC.spatial_l1(rB, x), TC.temporal_error(rB, x), TC.temporal_mse(rB, x)] xg = np.arange(len(mets)) w = 0.36 ba = ax.bar(xg - w / 2, va, w, color=C_A, label="A 闪") bb = ax.bar(xg + w / 2, vb, w, color=C_B, label="B 稳") ax.set_yscale("log") ax.set_ylim(1e-5, 1) for bars in (ba, bb): for b_ in bars: ax.text(b_.get_x() + b_.get_width() / 2, b_.get_height() * 1.5, f"{b_.get_height():.4g}", ha="center", fontsize=8, color=INK) ax.set_xticks(xg) ax.set_xticklabels(mets) ax.legend(fontsize=8.5, frameon=False, loc="upper left") style(ax, "三种指标下 A、B 的差距:第一列分不开,第三列差 393 倍", None, "误差(对数轴)") fig.suptitle("图 2 单帧 L1 完全分不开的两种重建,时序差分 MSE 差了 393 倍", fontsize=12, color=INK, x=0.012, ha="left", y=1.02) footer(fig, "要看什么:左图是同一个像素在 12 帧上的取值——A(红)每一帧都在真值附近独立抖动," "连起来播放就是闪烁;B(绿)带一个恒定偏移但帧间几乎不动。" "右图是量化结论:单帧空间 L1 下 A=0.0639、B=0.0634,指标根本分不出谁好谁坏;" "换时序差分 MSE,A 是 B 的 393 倍。只看单帧指标,闪烁是隐形的。") fig.savefig(OUT / "temporal_flicker.png") plt.close(fig) print(f"[图2] 单帧L1 A={va[0]:.5f}/B={vb[0]:.5f};时序MSE A={va[2]:.5f}/B={vb[2]:.5f}" f"({va[2] / max(vb[2], 1e-12):.0f}x)") # ────────────────────────────────────────────────────────────────────── # 图 3:自适应 GAN 权重随训练自动变小 # ────────────────────────────────────────────────────────────────────── def fig_adaptive(): b = np.zeros(AW.D) # 先按 adaptive_weight.main 的抽随机顺序取两个阶段点,保证与正文数字一致 W_early = AW.rng.standard_normal((AW.dz, AW.D)) * 0.05 W_late = AW.W_star + AW.rng.standard_normal((AW.dz, AW.D)) * 0.01 g_rec, g_gan = AW.grads(W_early, b) w_early = AW.adaptive_weight(g_rec, g_gan) g_rec2, g_gan2 = AW.grads(W_late, b) w_late = AW.adaptive_weight(g_rec2, g_gan2) # 扫描曲线用独立的 rng,不扰动上面的阶段点 sweep_rng = np.random.default_rng(7) scales = np.logspace(-4, -0.3, 14) weights = [] for s in scales: W = AW.W_star + s * sweep_rng.standard_normal((AW.dz, AW.D)) g_rec, g_gan = AW.grads(W, b) weights.append(AW.adaptive_weight(g_rec, g_gan)) weights = np.array(weights) fig, ax = plt.subplots(figsize=(7.8, 4.4)) ax.plot(scales, weights, marker="o", ms=4.5, lw=2.0, color=C_GAN, label=r"自适应权重 $d_{\mathrm{weight}}$") ax.axvline(0.05, color=MUTE, ls=":", lw=1.0) ax.text(0.055, w_early * 2.2, "训练早期\n(初始化附近)", fontsize=8.5, color=MUTE) ax.axvline(0.01, color=MUTE, ls=":", lw=1.0) ax.text(0.0105, w_late * 0.06, "训练后期\n(近收敛)", fontsize=8.5, color=MUTE) for s, w_ in [(0.05, w_early), (0.01, w_late)]: ax.plot([s], [w_], marker="s", ms=7, color=C_PIX, zorder=5) ax.annotate(f"{w_early:.3f}", xy=(0.05, w_early), xytext=(0.075, w_early * 1.6), fontsize=9, color=C_PIX) ax.annotate(f"{w_late:.4f}", xy=(0.01, w_late), xytext=(0.0125, w_late * 0.4), fontsize=9, color=C_PIX) ax.set_xscale("log") ax.set_yscale("log") ax.legend(fontsize=9, frameon=False, loc="upper right") style(ax, "重建越接近收敛,GAN 项被自动调得越小", "解码器离最优解的距离(权重扰动幅度,对数轴)", r"自适应权重 $d_{\mathrm{weight}}$(对数轴)") fig.suptitle("图 3 固定 λ 会两头翻车:早期压不住重建、后期压不住 GAN", fontsize=12, color=INK, x=0.012, ha="left", y=1.03) footer(fig, "要看什么:横轴是解码器离最优解有多远(越靠左越接近收敛),纵轴是 VQGAN 的" "自适应权重 ||∇L_rec|| / ||∇L_GAN||。它随残差缩小近似线性下降——重建收敛后" f"从 {w_early:.3f} 掉到 {w_late:.4f}(约 {w_early / max(w_late, 1e-12):.0f} 倍),GAN 项被同步调小。" "若用固定权重,早期 GAN 抢不过巨大的重建梯度,后期重建梯度变小、GAN 又反过来" "压过重建——这条斜线就是在消除这个漂移。") fig.savefig(OUT / "adaptive_weight_curve.png") plt.close(fig) print(f"[图3] d_weight: 早期 {w_early:.4f} -> 后期 {w_late:.4f}" f"({w_early / max(w_late, 1e-12):.0f}x)") def main(): fig_magnitudes() fig_temporal() fig_adaptive() print(f"[OK] 三张图已写入 {OUT}") for f in sorted(OUT.glob("*.png")): print(f" {f.name} {f.stat().st_size / 1024:.0f} KB") if __name__ == "__main__": main() pixel_vs_perceptual.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """复现 LPIPS 论文最核心的观察:相同的像素误差,感知质量可以天差地别。 只依赖 numpy,直接 `python pixel_vs_perceptual.py` 复现。 我们构造三种「重建」,让其中两种的像素 MSE 完全相等: A. 全局亮度偏移 —— 人眼对缓慢的整体明暗变化相当不敏感; B. 同等 L2 能量的逐点噪声 —— 直接污染纹理与细节,观感很差; C. 高斯模糊 —— 抹掉高频细节。 然后分别在「像素空间」和一个多尺度特征空间里比较距离。 """ import numpy as np rng = np.random.default_rng(1) def _gauss_kernel(size=5, sigma=1.0): ax = np.arange(size) - (size - 1) / 2.0 k = np.exp(-(ax ** 2) / (2 * sigma ** 2)) return k / k.sum() def _conv_separable(img, k): pad = len(k) // 2 x = np.pad(img, ((0, 0),) * (img.ndim - 2) + ((pad, pad), (0, 0)), mode="reflect") acc = np.zeros_like(img) for i, w in enumerate(k): acc += w * x[..., i:i + img.shape[-2], :] x = np.pad(acc, ((0, 0),) * (img.ndim - 2) + ((0, 0), (pad, pad)), mode="reflect") acc = np.zeros_like(img) for j, w in enumerate(k): acc += w * x[..., :, j:j + img.shape[-1]] return acc def _norm(feat, eps=1e-10): return feat / (np.sqrt(np.sum(feat ** 2, axis=1, keepdims=True)) + eps) def perceptual_proxy(x, y): """多尺度高通特征上的归一化距离(LPIPS 的结构代理,非本体)。""" def feats(z): k = _gauss_kernel(5, 1.0) b1 = _conv_separable(z, k) b2 = _conv_separable(_conv_separable(b1, k), k) return [_norm(z), _norm(z - b1), _norm(b1 - b2)] return sum(np.mean((a - b) ** 2) for a, b in zip(feats(x), feats(y))) def make_textured_image(H=64, W=64): """一张同时含锐边、细密纹理和平滑区域的图。""" yy, xx = np.mgrid[0:H, 0:W].astype(np.float64) img = 0.5 + 0.25 * np.sin(xx / 2.0) * np.sin(yy / 2.0) # 细密纹理 img += 0.3 * (xx > W / 2) # 一条锐边 checker = ((xx // 4 + yy // 4) % 2) * 0.15 # 棋盘格 img += checker img = np.clip(img, 0, 1) return np.stack([img, img * 0.92, img * 0.8], axis=0)[None] # [1,C,H,W] def mse(a, b): return float(np.mean((a - b) ** 2)) def main(): x = make_textured_image() delta = 0.12 # 统一的 L2 误差能量 rA = x + delta # A:全局亮度偏移(不 clip,保证误差严格等于 delta) rB = x + rng.normal(0, delta, x.shape) # B:零均值逐点噪声,标准差=delta rC = _conv_separable(x, _gauss_kernel(7, 1.6)) # C:模糊 rows = [ ("A 亮度偏移(整体+%.2f)" % delta, rA), ("B 逐点噪声(σ=%.2f)" % delta, rB), ("C 高斯模糊", rC), ] print("%-26s %12s %12s %14s" % ("重建方式", "像素MSE", "感知距离", "感知/像素 比")) print("-" * 68) for name, r in rows: pm = mse(x, r) pp = perceptual_proxy(x, r) print("%-26s %12.5f %12.5f %14.2f" % (name, pm, pp, pp / (pm + 1e-12))) mseA, mseB = mse(x, rA), mse(x, rB) pA, pB = perceptual_proxy(x, rA), perceptual_proxy(x, rB) print("\n关键对照:A 与 B 的像素 MSE 相差仅 %.2f%%," % (100 * abs(mseA - mseB) / mseA)) print("但特征空间里 B(噪声)的感知距离是 A(亮度偏移)的 %.1f 倍。" % (pB / pA)) print("高通/多尺度特征天然滤掉直流亮度、放大纹理污染——这正是 LPIPS 比 MSE 更贴人眼的原因。") if __name__ == "__main__": main() gan_schedule.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """GAN 两项的最小演示:hinge / vanilla 损失,以及判别器的 warm-up 调度。 只依赖 numpy,直接 `python gan_schedule.py` 复现。 视频 VAE 里对抗损失通常不是从头开:先用 L1+KL+LPIPS 把解码器训到大致成形, 第 disc_start 步才把判别器权重从 0 抬起来(taming 里的 adopt_weight)。 本脚本同时给出两种常见判别损失的数值,并展示「判别器过强」时的失衡信号。 """ import numpy as np rng = np.random.default_rng(3) def adopt_weight(weight, global_step, threshold=0, value=0.0): """taming-transformers 里的原逻辑:threshold 之前强制为 value(通常是0)。""" return value if global_step < threshold else weight def hinge_d(real, fake): return 0.5 * (np.mean(np.maximum(0.0, 1.0 - real)) + np.mean(np.maximum(0.0, 1.0 + fake))) def hinge_g(fake): return float(-np.mean(fake)) def vanilla_d(real, fake): softplus = lambda z: np.log1p(np.exp(-np.abs(z))) + np.maximum(z, 0.0) return 0.5 * (np.mean(softplus(-real)) + np.mean(softplus(fake))) def evaluate(real, fake, tag): print("%-28s D(hinge)=%.4f D(vanilla)=%.4f G(hinge)=%.4f 均值打分 real=%.2f fake=%.2f" % (tag, hinge_d(real, fake), vanilla_d(real, fake), hinge_g(fake), real.mean(), fake.mean())) def main(): # ① warm-up 调度:disc_start=2000 print("== adopt_weight 调度(disc_start=2000, 目标权重 0.05) ==") for step in (0, 1999, 2000, 5000): w = adopt_weight(0.05, step, threshold=2000) print(" step=%5d -> 对抗权重 = %.3f" % (step, w)) # ② 判别器在三种强弱下的损失 print("\n== 判别器强弱与损失信号 ==") # 平衡初期:真假打分都在 0 附近 real0 = rng.normal(0.0, 0.3, 256) fake0 = rng.normal(0.0, 0.3, 256) evaluate(real0, fake0, "起步:判别器分不清") # 健康:真≈+1,假≈-1,仍留梯度 real1 = rng.normal(1.2, 0.4, 256) fake1 = rng.normal(-1.2, 0.4, 256) evaluate(real1, fake1, "健康:适度拉开") # 过强:真≈+6,假≈-6,判别器 hinge 梯度为 0,但生成器对 fake logit 仍有梯度 real2 = rng.normal(6.0, 0.5, 256) fake2 = rng.normal(-6.0, 0.5, 256) evaluate(real2, fake2, "过强:margin 已饱和") d_fake_grad = 0.5 * (fake2 > -1.0) / fake2.size g_fake_grad = -np.ones_like(fake2) / fake2.size print("\nD 对 fake logit 梯度范数 = %.6f" % np.linalg.norm(d_fake_grad)) print("G 对 fake logit 梯度范数 = %.6f" % np.linalg.norm(g_fake_grad)) print("D margin 饱和不表示 G 梯度消失;参数梯度还取决于判别器输入雅可比。") if __name__ == "__main__": main() adaptive_weight.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """复现 VQGAN 的自适应 GAN 权重 calculate_adaptive_weight。 只依赖 numpy,直接 `python adaptive_weight.py` 复现。 固定 GAN 权重的痛点:训练早期解码器很烂,重建梯度很大;训练后期重建收敛、 梯度变小。若 λ 固定,GAN 梯度在后期会相对越来越强,甚至把重建带偏。 VQGAN 的做法是让两项在「最后一层」上的梯度范数之比来决定权重: d_weight = ||∇_last L_rec|| / (||∇_last L_gan|| + 1e-4),再 clamp 到 [0,1e4] 我们用一个可解析求导的迷你线性解码器,精确算出两个梯度范数, 并用有限差分抽查一个分量,证明数值不是凑出来的。 """ import numpy as np rng = np.random.default_rng(2) N, D, dz = 16, 24, 8 # 样本数、像素维、latent 维 Z = rng.standard_normal((N, dz)) W_star = rng.standard_normal((dz, D)) * 0.3 X = Z @ W_star # 真实数据:线性可完美拟合 # 一个固定的微型判别器打分函数 s(x)=x @ a(教学用,参数冻结) a = rng.standard_normal((D,)) def decode(W, b): return Z @ W + b def rec_loss(W, b): """重建项:L2(L1 同理,范数比机制不变)。""" return float(np.mean((decode(W, b) - X) ** 2)) def gan_g_loss(W, b): """生成器 hinge 损失 -mean(D(r))。""" s = decode(W, b) @ a return float(-np.mean(s)) def grads(W, b): R = decode(W, b) - X g_rec = (2.0 / (N * D)) * Z.T @ R # ∂L_rec/∂W(mean 对 N*D) # g=-mean_n(s_n),s_n=Σ_d xhat_nd·a_d:mean 只对 N,分母是 N,不是 N*D s_coef = -(1.0 / N) * Z.sum(axis=0) # ∂g/∂W_kd = -(1/N)(Σ_n z_nk) a_d g_gan = np.outer(s_coef, a) return g_rec, g_gan def adaptive_weight(g_rec, g_gan, cap=1e4): w = np.linalg.norm(g_rec) / (np.linalg.norm(g_gan) + 1e-4) return float(np.clip(w, 0.0, cap)) def finite_diff_check(W, b, eps=1e-6): """用中心差分抽查 W[0,0] 上的两个梯度,验证解析解。""" out = [] for fn in (rec_loss, gan_g_loss): Wp, Wm = W.copy(), W.copy() Wp[0, 0] += eps Wm[0, 0] -= eps out.append((fn(Wp, b) - fn(Wm, b)) / (2 * eps)) return out def stage(W, b, name): lrec = rec_loss(W, b) lgan = gan_g_loss(W, b) g_rec, g_gan = grads(W, b) w = adaptive_weight(g_rec, g_gan) print("%-22s L_rec=%8.5f L_gan=%8.4f ||∇rec||=%8.4f ||∇gan||=%7.4f d_weight=%7.4f" % (name, lrec, lgan, np.linalg.norm(g_rec), np.linalg.norm(g_gan), w)) return w def main(): b = np.zeros(D) # 阶段一:解码器刚初始化,离最优很远(重建残差大) W_early = rng.standard_normal((dz, D)) * 0.05 # 阶段二:接近收敛(在最优解上加很小扰动,残差小) W_late = W_star + rng.standard_normal((dz, D)) * 0.01 print("== 有限差分校验(W[0,0]) ==") fd_rec, fd_gan = finite_diff_check(W_early, b) g_rec, g_gan = grads(W_early, b) print("解析 ∂Lrec/∂W00=%.6f 差分=%.6f" % (g_rec[0, 0], fd_rec)) print("解析 ∂Lgan/∂W00=%.6f 差分=%.6f\n" % (g_gan[0, 0], fd_gan)) print("== 自适应权重随训练阶段变化 ==") w_early = stage(W_early, b, "训练早期(未收敛)") w_late = stage(W_late, b, "训练后期(近收敛)") print("\n重建收敛后 d_weight 从 %.3f 降到 %.4f(约 %.0f 倍),GAN 项被自动调小。" % (w_early, w_late, w_early / max(w_late, 1e-12))) print("若用固定 λ:早期 GAN 抢不过重建、后期 GAN 又可能压过重建——自适应权重就是在消除这个漂移。") if __name__ == "__main__": main() temporal_consistency.py #!/usr/bin/env python3 # -*- coding: utf-8 -*- """视频 VAE 为什么必须额外盯「时间维」:同样的单帧画质,闪烁程度可以天差地别。 只依赖 numpy,直接 `python temporal_consistency.py` 复现。 构造一段运动视频,再给两种重建: A. 逐帧独立噪声 —— 每一帧单独看误差不大,连起来播放却疯狂闪烁; B. 跨帧恒定的退化(同一固定纹理噪声)—— 单帧误差与 A 相同,但不闪。 单帧空间 L1 完全分不清 A、B;相邻帧差分的时序一致性损失一眼区分。 """ import numpy as np rng = np.random.default_rng(4) T, H, W = 12, 48, 64 SIGMA = 0.08 def make_sequence(): yy, xx = np.mgrid[0:H, 0:W].astype(np.float64) x = np.zeros((T, H, W)) for t in range(T): cx = 8 + 4 * t # 一条匀速移动的竖边 frame = 0.25 + 0.5 * (xx >= cx) frame += 0.1 * ((xx // 6 + yy // 6) % 2) x[t] = frame return np.clip(x, 0, 1)[:, None] # [T,1,H,W] def spatial_l1(r, x): return float(np.mean(np.abs(r - x))) def frame_diff(v): return v[1:] - v[:-1] def temporal_error(r, x): """相邻帧差分一致性:||Δr - Δx||(静止背景上 GT 帧差为0,闪烁直接显现)。""" return float(np.mean(np.abs(frame_diff(r) - frame_diff(x)))) def temporal_mse(r, x): return float(np.mean((frame_diff(r) - frame_diff(x)) ** 2)) def main(): x = make_sequence() fixed_noise = SIGMA * rng.standard_normal(x.shape[1:])[None] # [1,1,H,W] 跨帧恒定 rA = x + SIGMA * rng.standard_normal(x.shape) # 逐帧独立噪声 -> 闪烁 # 稳定退化:以跨帧恒定噪声为主,只掺 5% 的逐帧抖动(更贴近真实解码器) rB = x + fixed_noise + 0.05 * SIGMA * rng.standard_normal(x.shape) print("退化能量相同(σ=%.2f),比较单帧空间误差与时间维误差:\n" % SIGMA) print("%-22s %14s %16s %16s" % ("重建", "单帧空间L1", "时序差分L1", "时序差分MSE")) print("-" * 72) for name, r in (("A 逐帧独立噪声(闪)", rA), ("B 跨帧恒定退化(稳)", rB)): print("%-22s %14.5f %16.5f %16.5f" % (name, spatial_l1(r, x), temporal_error(r, x), temporal_mse(r, x))) print("\n单帧空间 L1 几乎相等(A=%.4f vs B=%.4f),但时序 MSE 上 A 约为 B 的 %.0f 倍。" % (spatial_l1(rA, x), spatial_l1(rB, x), temporal_mse(rA, x) / max(temporal_mse(rB, x), 1e-12))) print("这解释了视频 VAE 为何要在 2D 的 L1/LPIPS/GAN 之外:") print(" · 把判别器扩到时空(3D/PatchGAN、LTX 的 Video-DWT);") print(" · 或加相邻帧/warp 一致性损失,专门惩罚帧间抖动与闪烁。") if __name__ == "__main__": main()
2026年09月20日
8 阅读
0 评论
1 点赞
2026-09-18
AIGC 每日速读|2026-09-18|高通揪出长视频KV错配-Recency Forcing
今日 AIGC 论文速览 今日共 10 篇 · 视频生成与长时程一致性 2 篇 · 生成模型评测与审计 2 篇 · 推理与采样加速 2 篇 · 语音与音乐生成 2 篇 · 视频修复与图像应用 2 篇 重点论文标题列表 Recency Forcing(Qualcomm AI Research):按距离衰减注意力救长视频 MiniMax-H3-Reason(National University of Singapore;Fudan University;Tencent):全模态模型物理推理仅四成 ⚡ Diagonal Sparsity(KAIST;Agency for Defense Development;Seoul National University):对角稀疏把吞吐拉到3.1倍 VibeAvatar(Peking University):口型与美感分阶段各管各的 AVR(New York University;University of Technology Sydney):能抄就抄抄不到才让扩散补 今日论文速览 1. Recency Forcing:按距离衰减注意力救长视频 Recency Forcing: Bridging the Long-Horizon Gap in Autoregressive Video Generation | Qualcomm AI Research | arXiv:2609.19729 关键词:自回归视频生成,长时程一致性,KV cache,注意力偏置 ⚠️ 前序问题:自回归视频生成一拉长就退化。论文把根因定位到一处被忽视的训练-推理不一致:训练时短片段的全部上下文帧都在 KV cache 里,推理时显存约束却会把远处帧逐出 cache,等于抽走了模型当初赖以生成的上下文。作者称之为 KV 淘汰错配(KV eviction mismatch)。 本文贡献:思路不是用截断上下文去模拟淘汰——那会连模型仍需要的时间信息一起丢掉、破坏运动连贯——而是保留上下文,同时让远处帧的影响逐步衰减,等它们真被淘汰时影响已可忽略。论文先用扰动式敏感度指标 positional response R(Δ, t_denoise) 量出上下文影响随时间距离陡峭衰减、且在不同去噪步上系统性变化;据此提出 Temporal Response Bias(TRB),把一个非正的、随时间步变化的偏置直接加在 softmax 之前的注意力 logits 上。配套的 Biased Attention Reparameterization(BAR)是个等价改写,把偏置移到 softmax 之外,使 TRB 退化成一次标准 FlashAttention 调用。 实验效果:在 VBench 与 VBench-Long 上取得长时程生成的 state-of-the-art 质量,且不增加推理成本。60 秒视频的定性对比里,Self-Forcing 出现色彩漂移,LongLive 产生场景重复,DeepForcing 引入模糊,Recency Forcing 则保持了画面稳定与运动连贯。方法不改动上下文长度与训练目标,免训练与训练两种模式都能用,BAR 的等价改写让偏置的额外开销为零。 批判点评:把长视频退化归到 KV cache 淘汰这个具体的工程约束上,而不是笼统归因于「误差累积」,定位比多数同类工作清晰,BAR 让它零开销落地也务实。但摘要只给 state-of-the-art 的结论、未披露具体分数,免训练与训练两种模式的增益差距、TRB 偏置形状换骨干后是否还成立,都要看正文;VBench 系列本身偏重视觉质量,长时程的叙事与物体持久性未必能被它完全反映。 2. MiniMax-H3-Reason:全模态模型物理推理仅四成 Can MiniMax-H3 Reason About the Physical World? An Evaluation of Omni-Modal Generative Model | National University of Singapore;Fudan University;Tencent | arXiv:2609.18323 关键词:全模态生成模型,物理世界推理,跨模态整合,评测基准 ⚠️ 前序问题:全模态生成模型把文本、图像、视频、音频统一建模,MiniMax-H3 正是在共享潜空间里同时做多模态理解与音视频联合生成的代表。但多模态对齐到底有没有提升模型的世界推理能力?现有面向视频生成与世界模型的评测大多受限于单一输入模态,且提示词与目标视频内容高度重合,回答不了这个问题。 本文贡献:论文构建了围绕物理世界推理四个互补维度的评测框架,专门为全模态输入设计任务:隐式提示搭配多帧、音频-图像、前缀视频、音频-视频四类场景。关键设计是每一路模态只提供事件的部分证据,模型必须跨模态整合互补线索,才能推断出潜在的事件状态与未来动态。 实验效果:517 个评测实例上,MiniMax-H3 总体成功率 41.97%。分维度看,基于视频的决策推理最高、为 56.00%,基于音频的消歧推理最弱、只有 27.40%。作者据此认为要真正吃到多模态输入的红利,关键仍在有效的跨模态整合。 批判点评:用「每种模态只给部分证据」的设定逼出跨模态整合能力,比提示词与目标高度重合的常规评测更有区分度,音频维度 27.40% 也确实点出了当前全模态模型的短板。但全文只评了 MiniMax-H3 一个模型,没有同代横向对比,41.97% 这个绝对值究竟是模型不行还是任务过难无从判断;517 个实例摊到四个维度上,每个维度的样本量也偏薄。 3. Diagonal Sparsity:对角稀疏把吞吐拉到3.1倍 Understanding and Exploiting Diagonal Attention Sparsity in Autoregressive Image Generation | KAIST;Agency for Defense Development;Seoul National University | arXiv:2609.19702 关键词:自回归图像生成,稀疏注意力,KV cache,推理加速 ⚠️ 前序问题:自回归图像生成因为能直接复用基于 transformer 的 LLM 服务设施而成为多模态系统的一种范式,但每次请求要吐出上千个视觉 token,解码越来越被注意力计算中的 KV cache 访问卡住。稀疏注意力在文本 LLM 推理上已被充分研究,可它的稀疏性假设能不能迁移到自回归图像生成,此前没人说得清。 本文贡献:论文首次系统刻画了自回归图像生成中的注意力稀疏性,覆盖多种负载与代表性开源模型,发现三个和文本场景不一样的性质:prefill 与 decode 阶段存在明显不对称、注意力强烈集中在 prompt token 与局部 token 上、以及视觉 token 的空间局部性带来一种独特的对角稀疏模式。据此提出对角感知稀疏注意力,在最近窗口内沿对角方向选择性跳过 KV 条目。 实验效果:基于 FlexGen、FlashAttention-2 与自定义 kernel 实现在 GPU 服务系统上,相比稠密推理最高取得 3.1 倍吞吐与 1.19 倍延迟改善,质量损失低于 2%。 批判点评:把稀疏模式追到「视觉 token 的空间局部性」这一来源、而不是照搬文本 LLM 的稀疏假设,是这篇最扎实的地方,3.1 倍吞吐也有工程说服力。但「质量损失低于 2%」用的是哪个指标、在哪些负载上成立,摘要没交代;对角模式依赖光栅扫描式的 token 排布,换扫描顺序或改分辨率后是否还成立同样存疑。 4. VibeAvatar:口型与美感分阶段各管各的 VibeAvatar: Aligning Phonetic Kinematics and Human Aesthetics for High-Fidelity Talking Avatar Synthesis | Peking University | arXiv:2609.18632 关键词:说话人视频合成,语音驱动,GRPO,流匹配 ⚠️ 前序问题:多模态说话人视频合成要从一张参考肖像加一段语音生成逼真视频。扩散类方法进步很快,却始终难以同时拿到准确的唇部发音、符合人类偏好的运动美感和高效推理这三样。 本文贡献:论文的判断是:发音准确性和运动美感来自根本不同的源头,应该在互补的阶段分别处理,而不是塞给单个生成器隐式地一起学。VibeAvatar 因此把两个目标解耦——条件阶段用 Phonetic Kinematics Adapter(PKA)把面向识别的语音特征转成发音-运动条件;后训练阶段用 Aesthetic Motion Policy(AMP),以 GRPO 优化一个流一致的随机采样策略。运动生成器本身是轻量流模型,工作在紧凑的一维 warp 潜空间里。 实验效果:客观指标与用户研究上,发音、美感、效率三项均达到 state-of-the-art;生成 10 秒 512px 视频耗时不到 10 秒,显存占用约 3GB。 批判点评:「发音归条件阶段、美感归后训练」是个干净的分工,把 GRPO 用在采样策略而不是生成器权重上也省掉了重训成本,约 3GB 显存意味着消费级显卡能跑。但美感由 GRPO 所对齐的偏好定义,口味会不会过拟合到特定标注群体值得警惕;摘要只给了效率数字,发音与美感的具体分数、以及更长时长下的身份漂移都得看正文。 5. AVR:能抄就抄抄不到才让扩散补 Copy What Is Seen, Generate What Is Not: Training-Free Anomaly-Aware Video Restoration | New York University;University of Technology Sydney | arXiv:2609.18836 关键词:视频修复,异常检测,免训练方法,扩散修补 ⚠️ 前序问题:监控系统检测到异常后往往还得把画面修好,但这两件事一直被分开研究:免训练的异常检测止步于一个分数或标签,免训练的视频编辑则只听用户提示词、不听检测器。 本文贡献:AVR 只用冻结的预训练模型就把两端接上,而且只在画面确实没有证据可抄的地方才生成内容。先由运动证据把开放词表提案收敛成时空掩码;再用片段自身算出的背景先验,填掉异常曾经遮挡过的每一个像素,只把「任何一帧都没出现过」的部分交给扩散合成;最后一个冻结的验证器逐片段决定该信任经典修复、先验锚定修复还是背景条件修复。 实验效果:三个监控数据集上、覆盖全参考异常注入与真实异常两种设定:oracle 掩码下 AVR 的全帧保真度领先,在编辑区域内与三个训练过的视频 inpainting 方法持平,用自己产出的掩码时优于「先检测再生成」的流水线,同时压住了残留异常与自由扩散带来的闪烁。 批判点评:「能抄就抄、抄不到才生成」的取舍很贴监控取证的需求——少生成一分就少一分臆造风险,全程冻结模型也免了训练成本。但它靠运动证据圈定异常,静止或缓慢的异常可能整个漏掉;oracle 掩码下的领先说明上限不错,自产掩码下的成绩才是实际水平,三个监控数据集的场景多样性也仍有限。 6. CPR:作曲演奏精修三段渲染钢琴 CPR: Combining global composing, local performing and full-sequence refining in piano rendering with continuous autoregressive modelling | The Chinese University of Hong Kong, Shenzhen | arXiv:2609.18216 关键词:钢琴渲染,连续自回归,流匹配,音乐生成 ⚠️ 前序问题:提示词条件下的钢琴 MIDI-to-Music 渲染,要在忠实还原目标音符的同时复现参考录音的音色。现有方法分两条路:离散 codec 的自回归模型有因果时序建模能力,但量化会丢掉声学细节;流匹配更能保住声学结构,代价是全序列注意力开销大、语义结构反而更差。 本文贡献:连续自回归模型直接在连续表示上工作,既兼有自回归的条件跟随能力与流匹配的分布建模能力,又绕开量化瓶颈、计算成本更低。CPR 据此拆成三段:Composer 自回归预测连续隐状态,Performer 经局部流匹配生成 24kHz 声学潜变量,Refiner 再把波形上采样到 48kHz。论文另引入 Bottlenecked Representation Alignment(BREPA)与 Modality-Time RoPE(MT-RoPE),分别强化 Composer 隐状态里的音乐语义结构和跨模态时间对齐。 实验效果:P-MUSE-eval 上,CPR 只用 4 步(NFE=4)就输出 48kHz 音频,onset F1 78.4 与 FAD_clap 0.129 均为最佳;音色相似度 89.3 略低于需要 25 步的 P-MUSE(90.6),但明显优于 MIDI-VALLE 的 83.4。另有 13 位听众对 14 个样本的 MOS 主观评测。 批判点评:三段分工把「写什么音」「怎么弹出来」「采样率补齐」拆开,绕过离散量化这个老瓶颈,4 步出 48kHz 相对 P-MUSE 的 25 步是实打实的效率优势。但音色相似度并未超过 P-MUSE,说明连续表示在音色保真上还没占到便宜;MOS 只有 13 位听众评 14 个样本,样本量偏小,评测也集中在钢琴单一乐器上。 7. Self-Distilled TTS:模型自己教自己念生词 Self-Distilled Pronunciation and Accent Control for Neural Text-to-Speech | Independent Researcher;KOWRO Inc. | arXiv:2609.17234 关键词:语音合成,自蒸馏,发音控制,口音标注 ⚠️ 前序问题:直接读原始文本的 TTS 没有词典,碰上生僻词只能靠猜。已有的补救要么拿录音训练一条专门的发音与口音通道,要么从示例里一个词一个词地改。 本文贡献:这篇两样都不做。让冻结的骨干去读一个含有它本来就念得对的常见词的句子,再把它自己的输出当作同一句话的教师——只是把那个词替换成带标注的、指定口音的读法。这一对训练样本就是方法的全部。 实验效果:Sarashina2.2-TTS 上,经筛选的评测者(Fleiss' kappa = 0.85)在 0.89 的未见词上听出了指定口音,而无法表达口音的假名方案只有 0.57、没有赢下任何一组对比,自然度没有可测量的损害。未经调整迁移到自回归、扩散和编码器-解码器骨干后,319 个词上的读法正确率比不做编辑高出 0.25 到 0.47;CosyVoice 2 上三个词里有两个能带上口音,但在 Irodori 上不行,论文分析了原因。 批判点评:用模型自己的输出当教师、只换一个词构成训练对,简洁到几乎没有额外数据成本,跨骨干迁移也说明思路不绑架构。但它天然受限于「骨干本来就念得对的常见词」这个起点,词表外的音素组合未必覆盖得到;Irodori 上失效说明迁移并非无条件成立,日语音高重音之外的语言是否同样奏效也还没有证据。 8. JewelTry:免掩码珠宝试戴还管尺寸 JewelTry: Mask-Free Scale Aware Jewelry Virtual Try-On | Amazon | arXiv:2609.16626 关键词:虚拟试戴,免掩码扩散,尺寸感知,评测基准 ⚠️ 前序问题:虚拟试穿让顾客看到上身效果,已是网购的重要技术。服装类进展很大,珠宝却因为体积小、结构刚性、对细节高度敏感而一直是难啃的品类:既要忠实迁移外观,还得相对佩戴者有正确的尺寸和位置。现有珠宝试戴多依赖掩码引导,而免掩码方法又缺乏建模产品尺寸的显式引导。 本文贡献:论文先建了 JVTO-Bench,提供带真实产品尺寸标注的「参考-源-目标」三元组,覆盖四大珠宝品类。在此之上提出免掩码扩散框架 JewelTry:尺寸适配器把产品实际尺寸编码成一个 scale token,让模型在上下文中学到珠宝与人体解剖结构之间的尺寸关系;再用单向条件注意力与注意力精修损失,同时保住参考珠宝的粗粒度几何与细粒度结构细节。 实验效果:实验显示 JewelTry 在视觉保真度、背景保持、物体一致性与尺寸准确性之间取得平衡,为免掩码、尺寸感知的珠宝虚拟试戴立了一个较强基线。 批判点评:把「尺寸对不对」变成可编码的条件而不是靠掩码兜底,抓住了珠宝区别于服装的真正难点,配套放出带真实尺寸标注的 benchmark 也补上了这一方向的评测空白。但摘要通篇没给量化指标,「取得平衡」这类表述无法判断相对现有方法的实际差距;四个品类对珠宝这种长尾品类而言覆盖面也偏窄。 9. MSR:多张参考图各占一个 slot MSR: Multiple Subject Reference for Video Generation | Licon Studio | arXiv:2609.18393 关键词:多主体视频生成,参考图条件,LoRA,流匹配 ⚠️ 前序问题:让视频生成器同时吃多张参考图时,既要保住每个主体的外观,又要把每张图和它该扮演的角色对上,否则属性就会串到别的对象身上。 本文贡献:MSR 是面向 LTX 视频生成的 slot 感知条件方案:每张参考图被独立编码成一段静态片段、对应一组单独的潜 token;一个紧凑的傅里叶特征 MLP 给它加上数值化的 slot 嵌入,再用随 slot 变化的时间偏移改写这组 token 的旋转位置坐标。这些参考组被前置到带噪的目标 token 之前,在仅针对目标的流匹配训练中充当干净上下文。整套方案用 LoRA 实现,权重与推理工作流已开源。 实验效果:论文给的是定性结果:示例展示了包含不同角色与参考环境的写实及风格化场景组合;开发过程中的观察显示,相对更早的连续参考基线,参考混淆有所减少,但相似服装、复杂衣物和视角变化仍然困难。另有一个补充实验在冻结视觉参数的前提下加入了语音参考条件。 批判点评:用独立 token 组加 slot 嵌入来绑定「哪张图对应哪个角色」,思路直接,靠 LoRA 就能低成本落地,开源权重与工作流对想复现的人友好。但这更像一份技术报告而非完整论文:全文没有定量实验,「参考混淆减少」只是开发观察而非受控对比,缺少与基线的可比指标;作者自己也承认相似服装与视角变化仍未解决。 10. Newer Is Not Fairer:新版文生图并没有更公平 Newer Is Not Fairer: Gender Stereotyping in Text-to-Image AI Across Model Generations | Northeastern University | arXiv:2609.18007 关键词:文生图,性别偏见,模型审计,公平性评测 ⚠️ 前序问题:文生图模型已经在专业与创意场景里广泛使用,但它们如何呈现不同职业的性别、以及更新的模型是不是更公平,跨代际地看仍缺乏系统证据。 本文贡献:论文评测了 20 个职业、5 种提示词模板、4 代 Stable Diffusion(SD 1.5、SD 2.1、SDXL、SD 3 Medium),每个「职业×模型」单元生成 100 张、共 8000 张图,全部用 DeepFace 判定性别,并与美国劳工统计局(BLS)的真实劳动力数据对照,所有显著性检验都做了 Benjamini-Hochberg 多重校正。 实验效果:8000 张开源模型图像中 76.4% 为男性(95% CI [75.1%, 78.7%]);更扎眼的是历史上女性主导的职业里仍有 57.6% 生成男性。相对 BLS 数据,模型平均低估女性 20–46 个百分点,本就接近性别均衡的职业偏差最大——科学家 BLS 女性占 48%、模型输出 82–99% 为男性,清洁工 BLS 女性占 46%、模型输出 80–92% 为男性。代际并非稳步改善:偏差从 SD 1.5 一路恶化到 SDXL,到 SD 3 Medium 才部分回落。与 GPT-image-1 在五个职业上的初步对比显示商业模型偏差更低,但实际效应很小(Cramér's V = 0.080)且属探索性。没有任何一个模型达到性别均衡。 批判点评:8000 张图、多重检验校正、对照 BLS 真实就业数据,这套设计让结论比常见的偏见讨论扎实得多,「更新不等于更公平」的代际发现尤其有价值。但性别判定依赖 DeepFace 的二元分类,本身有误差且无法涵盖非二元表达;20 个职业、5 个模板的提示词空间偏窄,与 GPT-image-1 的对比作者也标注为探索性,撑不起开源与商业模型孰优孰劣的强结论。 趋势观察 长上下文的瓶颈,从「记不记得住」变成「cache 放不放得下」 Recency Forcing 把长视频退化追到 KV cache 淘汰这一工程约束,Diagonal Sparsity 直接冲着自回归图像生成的 KV 访问瓶颈去。两篇都不再泛泛谈误差累积,而是把注意力的显存与访存代价当成一等问题来解。 生成质量的定义权,正在交给后训练与采样策略 VibeAvatar 把运动美感交给 GRPO 在后训练阶段对齐,而不是让生成器隐式地学;CPR 用连续自回归绕开量化,再分三段把语义结构、局部声学与采样率分工。共同前提是承认单个生成器学不好多个互相冲突的目标。 评测开始审问模型「到底懂不懂」,而不只是「像不像」 MiniMax-H3 的评测让每种模态只给部分证据,逼模型跨模态整合才能答对;Newer Is Not Fairer 用 8000 张图对照真实就业数据,把「新版是否更公平」变成可证伪的问题。两者都不满足于保真度分数。 免训练与轻量适配,正在成为落地的默认路径 AVR 全程冻结预训练模型、只在没有证据可抄的地方才生成;MSR 用 LoRA 给现成视频模型加多主体条件;Self-Distilled TTS 让模型自己的输出当教师。共同动机都是避开重训成本。 人工智能炼丹君 整理 | 2026-09-18 更多 AIGC 论文解读,关注微信公众号「人工智能炼丹君」 每日更新 · 论文精选 · 深度解读 · 技术脉络 微信搜索 人工智能炼丹君 或扫描下方二维码关注
2026年09月18日
16 阅读
0 评论
0 点赞
1
...
4
5
6
...
19
粤ICP备2021042327号