跳转至

Efficient Image Synthesis with Sphere Latent Encoder

会议: NeurIPS 2026
arXiv: 2605.15592
项目主页: https://sphere-latent-encoder.github.io
领域: 图像生成
关键词: 图像合成, 极少步采样, 超球面隐空间, 表征自编码器, 解耦生成架构

一句话总结

提出 Sphere Latent Encoder 框架,通过引入预训练表征自编码器将重建与生成解耦,在纯超球面隐空间内完成去噪与采样迭代,避免了像素与隐空间频繁往复计算,以极低计算代价实现了高质量的极少步图像生成。

研究背景与动机

流匹配(Flow Matching)与扩散模型为现代高分辨率图像合成奠定了扎实基础,但在标准推理时依赖数十乃至上百步的迭代数值求解,带来巨大的计算开销。为提升推理效率,Consistency Models 与 MeanFlow 等少步生成方案相继被提出,力求在个位数步数甚至单步内完成图像采样。然而,这些方法通常面临严重的训练不稳定性与超参数敏感性,尤其是 MeanFlow 等框架涉及雅可比-向量积(JVP)一阶近似计算,优化信号极易产生冲突与退化,制约了少步生成模型的扩展性与训练鲁棒性。

近期提出的 Sphere Encoder 另辟蹊径,抛弃了变分自编码器(VAE)中易致后验坍塌的 KL 散度约束,改用 RMSNorm 将潜变量投影至高维超球面上,仅需 4 到 8 步高斯噪声演化即可合成高保真图像。然而,Sphere Encoder 存在两个显著瓶颈:其一,其采样与训练过程需在像素空间与潜空间之间频繁反复穿梭(即潜在变量解码至图像后再重新编码为潜在变量),每次迭代都消耗巨额的编码器与解码器 FLOPs;其二,它在单一编码器-解码器架构中联合优化重建与生成目标,而先验研究表明重建要求保留底层纹理细节,生成则倚赖紧凑的高阶语义表征,两者在优化目标上存在固有的目标冲突。当噪声水平较低时重建变好,超球面覆盖度却严重不足,导致模型不得不使用数十亿参数的庞大网络来勉强维系生成多样性与保真度。

面对上述矛盾,本文提出彻底将“像素重建”与“潜在生成”解耦,在纯超球面隐空间中完成全部生成迭代。核心 idea:利用冻结的高保真表征自编码器将像素映射解耦到固定边界,构建独立的潜在去噪网络全流程在超球面隐空间内实施单向多步细化,仅在采样末尾单次解码成图。

方法详解

整体框架

Sphere Latent Encoder 将图像合成管线清晰地解耦为两部分:一个是预训练且权重冻结的表征自编码器(Representation Autoencoder, RAE),负责在像素空间与高维隐表征之间建立高保真映射;另一个是专门在超球面隐空间上训练的 Transformer 去噪模型(基于 SiT 架构),负责完成从随机高斯噪声到干净语义潜在编码的平滑生成。输入图像首先由冻结的编码器投影到连续潜在空间,经球面归一化后作为监督信号;而在采样推理时,去噪网络直接在潜在空间内部逐级细化,直到完成预设步数后,仅调用一次冻结解码器输出最终图像。

整个方法的处理流程如下所示:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["高斯噪声潜在编码<br/>z ~ N(0, I)"] --> B["超球面投影<br/>RMSNorm 归一化 v"]
    B --> C["潜在去噪网络预测<br/>SiT 架构 + CFG 引导"]
    C --> D["衰减噪声扰动与再球面化<br/>v' = F(z_guided)"]
    D -->|未达到步数 T| B
    D -->|达到指定步数 T| E["单次像素解码<br/>RAE ViT 解码器"]
    E --> F["输出高保真图像"]

关键设计

1. 解耦生成与重建:固定表征自编码器与独立隐空间去噪

针对原生 Sphere Encoder 中重建与生成互相干扰的目标冲突,本文选用基于 DINOv2-B 编码器与 ViT-XL 解码器构建的预训练表征自编码器(RAE)作为底层视觉分词器。该自编码器具备两个关键特性:一是 DINOv2 提取的特征蕴含极其丰富的全局判别与语义结构,为生成模型提供了结构分明且连续的隐空间;二是其 ViT 解码器仅需 128 GFLOPs 即可达成 0.49 的高保真重建 FID(rFID)。在训练阶段,编码器与解码器全部保持冻结,去噪网络仅专注于学习潜表征流形上的概率转移,使得重建的像素保真度与生成的语义收敛各司其职,彻底消除了低噪声重建与全球面覆盖度之间的权衡困境。

