非对称流模型:用低秩噪声预测突破像素空间扩散瓶颈 Asymmetric Flow Models
非对称参数化流速度(数据全维、噪声低秩),实现像素空间 DiT 训练及潜空间→像素微调
前置知识
流匹配(Flow Matching)
流匹配是一类生成建模框架,定义样本 $x_0$ 与高斯噪声 $\epsilon \sim \mathcal{N}(0, I)$ 之间的线性插值 $x_t = \alpha_t x_0 + \sigma_t \epsilon$(典型取 $\alpha_t = 1-t, \sigma_t = t$),并训练神经网络预测速度场 $u = (\epsilon - x_0)/\sigma_t$,其等价于求解反向常微分方程把噪声搬运到数据。流匹配与 rectified flow、扩散模型在数学上密切相关,是当前 Stable Diffusion 3、FLUX 等主流文生图系统的核心训练目标。
AsymFlow 直接基于流匹配损失设计预测目标,因此读者必须先理解什么是 velocity $u$、什么是数据项 $-x_0$ 与噪声项 $\epsilon$,才能看懂 AsymFlow 如何对二者做不对称处理。
速度预测 vs 数据预测(u-prediction vs x0-prediction)
在扩散/流匹配训练中存在多种预测参数化选择:u-prediction 让网络直接输出速度 $u$;x0-prediction 让网络输出干净数据 $x_0$,再通过 $\hat u = (x_t - \hat x_0)/\sigma_t$ 换算回速度。x0-prediction 避免直接回归高斯噪声,但在低噪声($\sigma_t \to 0$)时除法病态;u-prediction 数值稳定但在像素空间下需要让网络内部表征承载高维噪声,造成容量浪费。
AsymFlow 的核心创新正是对二者做'空间维度'上的解耦——低秩子空间用 u-prediction、正交补用 x0-prediction——所以理解这两种参数化的优劣是看懂本文关键洞见的前提。
潜在空间扩散与自编码器
为降低像素空间维度爆炸带来的计算代价,主流系统(SDXL、FLUX 等)先用 VAE/GAN 式自编码器把图像压缩到低维潜空间 $z$,在潜空间做扩散,再用一个固定的 decoder 解码回像素。这套做法显著降低算力,但代价是细粒度细节完全由 decoder 决定,生成模型无法直接控制。
本文要把'潜空间→像素'微调打通,必须借助潜空间与像素空间之间的 Procrustes 线性对齐,因此读者需要了解什么是 latent code 与 decoder。
DiT / JiT 等 Plain Transformer 架构
Diffusion Transformer(DiT)以 patch token 序列为输入、用纯 Transformer 块(无 skip connection)建模,是当前 scaling law 最优的扩散架构。JiT 把 DiT 的预测目标改为 x0-prediction,在像素空间也能良好工作。但 plain DiT 缺少 U-Net 的 skip 通道,难以把高维噪声从输入'短路'到输出。
AsymFlow 声称不需要任何架构修改即可在 plain DiT 上做像素空间生成,论文的核心卖点就是和 plain transformer 兼容,因此需要了解 DiT 与 U-ViT / DDT 等带旁路架构的区别。
Procrustes 对齐与 PCA
Procrustes 对齐是在正交约束下求两个矩阵之间最优旋转的经典方法;PCA 则把数据沿方差最大的方向分解,得到一组正交基。AsymFlow 用 PCA 选数据子空间,用 Procrustes 把潜空间线性提升到像素 patch 子空间。
AsymFlow 的低秩子空间由 PCA 或 Procrustes 给出,理解这些线性代数工具是理解子空间如何选择、如何保证对齐的关键。
研究动机
当前最强的文生图模型(FLUX、Stable Diffusion 3、Qwen-Image 等)几乎全部在 VAE 自编码器压缩后的潜空间中做扩散生成,这一选择虽然换来了 scalability,但也把细节完全交给了一个固定 decoder,生成模型无法对像素级纹理、边缘、字体等做直接控制。作者指出,'回到像素空间'是社区重新燃起的兴趣,但一旦在像素空间直接训练 plain DiT,立刻会撞上 velocity prediction 的维度瓶颈:velocity $u = \epsilon - x_0$ 中噪声项 $\epsilon$ 是 $D$ 维高斯量(典型 patch 维度 $D = 768$),要让网络内部表征承载如此高维的随机量,会'污染'中间特征、形成信息瓶颈。已有缓解方案要么给 DiT 加上 U-ViT 风格的 skip connection,要么在末端加 DDT decoder head,这两种做法都破坏了 plain transformer 的简洁性,使得大规模 latent 模型难以直接搬到像素空间。此外,纯 x0-prediction 虽然回避了直接回归噪声,但 $\hat u = (x_t - \hat x_0)/\sigma_t$ 在低噪声水平 $\sigma_t \to 0$ 时数值病态;纯 u-prediction 又要在全维度上预测高斯噪声,因此这两种参数化在高维像素空间都不是最优解。
本文的目标是本文希望提出一种新的参数化方案,让 plain DiT 类架构无需任何架构改动就能在像素空间做流匹配训练,并把'高维噪声通过中间特征'这一负担压到可控范围;同时希望打通往像素空间的迁移路径——让已经训好的潜空间大模型(典型如 FLUX.2 klein 9B)可以作为像素空间模型的初始化而非从零训练,从而把'潜空间→像素空间'变成一项微调而非重新预训练。
与已有工作不同的是,现有工作几乎都把'velocity prediction 维度过高'问题归到架构侧(加 skip、换 decoder head),或粗暴切换到 x0-prediction 接受数值病态。本文的关键洞察是:问题其实出在预测目标的'对称性'——velocity 中数据项 $-x_0$ 是有结构的、可学习的,噪声项 $\epsilon$ 是高斯随机量,二者对网络内部表征的'污染'程度完全不同。把噪声项限制到低秩子空间、把数据项保留全维度,可以同时吸收 u-prediction 与 x0-prediction 的优点:低秩子空间内做 u-prediction 仍能精细控制流场,正交补空间用 x0-prediction 不再受维度灾难影响。基于此,本文进一步设计 Procrustes 对齐把潜空间线性 lift 到像素 patch 子空间,使得潜空间预训练权重可以直接复用,从而实现了'像素空间模型的零架构微调'——这是本文区别于 DiP、PixelDiT、DeCo、PixNerd 等'加 decoder'方案的根本差异。
核心方法
AsymFlow 的整体思路分两层:(1) 设计一种新的 velocity 预测目标 $u_A = P\epsilon - x_0$,其中 $P = AA^T$ 是 patch-wise 低秩正交投影矩阵,把噪声项压到 $r$ 维子空间而把数据项保留为 $D$ 维全维度。在采样阶段,通过一个无成本的解析公式把不对称预测重新拼回全维度 velocity,无需修改 flow matching 训练/采样流程。子空间基底 $A$ 可以由图像 patch 的 PCA 给出(从零训练场景),也可以由 latent code 与像素 patch 的 Procrustes 对齐给出(微调场景)。(2) 利用这种参数化的灵活性,提出 latent-to-pixel 微调方法:把潜空间 $u$-prediction 模型的输入输出线性层与一个 Procrustes 投影融合,等价于把潜模型 lift 为一个 rank-$d$ 的像素流模型;微调时只需冻结基模型、用 LoRA 修补低层 patch 差距,并配合一个方差缩减损失 $L_{VR}$ 与 LPIPS 感知损失,让模型专注于'像素投影缝隙'的修复。
AsymFlow 的核心创新是'rank-asymmetric velocity parameterization':在低秩子空间 $\mathrm{Im}(P)$ 内沿用 $u$-prediction(精细控制流场),在正交补 $\mathrm{Im}(I-P)$ 内沿用 $x_0$-prediction(避免维度灾难),二者通过正交分解无缝拼接。从两端特例看,$r=0$ 退化为纯 x0-prediction(JiT 风格),$r=D$ 退化为纯 u-prediction,而最优 $r$ 介于二者之间。本文的关键数学观察是分解 $P u_A = P\epsilon - P x_0 = P u$ 与 $(I-P)u_A = -(I-P)x_0$,这意味着低秩部分直接是真实的低秩速度,正交补部分可以从 $x_0$-预测通过除以 $\sigma_t$ 换算——这就是 Eq. (5) 的 $u = P u_A + (I-P)x_t/\sigma_t$。这个看似简单的 trick 让 plain DiT 完全无需架构修改就能在像素空间工作,同时在 latent-to-pixel 场景下还能把潜模型'提升'为像素模型的初始化——这是和 PixelDiT/DDT 等所有'加 decoder'路线的本质区别。
方法步骤详情
AsymFlow 的训练和采样流程如下:(1) 选低秩子空间:从零训练时用 $D=768$ patch 上的 PCA 得到 rank-$r$ 正交基矩阵 $A \in \mathbb{R}^{D \times r}$、对应投影 $P = AA^T$;微调时改为在 latent code 与 pixel patch 之间做 Procrustes 对齐得到 $A$。每个 patch token 共享同一投影。(2) 重新定义目标:把 flow matching 目标从 $u = \epsilon - x_0$ 改为 $u_A = P\epsilon - x_0$,网络预测 $\hat u_A = G_\theta(x_t, t)$。这是论文 Eq. (3)。(3) 解析恢复到全维度 velocity:利用分解 $u = P u_A + (I-P) x_t / \sigma_t$,把网络输出映射回 $\hat u$;在低噪声处对正交补的 $1/\sigma_t$ 用 $\sigma_{\min}=0.04$ 做钳位,仅影响 $(I-P)$ 方向,比 JiT 的全维度钳位显著更稳定。(4) 用标准 flow matching 损失 $\mathcal{L}_{FM} = \mathbb{E}_{t,x_0,\epsilon}\|u - \hat u\|^2$ 训练;这是关键——loss 形式没变。(5) 微调路径:把潜空间模型 $\hat u_z = G_\phi(z_t, t)$ 的输入输出线性层分别融合 Procrustes 投影 $A^T$ 与 $A$,并对输出乘上尺度因子 $s$ 校正 latent 与 pixel 的方差差异,得到初始化后的像素模型 $\hat u_A = G_\theta(x_t, t)$。微调时冻结基模型,仅训 rank-256 LoRA。(6) 方差缩减损失:把 $-x_0 + \lambda (x_0^L - \mathbb{E}[x_0^L|x_t])$ 作为 control variate 注入损失,其中 $\hat x_0^L$ 由 frozen 初始化低秩模型从配对 $x_t^L = \alpha_t x_0^L + \sigma_t\epsilon$ 预测;自适应权重 $\lambda$ 用正交投影最小化梯度范数。(7) 感知修正:在扩散时间表上前段用方差缩减、后段切换到 LPIPS 感知损失,补偿方差缩减项在低噪声处的近似误差。具体训练数据:ImageNet 256×256 从零训 600 epoch;T2I 用 3M LAION-Aesthetics、1024×1024、Qwen2.5-VL 重新打标、UniPC + APG 采样。
技术新颖性
技术上 AsymFlow 有三个层次的贡献:(a) 参数化层面——首次在空间维度上做 asymmetric split,把 x0- 与 u-prediction 视为同一族参数化的两个端点(r=0 与 r=D),并通过正交分解给出闭式恢复公式。这意味着 plain transformer 可以无痛进入像素空间,无须 skip connection 或 decoder head。(b) 微调路径层面——首次实现 latent→pixel 微调。利用 Procrustes 把 latent flow ODE 严格 lift 为低秩 pixel flow ODE,证明轨迹耦合(Theorem 1),所以初始化模型生成的内容在语义与结构上与潜模型一致,只剩低层 patch 缝隙需要修补。这把'大模型从潜空间搬到像素空间'从不可能变成了工程上 3M 图就能微调的事。(c) 损失层面——把 control variate 思想用于 latent-to-pixel 配对数据,把 frozen 初始化模型的预测 $\hat x_0^L$ 当作降方差项,并配合 LPIPS 感知修正。这是把经典方差缩减技巧(已在 RL 与 MCMC 中广泛使用)首次适配到 pixel diffusion 微调场景。
实验结果
论文给出 ImageNet 256×256 与文生图两个场景的系统结果。在 ImageNet 上,与 JiT-H/16 同架构同 recipe(仅改预测参数化),AsymFlow($r=8$)最终达到 1.76 FID / IS 312.0,明显超过 JiT 1.90 FID / IS 300.8;加上标准 REPA 后达到 1.57 FID,是 plain-transformer 类像素扩散模型中的 state of the art(除更贵的 SiD2 UViT/1 之外)。Table 1 还显示 AsymFlow 对低噪声 $\sigma_{\min}$ 钳位不敏感——关掉钳位时 JiT 退化 1.37 FID(1.90→3.27),而 AsymFlow 只退化 0.52(1.76→2.28),证明其数值稳定性大幅优于 x0-prediction。Figure 5 的 patch rank 消融表明 $r=8$ 是最优,远好于 $r=0$(JiT),继续增大 $r$ 反而轻微退化;同图显示 PCA 子空间显著优于 random 子空间(PCA r=8 → 1.76 FID;random r=8 → 接近 JiT 水平),证明收益来自'有意义的低秩子空间'而非'单纯降维'。Figure 6 的收敛曲线显示 AsymFlow 比 JiT 收敛快约 40%,证实新参数化优化效率更高。在 T2I 场景,AsymFLUX.2 klein(9B FLUX.2 klein 潜模型微调而来)在 HPSv3 拿到 10.66,超越 FLUX.2 klein Base 的 9.50 和 FLUX.1 dev 的 10.43,DPG-Bench 拿到 86.8(超过 FLUX.2 klein Base 的 85.2、Qwen-Image 的 87.8 这一最强 latent 基线、PixelDiT-T2I 的 83.5),GenEval 0.82(优于 FLUX.2 klein Base 0.80、PixelDiT-T2I 0.74、SDXL 0.55)。Table 3 的消融对比 latent baseline 与 DDT finetune baseline:AsymFlow 在所有指标上均超过二者,特别 HPSv3 10.70→12.03→12.99→13.06 随方差缩减与 LPIPS 逐步提升,说明每个新组件都带来增益;DDT baseline 虽有更多参数但 HPSv3 仅 10.33,pFID 26.0 显著高于 AsymFlow 22.5。Figure 8 的视觉对比显示 DDT baseline 输出模糊且有 patch 拼接缝,纯 flow matching 已有细节、方差缩减进一步提升细节但引入过噪、加上 LPIPS 后最自然真实。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 256×256 class-conditional pixel generation | FID(ADM 评估协议) | AsymFlow-H/16 (r=8 + REPA) 1.57;AsymFlow-H/16 (r=8) 1.76 | JiT-H/16 x0-prediction 1.86*(官方),同 recipe 复现 1.90;PixelREPA-H/16 1.81* | FID 绝对降低 0.29(1.86→1.57),相对降幅约 15.6%;同 recipe 下比 JiT 降低 0.14(1.90→1.76),相对降幅约 7.4% |
| ImageNet 256×256 class-conditional pixel generation | Inception Score (IS) | AsymFlow-H/16 (r=8) 312.0 | JiT-H/16 同 recipe 300.8 | IS 提升 11.2(+3.7%) |
| T2I 1024×1024 generation, human preference | HPSv3 | AsymFLUX.2 klein (full) 13.06 | FLUX.2 klein Base 9.50;PixelDiT-T2I 8.95;FLUX.1 dev 10.43 | 比最强 latent 基线 FLUX.2 klein Base 提升 3.56(+37.5%),比 PixelDiT-T2I 提升 4.11 |
| T2I 1024×1024, fine-grained prompt alignment | DPG-Bench | AsymFLUX.2 klein (full) 86.8 | FLUX.2 klein Base 85.2;Qwen-Image 87.8;PixelDiT-T2I 83.5 | 比 FLUX.2 klein Base 提升 1.6;比 PixelDiT-T2I 提升 3.3 |
| T2I 1024×1024, compositional generation | GenEval | AsymFLUX.2 klein (full) 0.82 | FLUX.2 klein Base 0.80;Qwen-Image 0.86;PixelDiT-T2I 0.74 | 比 FLUX.2 klein Base 提升 0.02;比 PixelDiT-T2I 提升 0.08 |
| T2I patch-level realism | pFID (COCO 10K) | AsymFLUX.2 klein (full) 22.5 | DDT finetune baseline 26.0 | pFID 降低 3.5(−13.5%) |
局限与改进
作者明确承认一个关键局限:latent-to-pixel 微调假设潜空间与像素 patch 之间存在较好的 patch-level 线性 lift,这一条件在标准 VAE 潜空间(FLUX、SDXL)下基本成立,因为它们通过 pixel reconstruction loss 训练,结构与像素 patch 强相关;但对于 RAE(Representation AutoEncoder,弱像素重建约束、强调语义压缩)等新一代 latent 模型,潜空间可能不再保留像素 patch 的几何结构,因此 Procrustes 对齐可能失效,AsymFLUX 范式难以直接套用。这一限制直接影响 AsymFlow 在'自研 latent 编码器'生态里的可迁移性。论文没有进一步分析 latent 类型对微调质量的定量影响。我自己的额外观察:(1) 实验只在 ImageNet 256×256 与 LAION 3M 上验证,更复杂的训练数据分布(如高分辨率、多类目)下参数化收益是否稳定未知。(2) 论文 $\sigma_{\min}$ 选择依赖 grid search,缺少自动选择方法,限制了自动化部署。(3) 微调实验没有对比全部 PixelDiT-T2I 与最新 latent SOTA(如 FLUX.1.1 pro),且训练数据仅 3M,与工业级 T2I 训练规模差距很大,所以'pixel SOTA'的边界尚有讨论空间。(4) PCA 子空间对训练数据敏感,跨数据集迁移时是否需要重新计算 PCA 未说明。(5) AsymFlow 的 FID 1.57 与 latent 扩散模型(如 SDXL/FLUX 等在 ImageNet 上的低 FID)仍有较大差距,论文没有系统对比 latent 模型 FID,因此'pixel space SOTA'的措辞需限定在 pixel 模型范畴内理解。
独立分析的弱点
独立分析认为本文至少存在以下可改进点:(1) 子空间选择依赖数据。Table 1 的 $\sigma_{\min}=0.04$ 是手工 grid search 出来的,没有自适应机制;当 patch 数据分布偏移时,PCA 子空间的最优秩 $r$ 可能变化,缺少自动化 rank 选择方案——一个改进方向是用 singular value 的能量阈值(如累积解释 95% 方差)或在 loss 中加入 rank 正则项做端到端学习。(2) Variance-reduction 项的 $\lambda$ 用 orthogonal projection 求解,但 $\mathbb{E}[x_0^L|x_t] \approx \mathbb{E}[x_0^L|x_t^L]$ 的近似在 $t$ 较小时偏差显著,作者用 LPIPS 在后段时间修补,但这相当于'用一个固定时间表把两种 loss 拼接',不够优雅——改进方向是改为 learned schedule 或用 cross-attention 把 frozen 模型的特征重新校准到当前 $x_t$,而非依赖配对 $x_t^L$。(3) 微调阶段 rank-256 LoRA 是较小改造,但没有给出 rank 与微调质量的 scaling 关系,难以判断更激进的全参数微调是否会更优。(4) latent-to-pixel 路径要求潜空间与像素 patch 对齐良好,对 RAE 等语义 latent 无效;改进方向是引入非线性 aligner(如一个小 MLP)或直接学习 $A$,但这会牺牲 Procrustes 的解析性。(5) ImageNet 实验沿用 JiT-H/16 配置,没有在更大的 G/2 模型或更长训练下验证参数化收益是否饱和,可能的最优 $r$ 是否随模型规模变化未知。(6) T2I 实验只用 LAION 3M 微调,没有展示推理时长/显存代价与 latent 模型的对比,pixel-space 推理的工程代价是其能否取代 latent 模型的关键考量。
未来方向
作者在 Conclusion 提到的高层方向包括:高保真图像与视频生成、低层细节控制、扩展到其他高维数据模态。基于结果可延伸的具体方向有:(1) 把 AsymFlow 推广到 latent 扩散本身——既然低秩子空间在像素空间有用,是否也能在 4D 视频、3D 占据场、多视角图像等'天然高维'模态中复用?这是直接的低成本扩展。(2) 与最新 DiT scaling recipe(如 REPA-Subject、ReDi 等)做正交组合,看 AsymFlow 是否能成为下一代 plain transformer 像素扩散的默认参数化。(3) 在 T2I 之后接入 i2i、inpainting、edit 等下游任务:latent→pixel 微调过的模型天然对'像素级控制'更敏感,可能在可控编辑任务上取得新突破。(4) 把 Procrustes 对齐替换为可学习的 linear adapter 或低秩 adapter,使方法对 RAE 等弱像素约束的 latent 也生效,扩展适用面。(5) 自动化 $\sigma_{\min}$ 与 patch rank $r$ 的选择,避免 grid search,比如基于训练集上的 gradient norm 自适应调整。(6) 与 mixture-of-experts、蒸馏(consistency model)结合:pixel-space 模型推理成本更高,若能通过蒸馏把多步 ODE 压到 1-4 步,将极大提升 AsymFlow 的实用性。
复现评估
从可复现性角度评估:论文明确给出项目页 https://hanshengchen.com/asymflow 与 arXiv 链接(2605.12964v2),整体框架基于 JiT-H/16 与 FLUX.2 klein 公开实现,理论上具备较强复现可行性。但仍存在明显门槛:(1) 数据规模:T2I 实验需要 3M LAION-Aesthetics 高质量图像、Qwen2.5-VL 重打标、单 megapixel 分辨率,这套数据 pipeline 复现需要非平凡算力。(2) 算力:9B 模型在像素空间微调(即使 LoRA rank-256)也需要相当规模的多机多卡 GPU 集群,论文没有给出确切 GPU 数与训练时长。ImageNet 实验从零训 600 epoch 用 JiT-H/16(953M 参数),按 JiT 原论文估计需要数百到上千 A100/天量级。(3) 实现细节:Procrustes 对齐需要 paired latent-pixel 数据,意味着训练 pipeline 必须能同时存取 latent 与 pixel patch;附录 A.1-A.4 给出详细推导但实现细节仍较多,需要读者自行推导 patch-wise orthogonal projection 与 variance reduction 的具体写法。(4) 评估:HPSv3、DPG-Bench、GenEval 都是公开 benchmark,复现容易;FID 按 ADM 评估协议需要在 50000 样本上做 grid search,门槛较高。(5) 微调超参:$\sigma_{\min}$、$r$、LoRA rank、LPIPS 时间表都需要重新调优,对工程能力要求不低。综合评估:核心算法的可复现性较好(论文给出了充分理论与方程),但完整 T2I 实验的复现门槛高,更适合具备大模型训练经验的研究团队。
论文图表