跳转至

tunable latent generative priors for compressed sensing

会议: NeurIPS 2026
arXiv: 2603.07357
论文: https://arxiv.org/abs/2603.07357
领域: 模型压缩
关键词: 生成先验, 逆问题求解, 嵌套Dropout, 潜变量扩散, 压缩感知

一句话总结

针对传统生成先验在逆问题求解中潜空间复杂度固定且容易过拟合噪声的缺陷,本文利用嵌套Dropout训练出单模型内多尺度有序排布的可调生成先验,使下游压缩感知与逆问题能在推理时自由选择最优截断维度 \(k\),并在理论与多架构上验证了中间维度显著超越全尺寸固定基线。

研究背景与动机

深度生成模型(如 GAN、流模型、变分自编码器及扩散模型)已广泛作为无监督物理信号先验用于求解各类逆问题(如压缩感知、图像去噪、修补、超分辨率及相位恢复)。传统主流范式通常遵循“先独立训练生成模型、后冻结参数执行逆问题求解”的两阶段路线。在这一框架下,生成先验的模型复杂度(即表征能力与潜变量维度)在训练阶段即被永久固化:GAN 拥有固定低维潜空间,而正规化流和扩散模型则在与全图同等甚至更高维度的潜空间上维护概率密度。然而,逆问题的病态程度与测量噪声水平在不同实际应用中差异巨大。

这种固化复杂度的范式面临根本性矛盾:对于严重欠采样或高噪声的逆问题,潜空间容量过大会使重建算法过度拟合测量噪声或漂移到无意义的虚假高频结构;而对于测量充分的观测,固定过小的潜空间又会引入巨大的表示误差(Representation Error)。过往研究若想适配不同测量条件,往往必须针对多个不同的潜空间维度分别独立训练多套模型,带来极高的训练计算开销与存储成本。

本文的核心切入点是打破生成先验的“单固定复杂度”假设,让单个生成模型在训练时自适应学习出有序层级的潜表征,使得用户在推理阶段无需重新训练即可根据具体的正向算子与噪声水平自由选取最合适的有效潜变量维度 \(k\)。核心 idea:利用嵌套 Dropout(Nested Dropout)在单次训练中赋予潜空间坐标显式重要性排序,构建出涵盖 VAE、流模型与潜扩散模型的可调生成先验(Tunable Latent Generative Priors),不仅在线性去噪场景下推导出了最优复杂度截断的解析闭式解,而且在压缩感知等多元逆问题中实现了中间有效维度超越固定全量基线的显著增益。

方法详解

整体框架

本文提出的可调生成先验框架覆盖了变分自编码器(VAE)、正规化流(Normalizing Flows, NFs)与潜变量扩散模型(Latent Diffusion Models, LDMs)。以核心的潜变量扩散模型为例,框架分为两阶段训练与下游逆问题调优求解。在训练阶段,首先训练一个在不同潜截断维度下均保持重建鲁棒性的 VAE 骨干;随后在潜变量空间训练去噪 U-Net,通过凸组合插值标准扩散目标与嵌套截断扩散目标,使潜坐标天然形成“前缀捕获全局主干、后缀刻画细粒度细节”的有序金字塔。在下游求解阶段,针对给定的前向测量算子与噪声水平,直接指定截断维度 \(k\),在交替去噪与数据一致性投影的过程中持续将潜向量强制截断至前 \(k\) 维。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入真实数据 x ~ p_data"] --> B["嵌套Dropout自编码器训练<br/>前缀潜变量强制重构与感知损失"]
    B --> C["潜扩散模型插值联合训练<br/>标准扩散损失与截断潜向量去噪损失凸组合"]
    C --> D["推理时指定截断维度 k<br/>无需重训,前缀保留与后缀清零"]
    D --> E["交替先验去噪与数据一致性校正<br/>反向扩散更新 + 梯度/最小化投影"]
    E --> F["潜变量前 k 维投影截断<br/>(z)↓k 滤除高频噪声过拟合"]
    F --> G["解码输出高质量重构信号 x = D(z_0)"]

关键设计

1. 基于嵌套Dropout的有序层级潜空间构建:单模型承载多尺度复杂度