2. 纯潜在域双噪声级监督:重建与局部一致性损失

针对原生方法在像素空间计算 \(L_1\) 与 LPIPS 感知损失所导致的高显存占用与显存吞吐瓶颈,本文将所有监督信号全部移至隐空间内计算。给定输入潜变量 \(z\),经球面化映射 \(\mathcal{F}(z) = \text{RMSNorm}(z)\) 得到 \(v\)。训练时对 \(v\) 注入两个不同强度的独立高斯噪声并重新球面化,构建高噪声样本 \(v_{\text{NOISY}} = \mathcal{F}(v + \sigma \epsilon)\) 与低噪声样本 \(v_{\text{noisy}} = \mathcal{F}(v + \sigma_{\text{sub}} \epsilon)\),其中 \(\sigma \ge \sigma_{\text{sub}}\)。模型通过以下两项核心损失联合监督去噪网络 \(\mathcal{G}\):

\[ \mathcal{L}_{\text{recon}} = \|\mathcal{G}(v_{\text{noisy}}) - z\|_1 + \mathcal{L}_{\text{cos}}(\mathcal{G}(v_{\text{noisy}}), z) \]
\[ \mathcal{L}_{\text{cons}} = \|\mathcal{G}(v_{\text{NOISY}}) - \text{sg}(\mathcal{G}(v_{\text{noisy}}))\|_1 + \mathcal{L}_{\text{cos}}(\mathcal{G}(v_{\text{NOISY}}), \text{sg}(\mathcal{G}(v_{\text{noisy}}))) \]

其中 \(\text{sg}(\cdot)\) 表示阻断梯度反向传播(stop-gradient)。\(\mathcal{L}_{\text{recon}}\) 引导模型从低噪声球面点准确预测真实隐编码;\(\mathcal{L}_{\text{cons}}\) 将低噪声分支的预测作为固定目标,强迫高噪声扰动的输出向其对齐,从而在超球面几何流形上构建起局部的平滑过渡与拓扑一致性,确保邻近区域对应于语义连贯的视觉内容。

3. 对数正态自适应噪声采样与完全剥离潜在循环损失

原生 Sphere Encoder 采用保守的均匀噪声分布,并额外引入一项需要“去噪-再加噪-再次去噪”的循环潜在一致性损失(Latent Consistency Loss),导致每步训练迭代耗时激增 67%(从 201 ms 升至 336 ms)。本文发现,在专职去噪器与强判别隐空间的支撑下,循环一致性损失在训练初期反而会因为未充分收敛的粗糙预测带来误导性约束。因此,本文彻底移除了该循环损失。同时,噪声级别 \(\sigma\) 与 \(\sigma_{\text{sub}}\) 改为从对数正态分布 \(\text{LogNorm}(+0.4, 1.0)\) 中独立采样两组数值后取较大者为 \(\sigma\)、较小者为 \(\sigma_{\text{sub}}\)。这种非对称的对数正态采样让网络在更宽阔的噪声跨度下得到充分训练,显著强化了网络在低步数大步长场景下的全局流形去噪能力。

4. 纯潜域渐进衰减迭代采样机制

在测试推理阶段,模型完全避免了任何中间步骤的像素空间解码与重新编码。采样从纯高斯噪声 \(z \sim \mathcal{N}(0, I)\) 出发,在每一步 \(t \in [0, T-1]\),先将其投影为超球面向量 \(v = \mathcal{F}(z)\),随后通过无分类器引导(Classifier-Free Guidance, CFG)获取条件增强后的预测隐变量:

\[ z_{\text{guided}} = \mathcal{G}(v, \varnothing) + \omega (\mathcal{G}(v, y) - \mathcal{G}(v, \varnothing)) \]

