∞-Brush: Controllable Large Image Synthesis with Diffusion Models in Infinite Dimensions¶

会议: ECCV 2024
arXiv: 2407.14709
代码: https://github.com/cvlab-stonybrook/infinity-brush (有)
领域: 图像生成 / 扩散模型 / 函数空间生成
关键词: [无限维扩散模型, 函数空间, 大图像生成, 交叉注意力神经算子, 可控生成]

一句话总结¶

提出首个在无限维函数空间中的条件扩散模型 ∞-Brush，通过交叉注意力神经算子实现可控条件生成，仅用 0.4% 像素训练即可在任意分辨率（最高 4096×4096）上生成保持全局结构的大图像。

背景与动机¶

高分辨率大图像生成在数字病理学和遥感等领域有重要需求，但现有方法面临两大瓶颈：(1) 像素/潜空间扩散模型（如 SDXL）无法超越训练分辨率，且计算量随分辨率二次增长；(2) 基于 patch 的方法（如 MultiDiffusion、Graikos et al.）虽计算高效，但局限于局部信息，无法捕捉远距离空间依赖关系。而已有的无限维扩散模型（∞-Diff）虽解决了分辨率限制问题，但不支持条件控制生成。

核心问题¶

如何在函数空间中构建条件扩散模型，使其既能生成任意分辨率的大图像、保持全局结构一致性，又能通过外部条件（类别标签、嵌入向量等）实现可控生成？

方法详解¶

整体框架¶

∞-Brush 将图像建模为 Hilbert 空间 \(\mathcal{H} = L^2(\mathcal{X}, \mu)\) 中的连续函数，而非固定分辨率的像素矩阵。核心组件包括：

条件扩散过程：在函数空间中定义前向加噪和条件反向去噪过程，条件信息 \(\mathbf{e}\) 可以是有限维标签或嵌入向量
交叉注意力神经算子（CANO）：在稀疏层实现线性复杂度的跨模态信息融合
层级去噪器：稀疏层（细粒度细节）+ 网格层（全局信息）

关键设计¶

交叉注意力神经算子（CANO）： - 改变 vanilla attention 的计算顺序：先做 \(\tilde{\mathbf{k}}_i^l \odot \mathbf{v}_i^l\) 的逐点乘积，再与 \(\tilde{\mathbf{q}}_t\) 做点积 - 计算复杂度从 \(\mathcal{O}(N^2 d)\) 降为 \(\mathcal{O}((N + \sum_l N_l) d^2)\)，对采样点数 \(N\) 线性 - 同时融合三种条件嵌入：时间步嵌入 \(\mathbf{t}\)、条件嵌入 \(\mathbf{e}\)、坐标嵌入 \(\mathbf{c}\)

层级去噪器架构： - 稀疏层（蓝色）：依次使用稀疏神经算子 → CANO → 自注意力算子，处理随机采样的稀疏坐标点，捕捉细粒度信息 - 网格层（粉色）：通过 k-NN 线性插值将稀疏点转换为规则网格，接入 UNO（基于 UNet 的神经算子）聚合全局信息，在瓶颈层使用 vanilla cross-attention

高效训练策略： - 每次迭代仅随机采样 \(256 \times 256 = 65536\) 个像素（占 4096×4096 图像的 0.4%），而 ∞-Diff 需要 25% - CANO 中的坐标嵌入起到位置编码的作用，使模型即使在极少像素上训练也能推断完整图像

平滑算子：通过截断高斯核 \(\mathbf{A}: \mathcal{H} \to \mathcal{H}\) 将离散像素近似为光滑函数，确保函数空间中的表示合法性。

损失函数 / 训练策略¶

训练目标来源于无限维条件扩散的变分上界最小化，最终简化为：

\[\theta^* = \arg\min_\theta \mathbb{E}_{\mathbf{u}_0 \sim \mathbb{Q}_{\text{data}}, t} \lambda_t \|\mathbf{C}^{-1/2}(\mathbf{A}\boldsymbol{\xi} - \boldsymbol{\xi}_\theta(\sqrt{\bar\alpha_t}\mathbf{A}\mathbf{u}_0 + \sqrt{1-\bar\alpha_t}\mathbf{A}\boldsymbol{\xi}, \mathbf{e}, t))\|_{\mathcal{H}}^2\]

其中 \(\boldsymbol{\xi} \sim \mathcal{N}(\mathbf{0}, \mathbf{C})\) 为协方差算子 \(\mathbf{C}\) 下的高斯噪声，\(\mathbf{A}\) 为平滑算子。核心理论基础是 Feldman-Hájek 定理保证两个高斯测度的等价性，从而使 KL 散度可计算。

