从‘炼丹’到‘理解’：为什么说扩散模型的本质是流形学习？一个玩具实验带来的启发

从‘炼丹’到‘理解’扩散模型与流形学习的本质关联当我们在高维像素空间中训练一个神经网络去预测噪声时是否想过它真正需要学习的是什么近年来扩散模型在图像生成领域取得的突破性进展背后隐藏着一个被忽视的数学本质——流形学习。让我们从一个简单的实验开始假设真实数据存在于一个2维流形中而我们观测到的是它在512维空间中的投影。一个仅有256维隐藏层的MLP网络为何能完美预测干净数据却无法准确预测噪声这个看似矛盾的现象揭示了扩散模型成功的关键所在。1. 流形假设理解扩散模型的数学基础1.1 自然数据的低维本质流形假设认为现实世界中的高维数据如图像实际上分布在一个嵌入高维空间的低维流形上。想象一张人脸照片虽然它由数百万像素组成但所有可能的人脸变化表情、角度、光照等可以用相对较少的参数来描述。这种低维结构是神经网络能够在高维空间中有效工作的关键。在数学上我们可以将这个过程建模为# 假设真实数据是d维的观测数据是D维的d D import numpy as np d 2 # 真实维度 D 512 # 观测维度 P np.random.randn(D, d) # 随机投影矩阵 P P / np.linalg.norm(P, axis0) # 正交化 x_true np.random.randn(d) # 真实数据点 x_observed P x_true # 观测到的高维数据1.2 噪声的高维特性与自然数据不同噪声通常均匀分布在整个高维空间中。下表对比了干净数据与噪声的关键差异特性干净数据噪声维度低维流形全空间结构有规律无规律可预测性高低神经网络学习难度相对容易极其困难这个差异解释了为什么直接预测干净数据x-prediction比预测噪声ε-prediction更容易成功——网络只需要关注低维流形上的结构而不必处理整个高维空间中的随机波动。2. 玩具实验验证流形学习假说2.1 实验设计与实现为了验证上述理论我们设计了一个简单的合成实验生成d2维的真实数据点通过随机投影矩阵P将其映射到D维空间D ∈ {2,8,16,512}训练一个MLP网络预测干净数据或噪声评估不同维度下的预测效果实验使用的网络结构非常简单import torch import torch.nn as nn class SimpleMLP(nn.Module): def __init__(self, input_dim, hidden_dim256): super().__init__() self.net nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, input_dim) ) def forward(self, x, t): return self.net(x)2.2 实验结果与分析当D512而网络隐藏层只有256维时我们观察到以下现象x-prediction尽管网络容量不足仍能准确预测干净数据ε-prediction预测噪声完全失败v-prediction预测速度介于x和ε之间表现中等注意这里的容量不足是带引号的因为对于低维流形上的数据网络实际上有足够的表达能力。这个实验验证了我们的核心观点扩散模型的本质是学习数据流形的结构而非记忆高维空间中的所有细节。当网络直接预测干净数据时它只需要关注流形上的信息而预测噪声则需要处理整个高维空间这对有限容量的网络来说是不现实的。3. 从理论到实践Just Image Transformers3.1 传统扩散模型的局限当前主流的扩散模型通常采用以下两种预测方式ε-prediction预测噪声如DDPMv-prediction预测速度如Progressive Distillation然而这两种方式都要求网络处理高维噪声信息这与我们观察到的流形学习本质相矛盾。更合理的做法是让网络直接预测干净数据x-prediction因为它更符合数据的低维本质。3.2 Just Image Transformers (JiT)设计基于这一洞察我们提出了Just Image Transformers方法将图像分割为p×p的块如16×16每个块直接作为Transformer的输入token网络直接预测干净图像块使用简单的MSE损失训练这种方法的关键优势在于无需潜在空间直接在像素空间操作无需分词器避免引入额外复杂性无需辅助损失单一目标函数简化训练下表比较了不同预测方式在ImageNet 256×256上的表现预测方式FID (256×256)训练稳定性所需网络容量x-prediction3.21高低ε-prediction失败低极高v-prediction15.67中高3.3 瓶颈设计的意外优势与传统认知相反我们在实验中发现了瓶颈结构的积极作用在patch嵌入层引入降维如从768维降到32维然后再扩展到Transformer的隐藏维度这种设计反而提升了生成质量可能的解释是瓶颈迫使网络专注于学习数据流形的本质特征过滤掉无关的高维噪声。这与流形学习的核心理念高度一致——找到数据的内在低维结构。4. 超越图像生成流形学习的普适意义4.1 对其他领域的启示流形学习的视角不仅适用于图像生成还可以推广到音频生成语音和音乐也存在于低维流形上视频生成时间连续性进一步降低了有效维度分子设计化学空间虽然庞大但有效分子占据很小的子空间4.2 未来研究方向基于这一理论框架我们认为有以下值得探索的方向流形感知的架构设计显式建模数据流形的网络结构自适应降维根据数据复杂度动态调整网络容量跨域流形学习利用一个领域的流形知识加速另一领域的学习在ImageNet 512×512分辨率下我们的JiT模型仅用基础配置就达到了FID 5.34证明了这一思路的潜力。更令人惊讶的是当我们将patch大小增加到32×32每个token 3072维而Transformer隐藏层保持768维时模型仍然表现良好——这正是因为网络不需要处理全空间只需关注流形上的结构。提示在实际应用中x-prediction通常需要调整噪声调度使网络在训练时看到更多中等噪声水平的样本这有助于学习流形结构。从炼丹式的调参到真正理解扩散模型的工作原理流形学习提供了一个强有力的理论框架。它不仅解释了现有方法的成功更为设计下一代生成模型指明了方向——不是盲目增加网络容量而是更好地理解和利用数据的底层结构。