← 返回 2026-05-27

JLT：在潜空间扩散 Transformer 中直接预测干净潜码 JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Funing Fu, Tenghui Wang, Junyong Cen, Qichao Zhu, Guanyu Zhou 📅 2026-05-26 👍 33 2026-07-13 08:36
DiT VAE 扩散模型 潜空间生成 预测目标
arXiv HuggingFace PDF

在固定 VAE 潜空间中，对比直接预测干净潜码与速度预测的差距。

前置知识

扩散模型与流匹配

扩散模型通过学习逆转一个加噪前向过程来生成样本；流匹配则把生成视为在噪声与数据之间学习一个传输速度场 v=x−ϵ。两者最终都让模型拟合一个去噪目标，但选择的直接回归量可以是噪声 ϵ、干净数据 x 或速度 v。

本文核心正是比较这三种直接回归量在潜空间中的效果差异，必须先理解 x、ϵ、v 在数学上可线性互换但实际训练动态不同这一前提。

潜空间扩散与 VAE 编码器

LDM 把高分辨率图像先经一个预训练 VAE 编码到低维潜空间，再在该空间训练扩散模型，最后解码回像素。这样能把 Transformer 的序列长度大幅缩短并降低算力开销。FLUX.2 VAE 是本文固定使用的潜空间 tokenizer。

本文消融实验固定的就是 FLUX.2 VAE 编码后的潜码，需要理解为什么要把生成过程搬到潜空间、而 VAE 解码器在采样末段才被调用。

ViT/DiT 类 Transformer 主干

DiT 用 ViT 风格的块替换传统卷积 U-Net，把潜码切成 patch 序列后由多头自注意力处理。本文的 JLT-B/1 沿用 JiT-B/16 的 12 层 Transformer、隐维 768、12 头，瓶颈 patch 嵌入 128 维。

JLT 是基于 DiT 的目标消融而非新架构，需熟悉 patch 嵌入、AdaLN 条件化、CFG（无分类器引导）等 DiT 标准组件才能看懂实验设置。

局部高斯近似与 Bayes 后验

把数据分布局部线性化为 x∼N(0,Σ)，噪声 ϵ∼N(0,I)，则给定混合 zt=tx+(1−t)ϵ 时可解析求出 E[x|z]、E[v|z] 及其条件方差 Var(x|z)、Var(v|z)。这是本文给出机制解释的核心工具。

论文的理论分析全部建立在局部线性高斯假设上，必须理解 Σ 的各向异性和单位矩阵 I 之和如何放大低方差方向，才能跟上推导。

研究动机

扩散模型的直接回归目标在数学上可以互相线性转换：ϵ-预测、x-预测、v=x−ϵ 预测通过一个仿射读出即可互换。EDM 与多篇综述因此把目标选择视为"标注习惯"，认为只需关心损失权重、预条件与采样器。但 JiT 在像素空间证明，干净图像集中在低维流形附近，而噪声和速度目标含有不在流形上的环境分量，让有限容量的 Transformer 拟合噪声比拟合干净数据更难。于是出现一个尚未回答的问题：若模型早已离开像素空间、运行在 VAE 压缩后的潜码上，原本就稀疏的方向已被大幅压缩，目标选择是否还重要？现有 LDM、DiT、SiT 体系几乎都默认采用 v-或 ϵ-目标，没人系统比较过潜空间内的 x-目标与 v-目标。

本文的目标是本文希望在固定 FLUX.2 VAE 表示、相同 Transformer 规模（Base，130M 参数）、相同优化器与训练步数（250K 步 / 200 epoch）的严格控制条件下，对比干净潜码预测（JLT-B/1、B/2）和速度预测（DiT-B/1、B/2）的 ImageNet 256×256 生成质量，并给出可解释目标差异为何仍然存在的理论机制，让"目标即几何"这一观点从像素空间延伸到潜空间。

