APEX:面向极致压缩的异步前缀去噪 APEX: Asynchronous Prefix Denoising for Extreme Compression
把层级通道前缀放到全局时钟上错峰去噪,让 64×/128× 极端压缩成为扩散模型的算力优势
前置知识
潜在扩散与空间压缩率(f8/f64/f128)
潜在扩散模型(LDM)先用自编码器把图像压缩到小网格潜变量,再在潜空间做扩散。空间压缩率 f 表示边长缩小的倍数:$512^2$ 图像在 f8 下得到 $64\times64$ 网格即 1024 个 token,f64 下是 64 个 token,f128 下只剩 16 个。Transformer 去噪器的注意力开销随 token 数平方增长,所以 token 越少训练与采样越便宜。
本文的全部动机就是让 f64/f128 这种极端压缩变得可用——16 到 64 个 token 是理解其算力收益(对比 1024 token 的 SD-VAE)的前提。
重构-生成两难(reconstruction-generation dilemma)
在固定空间压缩率下,给潜变量加通道能提升重构质量(rFID 更低、保住纹理细节),但额外通道存储的是图像特有信息,扩散模型必须自己把这些信息生成出来,于是生成难度大增。这形成一条此消彼长的前沿:要么通道少、重构差但易生成,要么通道多、重构好但 gFID 崩坏。
APEX 的目标声明就是移动这条前沿:在通道数不断加宽时让生成质量退化得尽量慢,读懂实验部分(Figure 5)必须先理解这条权衡曲线。
流匹配(flow matching)
一种生成式训练范式:在干净数据 $x_0$ 与高斯噪声 $\epsilon$ 之间的线性插值路径 $(1-t)x_0 + t\epsilon$ 上,训练网络回归速度场 $v = \epsilon - x_0$;采样时从噪声出发用欧拉法沿速度场积分回到数据。SiT 正是基于流匹配的扩散 Transformer。
APEX 的错峰扩散完全建立在流匹配上:每个通道组在自己的局部时间 $t_\ell$ 上加噪、预测速度 $v_\ell = \epsilon_\ell - b_\ell$,推理用局部时间增量做欧拉更新。
DC-AE 与通道前缀层级
DC-AE 是支持 32×/64×/128× 空间压缩的深度压缩自编码器,把表示容量从空间位置转移到通道维。其后续 DC-AE 1.5 首次把每个 token 的通道组织成累计前缀(如 8-32-128 通道),训练时随机采样前缀,但所有层级仍在同一共享时间步上去噪。
DC-AE/DC-AE 1.5 既是 APEX 的基座架构也是最重要的对照组;理解它们"有层级但共享时间"的局限,才能体会 APEX 三要素耦合的必要性。
gFID 与 CFG
gFID 是在 5 万张类条件生成样本上按 ADM 协议计算的 Fréchet Inception 距离,越低越好;带 classifier-free guidance(CFG)的版本记作 gFID。CFG 在采样时用 $v = v_{null} + s(v_{class} - v_{null})$ 外推条件方向,$s$ 是引导尺度,典型取 1.5。
论文几乎所有结论都由固定协议下的 gFID 数字支撑(如 2.85 vs 3.22),且比较时在相同训练 FLOPs 与 NFE 下独立扫 CFG 取最优,这是理解其公平性声明的关键。
研究动机
现代扩散模型普遍采用 Transformer 去噪器,注意力开销随 token 数平方增长,因此社区追求更激进的空间压缩:$512^2$ 图像在 f8(SD-VAE)下是 1024 个 token,而 f64 只有 64 个、f128 只有 16 个。问题在于空间网格缩小时必须增加潜变量通道数才能保住重构质量,而这会撞上重构-生成两难:额外通道保存图像特有的纹理与细节,重构 rFID 确实变好(DC-AE 从 f64c128 的 0.22 到 f64c512 的 0.13),但生成模型必须凭空产出这些信息,gFID 急剧恶化——DC-AE 在 f64 下从 c128 的 3.22 涨到 c256 的 9.19、c512 的 26.10,f128 下 c1024 崩到 89.47。最接近的结构化潜变量方法 DC-AE 1.5 把通道组织成前缀层级,但所有层级仍在共享时间步上去噪,f128 下从 11.86 涨到 51.47。作者还实测:把错峰调度天真地套在 DC-AE 1.5 层级上会灾难性失败(gFID 52.89),说明这不是换个调度就能解决的问题。
本文的目标是本文的目标是:在不增加 token 数、不引入额外语义先验或显式信息瓶颈的前提下,把极端空间压缩(64×、128×)从"每步更快"升级为"质量-算力的真正优势"。具体拆成三个可检验的子目标:其一,在 f64 与 f128 的每一个被评测的通道数上(c128/c256/c512、c512/c1024)全面超过现有方法,并且随通道加宽退化得远比对手慢;其二,在训练算力对齐后,64-token 的层级模型在所有被评测的引导推理算力档位上都胜过压缩更弱的 256-token 基线,且算力越紧优势越大;其三,用 15% 的去噪器训练 FLOPs 复现广泛使用的 1024-token f8c4 SD-VAE + SiT-XL 的公开 gFID 成绩,证明极致压缩可以转化为训练与采样两端的净收益。
与已有工作不同的是,已有工作各自握有答案的一半:一条线是结构化/层级潜变量——DC-AE 1.5 的通道前缀、Matryoshka 表示、级联与嵌套扩散、多尺度 tokenizer;另一条线是异步/多时间尺度去噪——SFD 与 Learning When to Denoise 用 DINOv2 语义潜变量加独立的纹理潜变量按固定或可学习的相对调度先去噪语义,WaiT 在像素空间延迟小波细节。但本文指出两者并不能自动拼装:共享时间扩散在 APEX 层级上表现差,错峰局部时间套在 DC-AE 1.5 层级上直接崩坏(gFID 52.89)。APEX 的独特切入是把三件事——每层表示什么(分辨率监督的累计通道级)、何时去噪(错峰局部时间)、训练算力如何沿轨迹分配(全局时钟采样密度 $p_{L0}$)——作为一个整体来协同设计,并且不依赖任何外部语义表示或潜变量先验,让后面的层级自由保存图像特有细节。
核心方法
直觉上,粗结构必须凭空生成、最难;纹理与细节一旦场景确立就容易预测。因此 APEX 把生成分解为"先生成结构,再条件化细化",并在通道维而非空间维做文章。技术路线分两半:表示侧,单个自编码器把每个 token 的 $C$ 个通道切成从结构到细节的累计通道组 $z = [b_0 \| b_1 \| \cdots \| b_{L-1}]$,累计前缀 $z^{(\ell)}$ 称为一个层,并用逐级提高的目标分辨率监督(f64c8-32-128 的三层分别对应 $128^2/256^2/512^2$),迫使窄层保存连贯布局而非孤立高频。生成侧,一个共享 Transformer 在单一全局时钟 $\tau$ 上以错峰局部时间 $t_\ell(\tau) = \mathrm{clip}(1 + \ell g - \tau, 0, 1)$ 去噪各层,早层永远比晚层更干净,为细化提供条件上下文。整个过程沿通道维展开:token 数不变,每步一次网络评估服务所有层,单步成本与平坦模型相同。
核心创新是表示、时间、训练分配的三方耦合,缺一即败。第一,分辨率匹配的累计层级:训练 AE 时均匀采样一层,置零其后通道,解码 $\hat{x}^{(\ell)} = D([b_0\|\cdots\|b_\ell\|0])$,再把重构与真值缩放到该层目标分辨率后算损失 $\mathcal{L}_{AE}$;若全用全分辨率监督,早期重构常只有锐利边缘而无周围结构。第二,错峰局部时间:全局时钟映射出各层局部时间 $t_0 \le t_1 \le \cdots$,每层处于未开始(置零)、去噪中、已完成(干净上下文)三态之一,流匹配加噪 $\tilde{b}_\ell = (1-t_\ell)b_\ell + t_\ell\epsilon_\ell$。第三,训练分配:$q(\tau)$ 把 $p_{L0}=0.9$ 放在 $L_0$ 去噪的区间 $(0,1)$,因为 $L_0$ 背负无条件结构生成负担,后续层级只是更易的条件细化。DC-AE 1.5 只有第一件事,SFD/WaiT 只有第二件事,这正是与它们的本质区别。
方法步骤详情
流程分三阶段。阶段一(AE 训练):图像编码为 $z$ 并切成 $L$ 个通道组;每步均匀采样一层,保留累计前缀、置零后续通道,全分辨率解码后与真值一起缩放到该层目标分辨率(f64c128 为 128/256/512)计算重构+感知+对抗损失。阶段二(扩散训练):从 $q(\tau)$ 采样全局时钟($p_{L0}=0.9$ 落在 $(0,1)$、其余均匀落在 $(1,T)$),按 $t_\ell(\tau)$ 映射局部时间;未开始组置零、去噪中组加噪、已完成组作干净上下文;共享 DiT 预测速度 $v_\ell = \epsilon_\ell - b_\ell$,损失中 $w_\ell$ 仅在去噪态非零且跨层归一化。阶段三(推理):时钟离散化到含激活点 $\ell g$ 的网格,新组从高斯噪声初始化;每步一次网络评估,欧拉更新 $\tilde{b}_\ell \leftarrow \tilde{b}_\ell + \Delta t_\ell \hat{v}_\ell$;任一层到 $t=0$ 即可解码累计前缀出图。
技术新颖性
新颖性有四点。第一,首次系统证明"层级通道 × 异步去噪"的组合非平凡:共享时间扩散处理 APEX 层级使 gFID 从 8.12 恶化到 17.83(仅加层级甚至到 20.37),错峰调度套在 DC-AE 1.5 的均匀细层级上则崩溃到 52.89——收益来自表示-扩散协同设计而非任何单一模块,反驳了"把两个好组件拼起来就好"的直觉。第二,分辨率监督的通道前缀是新的表示侧设计:与 DC-AE 1.5 的均匀小层级不同,APEX 层级几何间隔递增(8-32-128)且各在匹配分辨率下监督,得到低分辨率下连贯可用的前缀(Figure 7 显示 DC-AE 1.5 早期前缀泄漏高频细节而 APEX 没有)。第三,异步生成不加 token、不加单步成本:不同于级联扩散(多模型串联)和 SFD(语义/纹理两套潜变量),APEX 用单一共享解码器与 Transformer,等间距调度下每步退化为对去噪块的一次共享欧拉更新。第四,训练分配 $p_{L0}$ 把结构生成负担显式化为超参数,且 $p_{L0}=1$(尾部无直接监督)反而以 30.77 优于均匀采样的 35.45。
实验结果
五组结果。(1)通道扩展(Figure 5):f64 下 APEX gFID 2.85/3.15/4.90,DC-AE 涨至 9.19/26.10;f128 下 APEX 8.98/18.77,DC-AE 14.44/89.47——APEX 全部第一且退化最缓(rFID 已对齐)。(2)对标 SD-VAE:APEX 以 $1.78\times10^{20}$ FLOPs 达 gFID 2.62,与 SiT-XL/2 的 $1.21\times10^{21}$ 持平,仅 15% 算力。(3)受限算力:1× 预算下 APEX 2.74 vs f32c32 2.95、平坦 f64 3.22,平坦 f64 在 1/2× 处被 f32 反超,APEX 全程领先。(4)消融(Table 2):DC-AE 13.27→+层级 20.37→+前缀扩散 17.83→+APEX 扩散 8.12;DC-AE 1.5+朴素错峰 52.89。(5)调度:$g=0.25$ 最优(31.38 vs 50.01),$p_{L0}=1$ 最优(30.77),层级贡献 19.83/9.04/8.12。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet-512 类条件生成(f64c128,64 token) | gFID(CFG 1.5) | 2.85 | 3.22(DC-AE)/ 3.15(DC-AE 1.5) | 降低 11.5% / 9.5% |
| ImageNet-512 生成(f64c512,64 token,通道加宽) | gFID(CFG 1.5) | 4.90 | 26.10(DC-AE)/ 14.76(DC-AE 1.5) | 降低 81.2% / 66.8%,退化最平缓 |
| ImageNet-512 生成(f128c512,16 token) | gFID(CFG 1.5) | 8.98 | 11.86(DC-AE 1.5)/ 14.44(DC-AE) | 降低 24.3% / 37.8%,所有对手均为两位数 |
| ImageNet-512 生成(f128c1024,16 token,极端宽度) | gFID(CFG 1.5) | 18.77 | 51.47(DC-AE 1.5)/ 89.47(DC-AE) | 降低 63.5% / 79.0% |
| 与 SD-VAE f8c4 基线达成同质量(1024 token vs 64 token) | gFID / 去噪器训练 FLOPs | 2.62 / $1.78\times10^{20}$(SiT-XL/1,1M 步) | 2.62 / $1.21\times10^{21}$(SiT-XL/2,3M 步) | 同 gFID 仅用 15% 训练算力 |
| 受限推理算力(1× 预算:f64 250 NFE / f32 64 NFE) | gFID(各模型独立选最优 CFG) | 2.74(APEX f64c8-32-128) | 2.95(DC-AE f32c32)/ 3.22(DC-AE f64c128) | 降低 7.1% / 14.9%,且全预算区间领先 |
局限与改进
作者承认的局限:APEX 仍需人工选择通道边界、每层目标分辨率、错峰间隙 $g$ 与全局时钟分布 $p_{L0}$,虽然 §4.4 给出了实用指引($g \approx 0.25$、$p_{L0}=0.9$),但这些最优选择可能随数据集、分辨率与模型规模变化,与自编码器和去噪器联合学习它们被留作未来方向。我的补充观察:其一,全部评测限于 ImageNet-512 类条件生成,文本到图像、更高分辨率与开放域分布上的泛化未经检验,而结构-细节分解在开放提示下未必同样干净;其二,层级深于三层时尾部覆盖依赖 $p_{L0}=0.9$ 的折中($p_{L0}=1$ 会使 $\tau \ge 1$ 后开始的层级几乎无直接监督),说明方法对层级深度有隐性敏感性;其三,自编码器训练被排除在算力预算之外,多分辨率损失与前缀采样带来的额外开销未被量化;其四,通道组内部所有通道共享同一局部时间,组内更细的渐进调度未被探索;其五,gFID 对纹理保真度的刻画有限,细化层引入的细节质量缺少人评或感知指标佐证。
独立分析的弱点
独立分析的弱点。第一,超参数敏感且需手工共调:间隙扫描呈 U 形($g=0$ 退化为同步扩散的 50.01、$g=0.25$ 的 31.38、$g=0.75$ 又升到 45.40),最优间隙在更宽潜变量或更小去噪器下可能漂移;改进方向是把 $g$ 与 $q(\tau)$ 参数化为可学习对象(元梯度、bandit 式在线分配)。第二,层级边界与分辨率分配依赖"累计通道 4× 对应边长 2×"的启发式,8-32-128 在别的压缩率下未必最优;改进方向是用可微架构搜索或逐通道可学习的分辨率权重。第三,推理轨迹尾部效率存疑:$p_{L0}=1$ 时后段时钟无直接监督仍能工作(30.77),暗示细化窗口训练欠采样但推理仍要走完整路径;改进方向是设计非均匀推理步分配,让 $L_2$ 窗口步数更少更准。第四,组内所有通道共享同一局部时间,是粗糙的时间量化,组内引入子调度可能进一步提质量。第五,方法完全在通道维展开,与视频时序维、文本 token 的联合异步设计未验证;视频扩散中 token 压力更大,既是机会也是当前框架未覆盖的盲区。
未来方向
作者明确提出:把通道边界、每层目标分辨率、错峰间隙与全局时钟分布同自编码器和去噪器联合学习,是最有希望的方向。基于其成果还可延伸:其一,迁移到文本到图像与视频生成——视频潜变量 token 数是最大瓶颈,16-64 token 的收益会被放大,且时序维度天然提供由粗到细的层级结构;其二,与语义对齐方法正交结合——APEX 不需要外部语义先验,但可在 $L_0$ 层挂 REPA 式辅助损失进一步加速结构生成收敛;其三,"任意时刻可停"的前缀解码(Figure 3/9 已展示中间层出图)配上为前缀专门训练的细化解码器,可以产品化为 anytime generation,用少量 NFE 先出草图再逐步精修;其四,理论侧解释为何 $g=0.25$ 最优——分析错峰调度下条件上下文的噪声水平与细化层学习信号之间的权衡,给出间隙的近似最优刻画;其五,把"训练算力沿生成轨迹非均匀分配"的思想推广到掩码生成或自回归模型,那里的早层-晚层不对称同样存在;其六,与少步蒸馏结合,把 16-token 模型的错峰采样蒸馏到 1-4 步,逼近实时生成。
复现评估
复现条件良好。代码已开源(https://daviddinkevich.github.io/APEX),附录提供完整材料:AE 与扩散训练全部细节(AdamW、学习率 $10^{-4}$、EMA 0.9999)、三个算法伪代码、全部层级配置(Table 5)、FLOP 记账公式 $F \simeq 3 F_{forward} B S$、评测协议(5 万样本、ADM、CFG 1.5),以及附录 F 逐图复现步骤。数据仅需 ImageNet-1K 的 $512^2$ 训练集。算力分层明显:主结果 SiT-XL/1 在 f64 上需 500K 步 batch 2048(约 $1.8\times10^{20}$ FLOPs)、f128 上 1M 步 batch 4096,完整复现需多卡长时训练;但组件消融与调度分析用 SiT-B/1 在 f64c128 完成,规模小得多,核心结论(层级与异步必须耦合、$g$ 与 $p_{L0}$ 的最优区间)可低成本验证;AE 按 DC-AE 配方 batch 64 训练。综合:复现难度中高,主结果级复现需可观算力,但方法论验证门槛友好,开源质量上乘。
论文图表
间隙扫描完整数据(CFG 1.5 / 无引导 / sFID / IS):g=0 为 50.01(即同步扩散,IS 仅 87)、g=0.10 为 40.10、g=0.25 为 31.38(IS 160)、g=0.50 为 33.48、g=0.75 延伸点为 45.40。
给出错峰间隙的完整 U 形曲线及无引导指标,量化"更干净条件上下文"与"固定 NFE 效率下降"间的权衡,是调参依据。
组件研究的 sFID 与 IS 补充:(a) DC-AE/同步 13.27/7.42/117;(b) APEX AE/通道均值 20.37;(c) APEX AE/层级均值 21.08;(d) APEX AE/前缀采样 17.83;(e) 全分辨率 AE/错峰 8.77;(f) DC-AE 1.5/错峰 52.89(IS 崩至 31);(g) 完整 APEX 8.12/6.78/163。
为 Table 2 的七个配置补齐多指标证据,特别是 (f) 的 IS 31 说明错误组合不仅 gFID 差而且样本多样性彻底坍塌。