为了避免为每个复杂度独立训练模型的昂贵开销,本文引入嵌套Dropout(Nested Dropout)。设潜空间总维度为 \(d\),从截断几何分布 \(p_k\) 中采样保留维度 \(k \in \{1, \dots, d\}\),定义前缀截断算子 \((\bm{z})_{\downarrow k} = [\bm{z}_1, \dots, \bm{z}_k, 0, \dots, 0]\)。在 VAE 阶段,对损失函数引入截断重构正则项: $\(\mathcal{L}_{\text{VAE}} = \min_{E, D} \max_{\mathcal{C}} \left[ \mathcal{L}_{\text{rec}}(\bm{x}, \mathcal{D}(\mathcal{E}(\bm{x}))) + \lambda_{\text{reg}}\mathcal{L}_{\text{reg}}(\mathcal{E}(\bm{x})) + \lambda_{\text{adv}}\mathcal{L}_{\text{adv}}(\mathcal{C}, \bm{x}, \mathcal{E}(\bm{x})) + \lambda_{\text{drop}}\mathcal{L}_{\text{drop}}(\bm{x}, \mathcal{D}(\mathcal{E}(\bm{x})_{\downarrow k})) \right]\)$ 该设计迫使编码器将最关键的语义能量集中在靠前的坐标维度,只有当维度递增时才补充更精细的纹理残差,实现了单套参数对任意前缀容量的自适应兼容。

2. 潜扩散去噪目标的凸组合插值:保证全谱段与截断状态下的去噪一致性

直接截断潜变量会导致普通预训练去噪网络出现特征域偏移崩溃。为此,本文在训练潜扩散模型的去噪网络 \(\bm{\epsilon}_\theta\) 时,提出了融合原始全维目标与截断目标的凸组合优化策略: $\(\mathcal{L}_{\text{LDM}} = \mathbb{E}_{\mathcal{E}(\bm{x}), \bm{\epsilon}\sim\mathcal{N}(0, \bm{I}), t\sim\mathcal{U}(\{1,\dots,T\})} \left[ (1-\lambda) \|\bm{\epsilon} - \bm{\epsilon}_\theta(\bm{z}_t, t)\|_2^2 + \lambda \|\bm{\epsilon} - \bm{\epsilon}_\theta((\bm{z}_t)_{\downarrow k}, t)\|_2^2 \right]\)$ 其中 \(\lambda \in [0, 1]\) 平衡全量生成精度与截断去噪能力,\(k\) 同样按分布 \(p_k\) 动态采样。这种设计保证了去噪网络不仅能够对完整潜编码精准去噪,而且在下游任意强行抹去高维分量时,依然能维持稳定可靠的逆向条件得分估计。

3. 结合截断算子的逆问题交替后验采样:自适应平抑测量噪声

在执行逆问题求解时(如测量模型 \(\bm{y} = \mathcal{A}(\bm{x}) + \bm{\eta}\)),算法将潜变量更新显式解耦为“先验去噪生成”、“数据一致性更新”与“空间截断”三部曲。在每个时间步 \(t\),首先根据 DDIM/DDPM 计算去噪估计 \(\hat{\bm{z}}_0\) 并生成提议状态 \(\bm{z}'_{t-1}\);随后利用优化或梯度步校正数据一致性: $\(\bm{z}_{t-1} \leftarrow \arg\min_{\bm{z}} \left( \|\bm{y} - \mathcal{A}(\mathcal{D}(\bm{z}))\|_2^2 + \frac{\|\bm{z} - \bm{z}'_{t-1}\|_2^2}{2\tau_t^2} \right)\)$ 最后执行强制截断 \(\bm{z}_{t-1} \leftarrow (\bm{z}_{t-1})_{\downarrow k}\)。截断操作充当了自适应低通正则化器,直接在流形切空间中滤除因前向算子求逆不适定而放大带入的高维噪声分量。

4. 线性生成先验最优复杂度的解析闭式解:阐明 \(k\) 与噪声能量的本质关联