随后计算扰动衰减系数 \(r = (1 - \frac{t+1}{T})^\gamma\)(ImageNet 上取 \(\gamma=0.5\),小数据集上取 \(\gamma=0.75\)),将 \(z_{\text{guided}}\) 再次球面化后注入随步数单调递减的高斯扰动:\(z \leftarrow \mathcal{F}(z_{\text{guided}}) + \epsilon \cdot \sigma_{\max} \cdot r\)。迭代完成后,仅在末尾调用一次 RAE 的 ViT 解码器得到最终高保真图像 \(x = \mathcal{D}(z)\),单步推理相比原生模型节省了高达 85% 的 FLOPs 计算量。

损失函数 / 训练策略

模型总优化目标为重建损失与一致性损失的加权和:

\[ \mathcal{L}_{\text{total}} = 50.0 \cdot \|\mathcal{G}(v_{\text{noisy}}) - z\|_1 + 1.0 \cdot \mathcal{L}_{\text{cos}}(\mathcal{G}(v_{\text{noisy}}), z) + 25.0 \cdot \|\mathcal{G}(v_{\text{NOISY}}) - \text{sg}(\mathcal{G}(v_{\text{noisy}}))\|_1 + 1.0 \cdot \mathcal{L}_{\text{cos}}(\mathcal{G}(v_{\text{NOISY}}), \text{sg}(\mathcal{G}(v_{\text{noisy}}))) \]

在 ImageNet-1K(\(256 \times 256\) 分辨率)上采用 SiT-XL/1 架构(28 层,隐藏维度 1152,16 头,参数量 676M),在 64 张 64GB 显存的 GPU 上以批大小 512 训练 600 个 epoch,学习率采用固定的 \(1 \times 10^{-4}\),优化器选用 AdamW(\(\beta_1=0.9, \beta_2=0.95\)),并启用 EMA 衰减系数 0.9995。在 Animal-Faces 与 Oxford-Flowers 上则采用轻量化的 SiT-B/1 架构(12 层,隐藏维度 768,参数量 130M),在 8 张 GPU 上分别训练 500 与 2000 个 epoch。

实验关键数据

主实验

在 50,000 张随机生成样本的评测标准下,模型在无条件小数据集与大规模类别条件 ImageNet-1K 上均展现了极强的采样质量与计算效率权衡。

数据集 指标 (NFE) 本文 Sphere Encoder 相对提升 / 降幅
Animal-Faces FID (2-step) 10.63 19.29 -8.66 FID (-44.9%)
Animal-Faces FID (4-step) 6.89 18.23 -11.34 FID (-62.2%)
Animal-Faces GFLOPs (4-step) 390 4554 -91.4% 计算量
Oxford-Flowers FID (4-step) 8.61 12.96 -4.35 FID (-33.6%)
Oxford-Flowers GFLOPs (4-step) 567 9118 -93.8% 计算量
ImageNet-1K FID (\(4 \times 2\) NFE) 2.25 4.02 -1.77 FID (-44.0%)
ImageNet-1K FID (\(6 \times 2\) NFE) 2.11 - 接近 250 步 DiT-XL/2 (2.27)
ImageNet-1K CMMD (\(4 \times 2\) NFE) 0.144 0.363 -60.3% CMMD
ImageNet-1K 采样 FLOPs (G) 2969 (6步) 13326 (4步) 步数更多仍节约 77.7% FLOPs

消融实验

在 ImageNet-100 子集与 Animal-Faces 上的核心机制消融结果如下表所示:

消融配置 变量因子 评测 FID 相对首项变化 说明
均匀分布采样 Baseline 噪声分布 6.43 - 原生 Sphere Encoder 采样策略
LogNorm(-0.4, 1.0) 噪声分布 5.56 -13.5% 中高噪声更充分
LogNorm(+0.4, 1.0) 噪声分布 5.31 -17.4% 高噪声区间覆盖更密,性能最优
仅重建损失 (R) 训练目标 8.97 - 缺乏跨噪声流形对齐
重建 + 一致性 (R+C) 训练目标 5.31 -40.8% 局部流形平滑大幅增益
R+C + 循环潜在一致性 (L*) 训练目标 4.82 -46.3% 引入额外复杂约束反倒阻碍微调
R+C* (持续微调 300 ep) 训练目标 4.68 -47.8% 剥离循环约束优化更充分
移除球面化 (w/o spherify) 几何投影 89.68 +1816.2% 欧氏空间直接去噪导致彻底崩塌
完整球面化 (w/ spherify) 几何投影 4.68 - RMSNorm 超球面几何流形必不可少
FLUX VAE (\(d=16\)) 自编码器 172.25 - 低维空间表达力受限且语义贫瘠
GAE (\(d=32\)) 自编码器 23.26 -86.5% 结合视觉基座后有明显提升
RAE DINOv2 (\(d=768\)) 自编码器 10.63 -93.8% 高维深层语义赋予极强去噪能力

