[论文解读] Online Deep Equilibrium Learning for Regularization by Denoising

Jiaming Liu, Xiaojian Xu+3 more|arXiv (Cornell University)|May 25, 2022

Numerical methods in inverse problems被引用 6

一句话总结

本文提出在线去噪深度平衡正则化（ODER），一种内存高效的框架，通过在成像逆问题的深度平衡模型（DEQ）中对测量模型应用随机近似，实现在线处理。与批量处理相比，ODER通过逐次处理测量数据，在稀疏视角CT、MRI和强度衍射断层扫描应用中显著降低了计算和内存复杂度，同时保持了与批量DEQ相当的重建质量。

ABSTRACT

Plug-and-Play Priors (PnP) and Regularization by Denoising (RED) are widely-used frameworks for solving imaging inverse problems by computing fixed-points of operators combining physical measurement models and learned image priors. While traditional PnP/RED formulations have focused on priors specified using image denoisers, there is a growing interest in learning PnP/RED priors that are end-to-end optimal. The recent Deep Equilibrium Models (DEQ) framework has enabled memory-efficient end-to-end learning of PnP/RED priors by implicitly differentiating through the fixed-point equations without storing intermediate activation values. However, the dependence of the computational/memory complexity of the measurement models in PnP/RED on the total number of measurements leaves DEQ impractical for many imaging applications. We propose ODER as a new strategy for improving the efficiency of DEQ through stochastic approximations of the measurement models. We theoretically analyze ODER giving insights into its convergence and ability to approximate the traditional DEQ approach. Our numerical results suggest the potential improvements in training/testing complexity due to ODER on three distinct imaging applications.

研究动机与目标

• 解决成像逆问题中批量处理在深度平衡模型（DEQ）中带来的高内存与计算成本问题。
• 在每次迭代中减少对完整测量集的依赖，实现通过DEQ端到端学习图像先验。
• 开发一种保持收敛性与性能的DEQ随机在线变体，同时提升训练与推理效率。
• 证明在PnP/RED框架中在线处理在大规模成像应用中的实际可行性。
• 提供在线方法相对于批量DEQ的收敛性与近似质量的理论分析。

提出的方法

• ODER在基于RED的DEQ中引入数据一致性层的随机近似，每次迭代仅处理测量集的一个子集，而非完整集合。
• 通过固定点迭代求解隐式网络方程，使用测量数据的小批量计算随机梯度。
• 通过固定点方程实现隐式反向传播，无需存储中间激活值。
• 将CNN先验与测量模型联合训练，以同时消除噪声和随机处理带来的伪影。
• 采用Adam优化算法，结合学习率调度与权重衰减，按应用（如CT、MRI）进行调整。
• 将测量数据划分为非重叠块，每次迭代从每个块中随机选择一个子集，用于计算数据一致性更新。

实验结果

研究问题

• RQ1 在基于DEQ的PnP/RED框架中，对测量数据进行随机处理是否能保持与批量处理相当的重建质量？
• RQ2 在线处理对隐式网络在图像重建中固定点迭代的收敛性与稳定性有何影响？
• RQ3 在大规模成像逆问题中，ODER能在多大程度上降低训练与推理复杂度，同时保持性能？
• RQ4 在每次迭代计算预算相当的情况下，ODER是否在SNR与SSIM指标上优于标准DEQ与基线方法？
• RQ5 测量子集大小与采样策略的选择如何影响ODER的最终重建质量？

主要发现

• ODER实现了与全批量DEQ相当的重建质量，在仅每次迭代使用50幅视图的稀疏视角CT中，SNR与SSIM性能与RED（DEQ 180/180）相当。
• 在稀疏视角CT中，ODER在每次迭代计算预算相同的情况下，相比使用50/180视图的批量DEQ，SNR提升约3 dB。
• 在10%采样率的并行MRI中，ODER在视觉重建质量与误差残差方面均优于RED（去噪）与RED（展开）。
• 在500次测量的强度衍射断层扫描中，ODER在SNR与SSIM上均达到与全批量RED（DEQ）相当的性能，且伪影抑制效果更优。
• 采用分块的随机采样测量子集策略，显著提升了ODER的重建质量与稳定性。
• 理论分析表明，ODER收敛至一个近似于批量DEQ固定点的解，在较弱假设下具有收敛性保证。