与已有工作不同的是，此前关于"目标几何"的工作（JiT、REPA、RiT 等）要么在像素空间讨论，要么通过更换 tokenizer、增强表示来提升效果，没有人在固定表示、固定架构、固定训练的纯单变量条件下单独剥离出"直接输出参数化"这一维度。本文的关键切入角度就是：让一切外部因素全部冻结，只让网络头输出的 y 切换为 x 还是 v，用配对的基线（B/1 与 B/2 两组 patch 尺度）报告 FID 与 IS 的差距，再用局部高斯分析解释这种差距的几何成因，从而把潜扩散里长期被忽视的目标选择重新摆回设计空间。

核心方法

作者构造一组"配对实验"：取一个 130M 的 Base-scale Transformer，把 FLUX.2 VAE 冻结作为潜码 tokenizer，仅在网络头把直接回归目标切换为干净潜码 $y_x=x$（JLT-B/1、B/2）或速度 $y_v=x-\epsilon$（DiT-B/1、B/2），其它一切（架构、AdaLN 时间条件、CFG、CFG 强度、Heun 50 步采样、batch=1024、AdamW lr=$2\times10^{-4}$、250K 步）完全相同。在 FID-50K 与 IS 上直接对打，再用一个局部线性高斯模型给出理论机制：从边缘协方差 $\mathrm{Cov}(y_v)=\Sigma+I$ 出发，说明 v-目标相对 x-目标多了各向同性的"单位地板"，并放大低方差方向。

核心创新在于把"目标参数化"从被忽视的实现细节重新定义为潜扩散里一个独立的"几何设计维度"。本质区别是：JiT 类方法在像素空间用大 patch 直接预测 x，本质是"低维流形假设"，而 JLT 把同样的逻辑迁移到固定 FLUX.2 VAE 潜空间，并在可控对比中证明它带来 FID-50K 从 6.56 降到 2.56、IS 从 132.12 升到 220.74 这种数量级改进；与之配套的局部高斯推导把经验差距量化成 $\mathrm{Var}(v_i|z_i)=\mathrm{Var}(x_i|z_i)/(1-t)^2$，并指出 x-预测在 $\lambda_i\to0$ 时把系数压成 0，v-预测则趋于 $-1/(1-t)$，首次给出 v-目标"放大噪声方向"的解析机制。

方法步骤详情

方法分四步。第一步取 ImageNet 经冻结 FLUX.2 VAE 编码得潜码 $x\in\mathbb{R}^D$，解码器仅在采样末段使用。第二步按 $z_t=tx+(1-t)\epsilon$ 加噪，$t\in[0,1]$ 依 JiT 时间分布采样。第三步网络输出：JLT 头输出 $\hat x_\theta$，DiT 头输出 $\hat v_\theta=x-\epsilon$；训练把 $\hat x$ 经 $\hat v^{(x)}_\theta=(\hat x_\theta-z_t)/(1-t)$ 读出再算损失，等价用 $(1-t)^{-2}$ 加权 x-误差；AdamW 优化器，$\beta_1=0.9,\beta_2=0.95$，无 weight decay，lr $5\times10^{-5}$ 按 batch 1024 缩放至 $2\times10^{-4}$，共 250K 步，关闭 JiT 的 class-token 拼接与辅助分类损失以保消融纯净。第四步 50 步 Heun 采样，分报有/无 CFG 的 FID-50K 与 IS，CFG 强度两组目标间一致。

技术新颖性

技术新颖性体现在三个层面。首先是实验设计上的"配对消融"：在 LDM/SiT/DiT 之后首次严格控制除直接输出参数化以外所有变量，让目标本身成为唯一自变量，FID 从 6.56 到 2.56 的差距无法用 tokenizer、骨干或训练时长解释。其次是理论配套：局部高斯框架下 $\mathrm{Cov}(y_x)=\Sigma$、$\mathrm{Cov}(y_v)=\Sigma+I$ 的对比以及 $\mathrm{Var}(v_i|z_i)=\mathrm{Var}(x_i|z_i)/(1-t)^2$ 的恒等式，把代数等价与有限容量训练动态的差别量化到条件方差层面。最后是实践含义：把"目标即几何"从像素空间搬到潜空间，意味着即便表示已经压缩，目标参数化仍是一个独立的设计旋钮，研究者不应再把它当成无害的标注。

实验结果

