Revisiting Diffusion Model Predictions Through Dimensionality¶

会议: NeurIPS 2026
arXiv: 2601.21419
OpenReview: https://openreview.net/forum?id=NAHtkZBZqR
论文: https://arxiv.org/abs/2601.21419
领域: 图像生成
关键词: 扩散模型, 预测目标, 流匹配, 几何维度, 自适应预测参数

一句话总结¶

本文通过极简线性扩散模型推导出预测目标误差在流形切空间与正交环境空间的严格解耦动力学，证明流匹配中最优预测目标解析解为 \(k^* = D / (D + d)\)，并提出仅需单一可学习标量的 \(k\)-Diff 框架，实现了从低维潜空间到高分辨率像素空间的自适应无缝过渡与生成增益。

研究背景与动机¶

扩散生成模型与流匹配（Flow Matching）的核心设计基石之一是网络预测目标的选择。早期工作主要聚焦于低分辨率或相对低维的数据生成，普遍采用噪声预测（\(\boldsymbol{\varepsilon}\)-prediction）；随后在追求保方差与直线概率路径的过程中，速度预测（\(\boldsymbol{v}\)-prediction）在 EDM 和流匹配框架中逐渐确立为主导范式。然而，当生成模型进一步向高分辨率、超大规模且直接在像素空间（Pixel Space）展开建模演进时（如 SiD2、JiT 等工作），研究者们在经验上发现直接预测干净数据（\(\boldsymbol{x}\)-prediction）展现出了远超速度预测的训练稳定性和细节保真度。这一经验漂移表明：扩散模型的最优预测目标并非静态常数，而是高度依赖于任务空间和数据形态。

尽管实验表象清晰，但整个领域的理论解释长期处于碎片化甚至经验猜想的状态。第一，学术界并不明确最优目标是否仅局限于 \(\boldsymbol{\varepsilon}\)、\(\boldsymbol{v}\)、\(\boldsymbol{x}\) 这三个离散的三岔路口，还是存在于更广阔的广义连续物态空间中；第二，现有文献往往仅诉诸于直觉性的“流形假设”（Manifold Hypothesis）做定性阐述，严重缺乏关于环境维度（Ambient Dimension）与内在流形维度（Intrinsic Dimension）如何决定网络学习稳定性的严格定量数学描述；第三，在工程实践中，由于高维非线性图像流形的内在维度 \(d\) 几乎无法精确计算，研究者往往依赖高昂的超参数网格搜索或纯启发式试错来选定预测目标，严重制约了新生成领域的拓展效率。

面对这些核心阻碍，本文系统性地从理论与算法两端破局。研究首先建立了一套能容纳任意凸组合形式的广义连续预测目标框架，并通过分析跨全部信噪比谱（SNR）的单层线性扩散模型，揭示出网络权重在切空间与正交余空间中截然相反的学习动力学。理论证明，最优预测目标取决于数据流形维度与环境维度的相对张力，其解析平衡点精确落在 \(k^* = D / (D + d)\)。针对内在维度不可知的问题，本文进一步设计出端到端自适应优化的 \(k\)-Diff 算法。核心 idea：将预测目标参数化为数据与噪声的连续插值 \(\boldsymbol{u} = k\boldsymbol{x} - (1-k)\boldsymbol{n}\)，理论证明最优解由环境维度 \(D\) 与内在维度 \(d\) 的比值 \(k^* = D/(D+d)\) 解析决定，并提出仅需引入一个可学习标量的 \(k\)-Diff 框架，让模型通过标准反向传播全自动学习最优目标。

方法详解¶

整体框架¶

