从‘预测噪声’到‘预测原图’：一个被忽视的扩散模型设计选择，如何让大块ViT在像素空间起飞？

weixin_30374009

502人浏览 · 2026-03-27 02:27:11

weixin_30374009 · 2026-03-27 02:27:11 发布

从预测噪声到预测原图：扩散模型设计选择的革命性突破

在生成式AI领域，扩散模型已成为图像合成的中流砥柱。然而，一个被大多数实践者忽视的设计选择——网络输出目标的选择——正在悄然改变着这一领域的技术格局。本文将深入探讨为何让模型直接预测干净图像（x-预测）而非噪声（ε-预测）或速度（v-预测），能够使简单的视觉Transformer（ViT）在像素空间实现惊人表现。

1. 扩散模型预测目标的根本性选择

扩散模型的核心思想是通过逐步去噪过程生成数据。传统方法通常让网络预测噪声（ε-预测）或速度场（v-预测），而x-预测则直接让网络输出干净数据。这三种看似等价的选择在实践中却产生截然不同的效果。

关键差异对比表：

预测类型	数学表达式	流形特性	网络容量需求
x-预测	直接输出x	位于低维流形	较低
ε-预测	预测噪声ε	分布在高维空间	极高
v-预测	预测速度v=x-ε	脱离流形	极高

在ImageNet 256×256分辨率实验中，使用768维块（16×16×3）的ViT模型表现如下：

x-预测：FID 2.1（优秀）
ε-预测：FID >100（完全失败）
v-预测：FID >100（完全失败）

这一现象揭示了流形假设的关键作用：自然数据本质上位于低维流形上，而噪声和速度则分布在整个高维空间。让网络预测低维流形上的数据（x）远比预测高维空间中的噪声或速度要简单得多。

2. 大块ViT在像素空间的惊人表现

传统观点认为，处理高维像素数据需要复杂的U-Net架构或潜空间表示。然而，x-预测的引入彻底改变了这一认知。

JiT（Just Image Transformer）架构优势：

极简设计：纯Transformer架构，无特殊模块
大块处理：直接处理16×16或32×32像素块
自包含：无需预训练、辅助损失或分词器

在512×512分辨率下，使用32×32块（3072维）的JiT模型实现了：

FID-50K: 1.78
Inception Score: 253.7

这一表现超越了多数需要复杂预训练的潜空间扩散模型，证明了x-预测在高维像素空间的有效性。

3. 网络容量与数据维度的深层关系

x-预测的成功源于其对网络容量的高效利用。当数据维度远高于网络隐藏维度时：

ε/v-预测的困境：
- 需要网络保留所有噪声信息
- 导致信息过载和训练不稳定
x-预测的优势：
- 只需保留低维流形信息
- 自然实现信息过滤和降维

实验验证： 在玩具实验中，当观测维度D从2增加到512时：

x-预测保持稳定性能
ε/v-预测在D=16时开始失效，D=512时完全崩溃

4. 实践指南：如何应用x-预测

对于希望在实际项目中应用这些发现的开发者，以下是关键实施步骤：

网络输出修改：

# 传统ε-预测
epsilon = model(noisy_img, t)
clean_img = (noisy_img - sigma * epsilon) / alpha

# 改为x-预测
clean_img = model(noisy_img, t)  # 直接输出干净图像

损失函数调整：

使用v-空间损失但保持x-预测

def loss_fn(x_pred, x_true, noisy, t):
    v_pred = (x_pred - noisy) / (1 - t)
    v_true = (x_true - noisy) / (1 - t)
    return mse_loss(v_pred, v_true)

架构选择建议：
- 优先考虑Transformer而非U-Net
- 大胆使用大块尺寸（16×16或更大）
- 可尝试瓶颈设计（隐藏维度<块维度）

5. 跨领域应用的广阔前景

x-预测的优势在非图像领域尤为显著：

蛋白质结构预测：直接建模3D坐标
分子生成：处理原始原子位置
气象数据：处理高维网格数据

这些领域通常缺乏现成的分词器或潜空间表示，而x-预测的ViT提供了一种通用解决方案。例如，在分子生成中：

# 分子坐标生成示例
def generate_molecule():
    # 初始化随机噪声
    z = torch.randn(1, num_atoms*3)
    
    # 扩散过程
    for t in reversed(range(steps)):
        # x-预测直接输出分子坐标
        coords = model(z, t)
        z = update_step(z, coords)
    
    return coords