为了从理论上揭示为何中间维度最优,论文探讨了线性生成模型 \(\bm{G}_n = \bm{U}\bm{\Sigma}\bm{V}^\top\)(奇异值为 \(s_1 \geq s_2 \geq \dots \geq s_n > 0\))在高斯白噪声 \(\bm{\eta} \sim \mathcal{N}(0, \sigma^2 \bm{I}_n)\) 下的去噪恢复问题。在对截断模型 \(\bm{G}_k\) 进行最大似然估计(MLE)时,精确均方误差推导为: $\(\mathbb{E}_{\bm{x}_0, \bm{\eta}} \left[ \|\hat{\bm{x}}(k) - \bm{x}_0\|^2 \right] = k\sigma^2 + \sum_{j=k+1}^n s_j^2\)$ 当 \(s_1 \geq \sigma\) 时,使重构均方误差最小的最优复杂度 \(k^*\) 解析闭式解恰好为: $\(k^* = \arg\min_{k \in [n]} \mathbb{E}[\|\hat{\bm{x}}(k) - \bm{x}_0\|^2] = \max \{k \mid s_k \geq \sigma\}\)$ 该结论清晰证明:最优复杂度取决于信号频谱与噪声强度的竞争关系——仅保留奇异值大于噪声方差的有效信号方向,盲目采用全维度(\(k=n\))会导致 \(n\sigma^2\) 的巨大方差累积,这为深度非线性模型中“中间潜维度最优”的 U 型误差曲线提供了坚实的理论解释。

损失函数 / 训练策略

对于 FFHQ(\(256 \times 256 \times 3\))与 ImageNet,VAE 使用 \(24 \times 64 \times 64\)(总维数 98304)潜空间,带有权重为 0.1 的 VGG 感知损失、\(\lambda_{\text{drop}} = 0.1\) 的嵌套Dropout正则化及对抗判别器。潜扩散模型 U-Net 采用 min-SNR-\(\gamma\)(\(\gamma=5\))加权,使用 AdamW 优化器(学习率 \(10^{-4}\),余弦退火)。在下游逆问题中,使用 DDIM 反向采样(步数 500,随机性参数 \(\eta=0.1\)),每个时间步内嵌 3 步 Adam 优化器完成数据一致性微调,截断参数 \(k\) 在微型验证集(12 张图像)上根据 LPIPS 离线网格选定。

实验关键数据

主实验

在 FFHQ 256×256 验证集上,评估五类代表性逆问题(压缩感知 CS 采样率 10%、随机修复 Inpainting 缺失 85%、超分辨率 4×、相位恢复 Phase Retrieval 12.5% 及高动态截断 HDR,均注入 \(\sigma=0.05\) 测量噪声)。对比固定维度基线与引入可调机制的对应算法:

逆问题任务 评价指标 ReSample 基线 Tunable ReSample (本文) Flow ICTM 基线 Tunable Flow ICTM (本文)
Compressed Sensing (10%) PSNR ↑
LPIPS ↓
24.49
0.51
28.43 (+3.94)
0.28 (-0.23)
22.98
0.57
26.13 (+3.15)
0.41 (-0.16)
Inpainting (85%) PSNR ↑
LPIPS ↓
22.26
0.48
24.84 (+2.58)
0.39 (-0.09)
21.62
0.58
22.60 (+0.98)
0.51 (-0.07)
Super Resolution (4×) PSNR ↑
LPIPS ↓
22.43
0.44
24.53 (+2.10)
0.39 (-0.05)
20.91
0.59
21.38 (+0.47)
0.56 (-0.03)
Phase Retrieval (12.5%) PSNR ↑
LPIPS ↓
22.68
0.58
28.51 (+5.83)
0.29 (-0.29)
—
—
—
—
HDR (scale 2.0) PSNR ↑
LPIPS ↓
20.53
0.38
20.70 (+0.17)
0.36 (-0.02)
—
—
—
—

消融实验

论文在 CelebA-HQ 上针对不同模型族(扩散模型 DPS、潜扩散模型 PSLD 与正规化流 NF)进行了模块化可调对比(主指标为 PSNR ↑ 与 LPIPS ↓):

架构 / 算法配置 压缩感知 (CS) PSNR / LPIPS 随机修复 (Inpaint) PSNR / LPIPS 高斯去模糊 (Deblur) PSNR / LPIPS 说明
NF (标准全维基线) 21.31 / 0.508 20.07 / 0.376 21.43 / 0.438 固化潜空间易过拟合噪声
Tunable NF (可调前缀) 27.16 / 0.238 23.24 / 0.167 24.85 / 0.215 截断中间维度,PSNR跃升5.8dB
PSLD (标准全维基线) 23.30 / 0.248 25.85 / 0.136 27.33 / 0.165 梯度后验采样基线
Tunable PSLD (可调潜变量) 24.69 / 0.164 26.75 / 0.095 28.50 / 0.115 保持原有采样流程,仅约束有效秩
Tunable LDM Prior (本文Alg 2) 25.49 / 0.179 25.30 / 0.122 29.52 / 0.085 结合隐空间硬一致性交替优化