关键发现

  • 超球面归一化是模型收敛的决定性前提:去掉 RMSNorm 球面化操作(w/o spherify)后,FID 剧烈暴增至 89.68,生成图像完全崩塌为不可辨认的杂色斑块。这表明在固定半径的超球面流形上规整潜在分布,是使简单去噪模型得以仅凭数步消除高斯噪声的核心几何先验。
  • 自编码器的表征维度与语义深度是生成性能的基石:FLUX VAE 仅有 16 维通道且缺乏显式深层语义监督,在少步生成下彻底失败(FID 172.25);而基于 DINOv2 的 RAE 拥有 768 维高维语义表征,将 FID 推进到 10.63,证明强判别基座模型提取的连续流形极大减轻了少步去噪器的生成负担。
  • 循环潜在一致性损失存在负面优化阻碍:在潜在空间额外添加原生 Sphere Encoder 的循环一致性约束不仅导致每步训练耗时增加 67%,最终 FID(4.82)反而劣于纯粹的重建加一致性组合(4.68)。消除不必要的循环损失有助于梯度直接对齐真实潜在表征。

亮点与洞察

  • 解耦范式打破重建与生成的零和博弈:通过固定高容量表征自编码器作为静态边界,彻底分流了低层像素重建与高层语义生成两类性质迥异的任务,避免了相互制约。
  • 全流程隐空间流形闭环:巧妙地将 Sphere Encoder 的几何概念与 Latent Diffusion 的高效率结合,消除了所有中间步骤的像素往返,以纯潜在运算实现了 6.5 倍以上的推理加速。
  • 无须一阶近似的少步去噪新路径:不同于 MeanFlow 或 Consistency Models 依赖不稳定的 JVP 或脆弱的切线边界匹配,该架构沿用成熟的加噪-去噪监督与一致性停止梯度,训练极其平稳。

局限与展望

  • 条件控制模式目前局限于类别生成:论文当前仅在 ImageNet 类别条件与无条件小数据集上完成实验,尚未拓展至文本到图像(Text-to-Image)或局部编辑等复杂的通用多模态提示场景。
  • 对高质量外部表征自编码器的高度依赖:系统生成性能直接受制于预训练分词器(如 DINOv2)的表征质量与覆盖分布,若自编码器在特定垂直领域存在域外泛化缺陷,生成模型本身难以自行校正。
  • 极限制动下尚未达成真正高质量单步采样:在 1-step(1 NFE)生成上仍逊色于具备显式速度向量匹配的流模型,未来探索自编码器与潜在流模型的联合表征微调有望弥补单步极限性能。

相关工作与启发

  • vs Sphere Encoder: 原生模型在统一网络内交替调用编解码器进行像素-隐变量多次往返,计算极其沉重且易遇重建与生成的冲突;本文将其解耦为冻结 RAE 与纯隐空间 SiT 去噪,以 15% 的计算代价取得大幅领先的图像保真度。
  • vs MeanFlow / α-Flow: 少步流模型依赖雅可比-向量积(JVP)一阶近似求解,训练过程脆弱敏感且容易发散;本文的球面潜在去噪采用标准回归与 stop-gradient 对齐,训练过程如常规扩散模型般平稳易收敛。
  • vs Latent Diffusion / DiT: 传统 LDM 需要 50 到 250 步采样才能收敛(如 DiT-XL/2 在 250 步下 FID 为 2.27);本文仅需 4 到 6 步采样即可达到 2.25 与 2.11 的 FID,将推理延迟降低了一个数量级。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 创造性地将超球面连续流形生成思想平移至解耦潜在表征空间,摆脱了像素反复往返。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖无条件与条件多尺度数据集,包含完备的消融实验、计算复杂度对比与高阶度量指标。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑推导层次清晰,图表完备,方法与动机的对应关系紧凑自洽。
  • 价值: ⭐⭐⭐⭐☆ 为少步高分辨率图像合成提供了一种避开 JVP 复杂优化的极具实用前景的新工程范式。