在连续时间扩散与流匹配过程中，含噪输入 \(\boldsymbol{z}\) 与模型建模目标 \(\boldsymbol{u}\) 通常可统一写为数据 \(\boldsymbol{x} \in \mathbb{R}^D\) 与标准高斯噪声 \(\boldsymbol{n} \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I}_D)\) 的线性组合：\(\boldsymbol{z} = \alpha\boldsymbol{x} + \sigma\boldsymbol{n}\)，\(\boldsymbol{u} = \varphi\boldsymbol{x} + \psi\boldsymbol{n}\)。对于标准的流匹配，设定为 \(\alpha=t, \sigma=1-t\)。为了打破离散目标选取的局限，本文引入标量参数 \(k \in [0, 1]\)，将广义目标统一参数化为 \(\boldsymbol{u} = k\boldsymbol{x} - (1-k)\boldsymbol{n}\)。当 \(k=0\) 时退化为噪声预测 \(\boldsymbol{\varepsilon}\)，当 \(k=0.5\) 时等价于标准速度预测 \(\boldsymbol{v}\)，当 \(k=1\) 时严格对应数据预测 \(\boldsymbol{x}\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入数据 x 与高斯噪声 n"] --> Proj["投影至含噪输入 z = t x + (1-t) n"]
    Proj --> Mode["流形正交模式解耦<br/>切空间平行分支 W|| vs 正交补空间垂直分支 W⊥"]
    Mode --> Theo["解析维度平衡推导<br/>理论最小化误差得出解析解 k* = D / (D + d)"]
    Theo --> KDiff["k-Diff 可学习标量参数化<br/>定义标量 wk 经 Sigmoid 映射为 k"]
    KDiff --> LossOpt["统一速度空间映射与反向传播<br/>将 u 还原为 v 速度场并计算 L2 损失更新网络与 wk"]
    LossOpt --> Out["自适应生成采样<br/>潜空间收敛至 k=0.66，像素空间快速收敛至 k=1.0"]

在系统架构中，整体 pipeline 包含理论层面的正交动力学解耦与推导，以及实践层面的端到端梯度优化流：训练时采样时间步 \(t \sim \text{logit-normal}\)，由当前的自适应可学习参数 \(k = \text{sigmoid}(w_k)\) 构建插值目标 \(\boldsymbol{u}\)，网络预测 \(\hat{\boldsymbol{u}}\) 后映射回等效速度向量 \(\hat{\boldsymbol{v}}\) 执行一致性反向传播，实现对任意数据表征维度特性的全自动适配。

关键设计¶

1. 流形正交模式解耦动力学：揭示切空间数据恢复与正交空间噪声消除的本质冲突

为严谨分析数据几何对预测目标的影响，理论模型采用嵌入在高维欧氏空间 \(\mathbb{R}^D\) 中的白化低维流形假设，数据 \(\boldsymbol{x} = \boldsymbol{P}\tilde{\boldsymbol{x}}\)，其中 \(\tilde{\boldsymbol{x}} \in \mathbb{R}^d\) 满足 \(\mathbb{E}[\tilde{\boldsymbol{x}}\tilde{\boldsymbol{x}}^\top] = \boldsymbol{I}_{d\times d}\)，且 \(\boldsymbol{P} \in \mathbb{R}^{D \times d}\) 是正交子矩阵满足 \(\boldsymbol{P}^\top\boldsymbol{P} = \boldsymbol{I}_{d\times d}\)。考虑单层全连接权重矩阵 \(\boldsymbol{W} \in \mathbb{R}^{D \times D}\) 在均方误差目标下的连续时间梯度下降流。利用流形投影算子 \(\boldsymbol{P}\boldsymbol{P}^\top\) 及其正交补算子 \(\boldsymbol{I} - \boldsymbol{P}\boldsymbol{P}^\top\)，网络权重精确分解为两个独立演化的正交动力学模式：

\[\tau \frac{\mathrm{d}\boldsymbol{W}_{\parallel}}{\mathrm{d}\tilde{t}} = -\int \widetilde{\mathcal{D}}\alpha \Big( (\alpha^2 + \sigma^2)\boldsymbol{W}_{\parallel} - (\varphi\alpha + \psi\sigma)\boldsymbol{P}\boldsymbol{P}^\top \Big)\]

\[\tau \frac{\mathrm{d}\boldsymbol{W}_{\perp}}{\mathrm{d}\tilde{t}} = -\int \widetilde{\mathcal{D}}\alpha \Big( \sigma^2 \boldsymbol{W}_{\perp} - \psi\sigma (\boldsymbol{I} - \boldsymbol{P}\boldsymbol{P}^\top) \Big)\]

这一动力学方程深刻地揭示了两个模式的物理本质：平行分量 \(\boldsymbol{W}_{\parallel}\) 沿着数据流形切空间演化，其动力学驱动力强烈依赖于信号系数 \(\alpha\) 与数据系数 \(\varphi\)，本质上执行的是从含噪特征中重构流形几何的数据恢复（Data Recovery）任务；反之，垂直分量 \(\boldsymbol{W}_{\perp}\) 处于数据流形之外的环境空间，其演化完全不受数据项 \(\alpha\) 和 \(\varphi\) 的影响，只由纯噪声的衰减率 \(\sigma\) 与权重 \(\psi\) 主导，本质上是在正交余空间执行噪声消除（Noise Elimination）。因此，任何对数据几何视而不见的固定预测目标设计，都必然在两类截然不同的动力学需求之间产生次优甚至劣质的折衷。

2. 最优连续预测目标解析平衡：证明维度比率决定最优解 \(k^* = D / (D + d)\)

当系统演化至梯度动力学不动点平衡态时，将平衡解代回全谱目标误差函数，总最优损失 \(\Delta^*\) 精确解构为两个由维度缩放的正交贡献项：

\[\Delta^* = \frac{1}{2}d \underbrace{\left( \int \widetilde{\mathcal{D}}\alpha (\varphi^2 + \psi^2) - \frac{\big(\int \widetilde{\mathcal{D}}\alpha(\varphi\alpha + \psi\sigma)\big)^2}{\int \widetilde{\mathcal{D}}\alpha(\alpha^2+\sigma^2)} \right)}_{\text{流形内部误差 (Intra-manifold Loss)}} + \frac{1}{2}(D - d) \underbrace{\left( \int \widetilde{\mathcal{D}}\alpha \psi^2 - \frac{\big(\int \widetilde{\mathcal{D}}\alpha \psi\sigma\big)^2}{\int \widetilde{\mathcal{D}}\alpha \sigma^2} \right)}_{\text{正交环境误差 (Residual Ambient Loss)}}\]

在流匹配设定 \(\alpha=t, \sigma=1-t, t \sim \mathcal{U}[0, 1]\) 下，代入一阶参数化 \(\varphi=k, \psi=-(1-k)\)。通过对积分项进行逐项解析求解，垂直项由于余空间维度 \((D-d)\) 的巨大缩放，会随 \((1-k)^2\) 剧烈放大。令总损失对常数标量 \(k\) 求导 \(\partial \Delta^* / \partial k = 0\)，最终推导出震撼且极具几何美感的解析极值点：

\[k^* = \frac{D}{D + d}\]

该公式清晰划分了扩散模型预测目标的三大几何演化机制： - 低维/密集数据区（\(D \approx d\)）：环境维度与流形维度几乎一致，\(k^* \approx 0.5\)，严格证明了传统速度预测 \(\boldsymbol{v}\) 在稠密低维空间的理论最优性； - 超高维/稀疏流形区（\(D \gg d\)）：在高分辨率像素空间中，图像真实流形自由度远小于像素物理维度（\(D \gg d\)），此时 \(k^* \to 1.0\)，公式从第一性原理严谨证明了为什么像素扩散模型必须转向数据预测 \(\boldsymbol{x}\)； - 中间过渡带：对于潜空间（Latent Space）等压缩后但未彻底消除冗余维度的表征空间，最优目标既非纯粹的 \(\boldsymbol{v}\) 亦非纯粹的 \(\boldsymbol{x}\)，而是必然处于 \(0.5 < k^* < 1.0\) 的连续介质带中。

3. k-Diff 端到端自适应学习框架：单一标量反向传播绕过内在维度估计黑箱

尽管解析公式给出了精准的理论指导，但现实世界中的复杂图像流形高度非线性且拓扑结构繁复，其内在有效维度 \(d\) 属于经典的数学不可解问题，无法在训练前预先计算。为此，本文将理论结论工程化落地，提出了 \(k\)-Diff 算法。定义一个可微标量参数 \(w_k\)，通过 Sigmoid 激活函数将其投影到保真区间：\(k = \text{sigmoid}(w_k)\)。

网络预测广义目标 \(\hat{\boldsymbol{u}} = \text{net}(\boldsymbol{z}, t)\) 后，通过下述代数变换将 \(\hat{\boldsymbol{u}}\) 投影至流匹配统一速度向量 \(\hat{\boldsymbol{v}}\)，并构造速度尺度的反向传播损失：

\[\hat{\boldsymbol{v}} = \frac{(1 - 2k)\boldsymbol{z} + \hat{\boldsymbol{u}}}{k(1-t) + (1-k)t}\]

梯度流直接同时反向穿透主干权重 \(\boldsymbol{\theta}\) 与标量 \(w_k\)。为了避免时间步两端的分母奇点，引入数值截断阈值 \(\max\big(k(1-t) + (1-k)t, \, 0.05\big)\)。整个设计仅引入 1 个额外参数（开销在 100M-800M 模型中完全可以忽略不计），使得模型无需任何人工先验即可在训练初期沿着几何梯度自发收敛至该数据空间的最优 \(k^*\)。

损失函数 / 训练策略¶

模型通过速度预测空间的重加权损失执行梯度更新：\(\mathcal{L}(\boldsymbol{\theta}, w_k) = \frac{1}{2} \mathbb{E}_{t, \boldsymbol{x}, \boldsymbol{n}} \big[ \|\hat{\boldsymbol{v}} - \boldsymbol{v}\|^2 \big]\)。在优化器选择上采用 Adam（\(\beta_1=0.9, \beta_2=0.95\)），学习率保持常数 \(2 \times 10^{-4}\)，Batch Size 设置为 1024。时间步采样遵循 Logit-Normal 分布 \(\text{logit}(t) \sim \mathcal{N}(\mu, \sigma^2)\)（潜空间为 \(\mu=0.0, \sigma=1.0\)；像素空间为 \(\mu=-0.8, \sigma=0.8\)）。采样推断采用二阶 Heun 求解器运行 50 步（NFE=99），结合无分类器引导（CFG）完成高保真图像生成。

实验关键数据¶

主实验¶

主实验在 ImageNet \(256 \times 256\) 和 \(512 \times 512\) 基准上全面铺开，分别涵盖潜空间模型 LightningDiT-XL/1 与像素空间模型 JiT-B/16、JiT-B/32，与当期最先进的多样化扩散模型进行定量对比。评估指标包括 FID-50k（越低越好）和 Inception Score（IS，越高越好）。

表 1：ImageNet \(256 \times 256\) 条件图像生成对比实验

方法类型	代表性基准模型	预测空间	参数量	GFLOPs	FID-50k ↓	Inception Score ↑
潜空间扩散 (Latent)	DiT-XL/2 (Peebles & Xie, 2023)	Latent	675+49M	119	2.27	278.2
潜空间扩散 (Latent)	SiT-XL/2 (Ma et al., 2024)	Latent	675+49M	119	2.06	277.5
潜空间扩散 (Latent)	REPA (SiT-XL/2) (Yu et al., 2024)	Latent	675+49M	119	1.42	305.7
潜空间扩散 (Latent)	LightningDiT-XL/1 (Yao et al., 2025)	Latent	675+49M	119	1.35	295.3
潜空间扩散 (Latent)	k-Diff (LightningDiT-XL/1) (本文)	Latent	675+49M	119	1.34	301.7
像素空间扩散 (Pixel)	ADM-G (Dhariwal & Nichol, 2021)	Pixel	554M	1120	4.59	186.7
像素空间扩散 (Pixel)	RIN (Jabri et al., 2022)	Pixel	410M	334	3.42	182.0
像素空间扩散 (Pixel)	SiD (UViT/2) (Hoogeboom et al., 2023)	Pixel	2B	555	2.44	256.3
像素空间扩散 (Pixel)	VDM++ (UViT/2) (Kingma & Gao, 2023)	Pixel	2B	555	2.12	267.7
像素空间扩散 (Pixel)	PixelFlow (XL/4) (Chen et al., 2025)	Pixel	677M	2909	1.98	282.1
像素空间扩散 (Pixel)	JiT-B/16 (Li & He, 2025)	Pixel	131M	25	3.66	275.1
像素空间扩散 (Pixel)	k-Diff (JiT-B/16) (本文)	Pixel	131M	25	3.64	274.9

表 2：ImageNet \(512 \times 512\) 高分辨率像素空间与潜空间生成对比

类别	架构与模型	参数量	GFLOPs	FID-50k ↓	Inception Score ↑
潜空间 (Latent)	DiT-XL/2	675+49M	525	3.04	240.8
潜空间 (Latent)	SiT-XL/2	675+49M	525	2.62	252.2
潜空间 (Latent)	REPA (SiT-XL/2)	675+49M	525	2.08	274.6
潜空间 (Latent)	DDT-XL/2	675+49M	525	1.28	305.1
像素空间 (Pixel)	ADM-G	559M	1983	7.72	172.7
像素空间 (Pixel)	RIN	320M	415	3.95	216.0
像素空间 (Pixel)	SiD (UViT/4)	2B	555	3.02	248.7
像素空间 (Pixel)	PixNerd (XL/16)	700M	583	2.84	245.6
像素空间 (Pixel)	JiT-B/32 (Li & He, 2025)	133M	26	4.02	271.0
像素空间 (Pixel)	k-Diff (JiT-B/32) (本文)	133M	26	4.03	277.1

消融实验¶

为论证 \(k\)-Diff 采用单一常数标量 \(k\) 设计的精简性与完备性，论文在 ImageNet \(256 \times 256\)（LightningDiT-XL/1，64 epochs）上针对参数独立性与时间依赖性开展了深度消融分析。

表 3：\(k\) 的参数化自由度与时间依赖性消融研究

实验配置	目标形式说明	时间 Bin 划分	学习到的收敛值 (第 64 轮)	FID-50k ↓
k-Diff (单一标量，本文默认)	严格单标量插值 \(k\boldsymbol{x} - (1-k)\boldsymbol{n}\)	1 (全时步恒定)	\(k = 0.522\)	2.05
双独立参数解耦 (\(k_0, k_1\))	解耦系数 \(k_0\boldsymbol{x} - k_1\boldsymbol{n}\)	1 (全时步恒定)	\(k_0=0.556, k_1=0.509\)	2.04
时变分段插值 \(k(t)\)	128 个时间 Bin 线性插值连续标量	128 (时间依赖)	随 \(t\) 呈现 0.56 降至 0.3 轨迹	2.17

关键发现¶

自适应维数轨迹与理论惊人契合：在像素空间（JiT-B/16）训练中，\(k\) 在前 5%（小于 30 个 Epoch）即呈现断崖式上升并死死钉在 \(k=1.0\)（即纯 \(\boldsymbol{x}\)-prediction）；而在经过 VAE 压缩的潜空间（LightningDiT-XL/1）中，\(k\) 则从初始化的 0.5 稳定缓慢攀升至稳态值 0.66。这表明潜空间压缩虽降低了维度，但仍保留了可观的环境冗余维度，纯 \(\boldsymbol{v}\) 目标在此并非绝对最优，介于 \(0.5\) 与 \(1.0\) 之间的目标能更好地平衡重建与去噪。
时间依赖目标反而损害泛化：引入 128 个分段可微节点构建时变目标 \(k(t)\)，模型 FID 反而退化到 2.17。原因在于时变目标导致不同信噪比下的学习目标动态漂移，破坏了跨时间步网络权重的连续收敛性；同时 Logit-Normal 时间采样的两端梯度消失问题也导致端点学习失真。
扩展长轮次下的鲁棒性：LightningDiT-XL/1 在潜空间训练扩展至 800 Epoch 时，在第 384 Epoch 取得极致的峰值 FID 1.22，最终收敛至 1.34，展示出 \(k\)-Diff 极强的长周期优化鲁棒性。

亮点与洞察¶

极简线性模型的深刻理论洞察：作者并未陷入复杂的非线性深层网络黑箱，而是返璞归真地通过一个单层白化线性扩散模型，将流形切空间与余空间的微分流动力学彻底剥离，揭示了“数据恢复”与“噪声消除”的本质权衡。
第一性原理统一预测目标三岔路：给出 \(k^* = D / (D + d)\) 解析解，终结了学术界多年来关于“流匹配为什么在像素空间必须预测 \(\boldsymbol{x}\)、在潜空间为什么用 \(\boldsymbol{v}\)”的经验猜想，建立了首个严格的量化几何理论桥梁。
极低落地代价的工程实用性：提出的 \(k\)-Diff 框架无需估计极难求解的数据内在流形维度 \(d\)，仅用 1 个单一浮点参数的反向传播便实现了对任意未探索高维空间的即插即用自适应，对工业界扩散模型开发极具实用价值。

局限与展望¶

理论假设与复杂非线性的差距：推导建立在全局线性正交投影 \(\boldsymbol{x} = \boldsymbol{P}\tilde{\boldsymbol{x}}\) 和白化统计假设上，未全面涵盖非线性曲率极大的真实复杂数据流形。
时变目标优化困境尚未完全破解：尽管本文发现恒定标量 \(k\) 优于经验分段 \(k(t)\)，但理论上高曲率流形在不同扩散时间步下的切空间曲率与有效维度应当是动态变化的，如何设计不破坏收敛稳定性的连续时变目标仍待探索。
拓展至多模态与离散扩散：现有工作验证集中在连续图像空间（潜空间与像素空间），对于视频时空流形、三维点云流形或跨模态对齐表征中的几何维度自适应仍有待验证。

评分¶

新颖性: ⭐⭐⭐⭐⭐（首次给出解析解 \(k^* = D/(D+d)\)，理论视角极其优美且直击本质）
实验充分度: ⭐⭐⭐⭐⭐（横跨潜空间与像素空间，覆盖 256 与 512 分辨率，具备严格的长轮次与参数消融）
写作质量: ⭐⭐⭐⭐⭐（数学推导严密清晰，附录模式展开完备，图表直观且论证层层递进）
价值: ⭐⭐⭐⭐⭐（为扩散模型最核心的超参数选择提供了坚实的理论基准和极简工程落地方案）