[论文解读] Practical Operator Sketching Framework for Accelerating Iterative Data-Driven Solutions in Inverse Problems

Junqi Tang, Xu, Guixian+2 more|arXiv (Cornell University)|Aug 31, 2022

Sparse and Compressive Sensing Techniques被引用 5

一句话总结

本文提出了一套实用的框架，通过随机优化与算子压缩技术，加速成像反问题中深度展开网络的计算。结合小批量梯度近似与通过压缩实现的降维，所提出的可学习随机对偶（LSPD）与压缩LSPD（SkLSPD）网络在显著降低计算与内存消耗的同时，保持了与全批量方法相当的重建精度，尤其在高维X射线CT与MRI任务中表现突出。

ABSTRACT

We propose a new operator-sketching paradigm for designing efficient iterative data-driven reconstruction (IDR) schemes, e.g. Plug-and-Play algorithms and deep unrolling networks. These IDR schemes are currently the state-of-the-art solutions for imaging inverse problems. However, for high-dimensional imaging tasks, especially X-ray CT and MRI imaging, these IDR schemes typically become inefficient both in terms of computation, due to the need of computing multiple times the high-dimensional forward and adjoint operators. In this work, we explore and propose a universal dimensionality reduction framework for accelerating IDR schemes in solving imaging inverse problems, based on leveraging the sketching techniques from stochastic optimization. Using this framework, we derive a number of accelerated IDR schemes, such as the plug-and-play multi-stage sketched gradient (PnP-MS2G) and sketching-based primal-dual (LSPD and Sk-LSPD) deep unrolling networks. Meanwhile, for fully accelerating PnP schemes when the denoisers are computationally expensive, we provide novel stochastic lazy denoising schemes (Lazy-PnP and Lazy-PnP-EQ), leveraging the ProxSkip scheme in optimization and equivariant image denoisers, which can massively accelerate the PnP algorithms with improved practicality. We provide theoretical analysis for recovery guarantees of instances of the proposed framework. Our numerical experiments on natural image processing and tomographic image reconstruction demonstrate the remarkable effectiveness of our sketched IDR schemes.

研究动机与目标

• 解决高维成像反问题（如X射线CT与MRI）中深度展开网络带来的高计算与内存开销问题。
• 通过整合随机一阶优化原理，开发一种可扩展、高效的全批量深度展开替代方案。
• 利用压缩技术实现降维，以近似昂贵的高维算子乘积。
• 为所提出的随机与压缩展开框架提供收敛性与估计误差的理论保证。
• 通过一种高效的实例自适应微调框架，支持分布外泛化能力。

提出的方法

• 提出基于小批量的可学习随机对偶（LSPD）网络，扩展了最先进的可学习对偶（LPD）框架，采用随机梯度下降方法。
• 引入压缩LSPD（SkLSPD），应用压缩技术以近似高维图像空间中的算子乘积，降低计算成本。
• 利用压缩矩阵近似大型前向与伴随算子的作用，实现对 $ A^T S_i^T S_i A $ 的高效计算，其中 $ S_i $ 为压缩算子。
• 将压缩技术与随机展开相结合，实现网络的联合加速与压缩，同时保持性能。
• 理论分析基于随机非凸复合优化，基于标准假设，推导了估计误差的上界。
• 采用带可学习去噪器与迭代更新的对偶展开结构，集成随机梯度估计与压缩近似。

实验结果

研究问题

• RQ1 在高维成像任务中，基于小批量梯度的随机展开能否实现与全批量深度展开相当的重建精度？
• RQ2 压缩技术在多大程度上可降低深度展开网络的计算与内存负担，同时不损害性能？
• RQ3 与全批量LPD相比，所提出的LSPD与SkLSPD网络在低剂量与稀疏视角X射线CT重建中的表现如何？
• RQ4 所提出的框架能否泛化至分布外数据？若能，其在新测量条件下的微调效率如何？
• RQ5 针对随机与压缩展开方案，可建立哪些关于收敛性与估计误差的理论保证？

主要发现

• LSPD与SkLSPD网络在仅需全批量LPD网络一小部分计算量的情况下，实现了与之相当的重建精度。
• 在低剂量与稀疏视角X射线CT任务中，所提方法在显著降低内存与计算负载的同时，保持了高质量的图像重建效果。
• 理论分析表明，在标准假设下，LSPD与SkLSPD的估计误差是有界的，且与全批量方法相当。
• SkLSPD通过压缩技术近似高维算子乘积，实现了进一步加速，且性能下降可忽略不计。
• 实例自适应框架可高效实现网络在分布外数据上的微调，展现出对噪声与模态分布偏移的强鲁棒性。
• 实证结果证实，随机展开与压缩技术的结合在不损失精度的前提下，实现了计算效率的‘免费午餐’。