三项核心证据。第一，Table 1 配对目标消融：VAE 网格 /1 下 JLT-B/1（x-目标，CFG）FID-50K 2.56、IS 220.74，速度基线 DiT-B/1 仅 FID 6.56、IS 132.12，绝对 FID 改进 4.00、相对降幅约 61%；/2 下 JLT-B/2 把 FID 从 DiT-B/2 的 28.71 拉到 14.81、IS 从 58.46 升到 107.29，相对降幅约 48%，优势不是 patch 大小的偶然产物。第二，Figure 2 训练曲线：两配对从首检点起就保持 FID 低、IS 高排序，JLT-B/1 在 100K 步进入低 FID 区并把差距维持到 250K 步全程。第三，Table 1 末尾最终评测：CFG 调到最终设置后 JLT-B/1 FID 降至 2.50、IS 232.51，无 CFG 时 14.00；Table 2 与 JiT-L/16（2.79）、LDM（3.60）、JiT-B/16（4.37）横向对比，凸显 130M 潜模型即可逼近 459M 像素模型水平。

Matched latent target ablation on ImageNet 256 × 256.
Table 1: Matched latent target ablation on ImageNet 256 × 256.
Guided ImageNet 256×256 comparison with representative baselines.
Table 2: Guided ImageNet 256×256 comparison with representative baselines.
ImageNet 256 × 256 samples from JLT-B/1 using 50-step Heun sampling.
Figure 1: ImageNet 256 × 256 samples from JLT-B/1 using 50-step Heun sampling.
Training curves for the matched target ablation.
Figure 2: Training curves for the matched target ablation.
查看结构化数据
任务指标本文基线提升
ImageNet 256×256 类别条件生成（CFG，VAE 网格 /1） FID-50K 2.56（最终设置 2.50） DiT-B/1 (v-目标) 6.56 绝对 -4.00，相对降幅约 61%
ImageNet 256×256 类别条件生成（CFG，VAE 网格 /1） Inception Score 220.74（最终 232.51） DiT-B/1 132.12 +88.62，相对提升约 67%
ImageNet 256×256 类别条件生成（CFG，VAE 网格 /2） FID-50K 14.81 DiT-B/2 28.71 绝对 -13.90，相对降幅约 48%
ImageNet 256×256 类别条件生成（CFG，VAE 网格 /2） Inception Score 107.29 DiT-B/2 58.46 +48.83，相对提升约 84%
无 CFG 的 JLT-B/1 最终评估 FID-50K 14.00 （无对应 DiT 无 CFG 配对基线） 体现无引导下 x-目标自身仍然可用

局限与改进

作者明确承认三点局限。其一，局部高斯分析是"局部线性"近似，并非完整生成理论，不能直接声称 x-目标对任意 tokenizer、噪声调度、损失权重或采样器都全局最优；现实潜分布是非高斯的，且局部协方差会随类别与空间位置变化。其二，实验集中在 ImageNet 256×256 和 130M 的 JLT-B/1 配置，不能外推到 XL 规模、其他数据集（如 LAION）或更大 VAE 网格。第三，没有使用 JiT 的重复 class-token 拼接和辅助分类损失，这两项都是潜在增强项，关掉它们虽然让消融更纯净，但也意味着本结果不等同于完全体的 JiT/SiT 复现，需要在更大规模下重新验证机制是否依然成立。

独立分析的弱点

独立观察可以补充几条弱点。第一，CFG 强度、Heun 步数、CFG schedule 没有做目标相关的扫描，FID 从 2.56 降到 2.50 的最后一公里可能蕴含 v-目标在更高 CFG 下也能补齐的潜在可能，论文没有直接对比 DiT 在最优 CFG 下的极值。第二，VAE 网格 /1 与 /2 之间缺乏对 token 数和有效秩的独立分析，没法判断 v-目标在 token 维度更紧的情况下是否还能保持 14.81 vs 28.71 的差距。第三，理论部分的局部高斯模型只覆盖一维独立坐标，没有给出端到端 FID 改进的上界，机制和指标之间缺一个直接量化桥。第四，论文没有公开 tokenizer 与采样随机种子，第三方复现时需要假设 FLUX.2 VAE 的具体解码配置，潜在影响不可忽略。改进方向：对 v-目标加上更细致的 CFG 与步数扫描、引入 token 维度缩放、给出更紧的 FID 上界，并在附录列出 VAE 与采样的全部超参。