关键发现

  • 非单调 U 型误差曲线普遍存在:跨流模型、VAE 和扩散模型的所有测试显示,重建误差随保留潜维度 \(k\) 呈典型的下凹曲线。在极小采样比(如 CS 10%)下,中间截断维度(如保留 40%~60% 潜变量)取得最佳感知与数值重构,完全保留 \(100\%\) 维度反而因噪声过拟合导致感知质量严重恶化。
  • 与显式优化求解器协同效果最佳:对于 ReSample 和基于流的 MAP 这类在每步显式求解测量一致性最小化的算法,可调潜先验带来的增益极其显著(PSNR 提升高达 3~6 dB);而对于纯一阶得分近似的梯度引导后验采样(如 LatentDAPS),效果相对更平缓或偶有起伏。
  • 最优维度随测量数与噪声单调平移:实验观测与线性理论高度一致——当测量数 \(m\) 增加或测量噪声 \(\sigma\) 降低时,经验最优 \(k^*\) 逐渐向高维平移;高噪声、少测量场景下最优 \(k^*\) 急剧收缩。

亮点与洞察

  • 将模型压缩技术转化为逆问题的连续正则化杠杆:将原本用于模型存储与传输的嵌套Dropout结构创新性地迁移为生成先验的自适应复杂度调节器,无需额外参数即赋予了模型测试时连续缩放能力。
  • 打通了线性理论与深度架构的对应机制:通过严格证明高斯线性去噪下截断维度的闭式解 \(s_k \geq \sigma\),深刻阐明了高维非线性生成模型在逆问题中产生“过拟合高频噪声”的根源,为工程实践中超参选择提供了数学直觉。
  • 极佳的算法通用性与即插即用特性:提出的可调框架无缝兼容现有的主流扩散反向采样器(DDIM、DPS、PSLD、ReSample)及正规化流,仅需在每步迭代加入一次坐标置零截断操作即可激活性能收益。

局限与展望

  • 理论推导局限于线性高斯模型:当前的闭式最优复杂度准则基于线性正交生成矩阵,未能涵盖非线性深度卷积/注意力骨干与非高斯、盲逆问题(Blind Inverse Problems)。
  • 最优截断维度依赖验证集离线搜索:目前 \(k\) 仍作为超参数在含真值的验证子集上选取,尚未实现测试时完全无监督、基于残差收敛曲线或测量谱密度的自适应在线确定。
  • 实验分辨率受限:主实验最高验证至 \(256 \times 256\) 分辨率,未在百万像素级高分辨率扩散模型(如 SDXL/Flux)超大潜空间下评估其截断行为与计算加速潜力。

相关工作与启发

  • vs Bora et al. (ICML 2017) [4]: 经典生成压缩感知使用固定低维度的 DCGAN,无法表达复杂自然图像的丰富纹理(表示误差过大);本文通过可调潜变量在保留大模型强生成力的同时按需收缩容量。
  • vs PSLD (NeurIPS 2023) [42] / ReSample (ICLR 2024) [44]: 既有扩散逆问题方法均在固定的全尺寸潜空间中交替优化,面对高病态算子容易陷入虚假噪声流形;本文直接为其赋能了维度可调能力,实验证明可在原有求解器上实现即插即用的显著涨点。
  • vs Bekasov & Murray (2020) [3]: 前人工作仅探讨了流模型在纯无监督密度估计下的嵌套排序;本文首次将其拓展至大规模潜变量扩散模型并全面系统地攻克线性与非线性逆问题。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统性提出在逆问题中利用嵌套Dropout调节生成先验潜维度,立意独到且兼具深度理论与实证分析。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 VAE、Normalizing Flows、LDM 三大架构,跨越 5 个经典数据集与 5 类主流逆问题,正负对比全面。
  • 写作质量: ⭐⭐⭐⭐⭐ 叙述逻辑严密,定理表述严谨,Mermaid 流程与公式推导自洽清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为解决逆问题中生成先验过拟合噪声的长期难题提供了极具落地可行性的轻量解决方案。