优化器：Adam，lr=\(5 \times 10^{-5}\)，\(\beta_1=0.9\)，\(\beta_2=0.99\)
EMA rate: 0.995
推理：DDIM 50步
硬件：4× NVIDIA A100，每 GPU batch size 20
FlashAttention-2 加速 CANO

实验关键数据¶

数据集	分辨率	方法	CLIP FID ↓	Crop FID ↓
CelebA-HQ	1024²	∞-Diff (无条件)	9.44	-
CelebA-HQ	1024²	∞-Brush	8.38	-
BRCA	4096²	Graikos et al.	2.75	11.30
BRCA	4096²	∞-Brush	2.63	14.76
BRCA 5×	1024²	SDXL	6.64	6.98
BRCA 5×	1024²	Graikos et al.	7.43	15.51
BRCA 5×	1024²	∞-Brush	3.74	17.87
NAIP	1024²	SDXL	10.90	11.50
NAIP	1024²	Graikos et al.	6.86	43.76
NAIP	1024²	∞-Brush	6.32	48.65

方法	参数量	1024² epoch 时间	4096² epoch 时间
SDXL	~2.6B	~300 hr	OOM / 不可行
Graikos et al.	~400M	~140 hr	~140 hr
∞-Brush	~78M	~12 hr	~12 hr

下游应用：BACH 测试集准确率从 79%（仅真实数据）提升至 83%（真实+合成数据）。

消融实验要点¶

CANO 的作用（BRCA 4096²）：有 CANO → CLIP FID 2.63 / Crop FID 14.76；无 CANO → 3.81 / 16.28。去除 CANO 后仅在 UNet 瓶颈层做 vanilla cross-attention，细节和全局结构均显著下降
训练像素比例：0.4% vs 1.6% 像素，更多像素提升生成质量，但 0.4% 已得到良好效果
零样本分类：混淆矩阵表明生成图像与文本提示语义一致（benign/in-situ/invasive/normal）

亮点¶

理论贡献扎实：严格推导了函数空间中条件扩散模型的变分目标，利用 Feldman-Hájek 定理、Radon-Nikodym 导数等泛函分析工具保证训练目标的合理性
极致训练效率：0.4% 像素采样训练，78M 参数量仅为 SDXL 的 1/33，训练时间仅为 1/25，且跨分辨率保持恒定
分辨率无关性：同一模型可在任意分辨率推理（256² 到 4096²），真正意义上的"无限维"生成
CANO 设计巧妙：通过改变注意力计算顺序实现线性复杂度，并将坐标嵌入作为隐式位置编码

局限性 / 可改进方向¶

局部细节不足：Crop FID 系统性劣于 patch-based 方法和 SDXL，全局结构好但精细纹理不够
参数量受限：78M 参数远小于竞品（SDXL 2.6B），模型容量不足
从零训练：无法利用已有大规模预训练模型（如 SDXL 在 LAION-5B 上的预训练），冷启动劣势明显
条件信息粒度粗：仅用单个全局嵌入向量作为条件（将 4096² 下采样到 256² 再提取），丢失了大量空间细节。对比 patch-based 方法用 16 个局部条件
专用领域：仅在病理/遥感图像上验证，未在自然图像上测试

与相关工作的对比¶

维度	∞-Brush	∞-Diff	SDXL	Patch-based (Graikos et al.)
函数空间	✅	✅	❌	❌
条件控制	✅	❌	✅	✅
超训练分辨率推理	✅	✅	❌	有限
全局结构保持	✅ 最优	-	中等	弱（远距离依赖差）
局部细节	较弱	-	强	较强
预训练利用	❌	❌	✅ (LAION-5B)	✅
最大分辨率	4096²	理论任意	1024²	~2048²

启发与关联¶

函数空间扩散 → 视频/3D：将函数空间表示扩展到时间维度（视频作为时空函数）或 3D 场景（NeRF 式连续表示），可能实现任意时长/分辨率的生成
知识迁移：论文自身提到的未来工作——从有限维预训练模型向无限维模型的知识迁移，可以结合蒸馏、adapter 等技术
多尺度条件：当前仅用全局嵌入的单一条件是性能瓶颈，设计分层条件（全局 + 区域 + 局部），类似 ControlNet 的多层条件注入
与 Neural Operator 社区的交叉：CANO 的设计思路可以反向应用于 PDE 求解等科学计算场景的条件控制

评分¶

新颖性: ⭐⭐⭐⭐⭐（首个函数空间条件扩散模型，理论推导完整）
实验充分度: ⭐⭐⭐⭐（多数据集、消融充分，但缺少自然图像实验）
写作质量: ⭐⭐⭐⭐⭐（理论推导清晰、结构完整、补充材料详尽）
价值: ⭐⭐⭐⭐（开辟新范式，但局部细节不足限制实用性）