未来方向

作者在附录 D 提出三条诊断方向值得跟进：一是比较 $y_x$ 与 $y_v$ 的有效秩（effective rank），预期 $y_v$ 在 Σ 各向异性时更高；二是用 kNN 等非参方法估计局部条件方差，验证 v-目标的真实后验不确定性大于 x-目标；三是用相同预算训练有限容量探针，看 $y_x$ 是否比 $y_v$ 更容易拟合。基于这些思路还可以延伸：把 JLT 拓展到 text-to-image 与视频潜空间；在 XL 规模与 SD3、FLUX 等更强 VAE 上检验机制；把 $y_v=\Sigma+I$ 的"单位地板"作为可学习 bias 或噪声注入项，缓解速度目标的方差膨胀；以及把这一目标几何视角推广到 consistency 模型与 rectified flow 的 few-step 蒸馏场景。

复现评估

复现门槛整体可控。代码仓库公开在 https://github.com/akatsuki-neo/JLT，主体结构沿用 JiT-B/16 的公开实现，模型仅 130M 参数，单阶段训练 250K 步（约 200 epoch）加有效 batch 1024。所需算力以 250K 步 AdamW、单卡 A100/H100 量级即可推进，但因为 ImageNet 全量 1.28M 图与 50K 步 Heun 采样下的 FID-50K 评估，完整复刻 Table 1/2 仍需要数天到一周级 GPU 时间。复现难度主要在三处：FLUX.2 VAE 的具体变体与解码配置需对齐、CFG 强度和 Heun 步数对最终数字敏感、以及 cfg 引导下需要单独跑两遍前向。整体属于中度可复现，第三方在给出 VAE 与采样超参的前提下应能复现主消融，精确数字可能因随机种子与 batch 顺序有小幅漂移。

论文图表

Equation (3): 边缘目标协方差 Cov(yx)=Σ, Cov(yϵ)=I, Cov(yv)=Σ+I

在局部高斯 x∼N(0,Σ)、ϵ∼N(0,I) 假设下，三种直接回归目标的协方差显式表达；其中 v-目标相对 x-目标多了各向同性的单位矩阵 I，体现单位地板。

是理论机制章节最核心的等式，把目标即几何具象为可比较的协方差，决定了后文 v-目标放大低方差方向的论证。

Equation (5): 单坐标条件方差 Var(xi|zi) 与 Var(vi|zi) 的 Bayes 残差

令 $z_i=tx_i+(1-t)\epsilon_i$、$D_i=t^2\lambda_i+(1-t)^2$，得到 $\mathrm{Var}(x_i|z_i)=\lambda_i(1-t)^2/D_i$，$\mathrm{Var}(v_i|z_i)=\lambda_i/D_i$，并由此推出 $\mathrm{Var}(v_i|z_i)=\mathrm{Var}(x_i|z_i)/(1-t)^2$。

给出 v-目标条件不确定性严格大于 x-目标的解析量化，是论文几何机制证明的关键一步。

Equation (6): 条件方差比值 Var(vi|zi)=Var(xi|zi)/(1−t)^2

对每个特征方向 $i$ 都成立的恒等式，说明 v-目标把回归残差按 $(1-t)^{-2}$ 放大，且当 $t\to1$ 时放大效应趋于无穷。

是支持 x-目标优选的最简洁机制命题，决定读者能否理解作者对为什么代数等价不等于训练等价的回答。

Figure A1 (Appendix C/D): Implementation details & additional geometry diagnostics

附录 C 列出 AdamW 优化器细节（β1=0.9、β2=0.95、无 weight decay、lr 5e-5、缩放后 2e-4、batch 1024）和训练设置；附录 D 提出三条未来诊断：有效秩比较、kNN 局部条件方差估计、有限容量探针拟合难度对比。

复现性核心配置和未来工作的方向集中于此，对决定是否复现、能否复现都非常关键。