Fast Equivariant Imaging: Acceleration for Unsupervised Learning via Augmented Lagrangian and Auxiliary PnP Denoisers

TL;DR

FEI tackles the slow convergence of unsupervised Equivariant Imaging by introducing a Lagrangian-based splitting into Latent-Reconstruction and Pseudo-Supervision steps. It presents two concrete implementations, HQS and linearized ADMM, and a plug-and-play variant (PnP-FEI) that leverages primal priors alongside dual EI priors. Across sparse-view CT and image inpainting tasks, FEI delivers about an order-of-magnitude speedup and improved generalization compared with vanilla EI, with EQPnP-FEI further boosting performance by incorporating equivariant priors. The work extends to other EI variants and establishes a framework that makes unsupervised deep imaging practical for large-scale problems.

Abstract

In this work, we propose Fast Equivariant Imaging (FEI), a novel unsupervised learning framework to rapidly and efficiently train deep imaging networks without ground-truth data. From the perspective of reformulating the Equivariant Imaging based optimization problem via the method of Lagrange multipliers and utilizing plug-and-play denoisers, this novel unsupervised scheme shows superior efficiency and performance compared to the vanilla Equivariant Imaging paradigm. In particular, our FEI schemes achieve an order-of-magnitude (10x) acceleration over standard EI on training U-Net for X-ray CT reconstruction and image inpainting, with improved generalization performance.

Figures (5)

• Figure 1: PSNR and MSE curves with respect to iteration during the training of Sparse-view CT reconstruction. We observe a 10-time acceleration can be achieved by PnP-FEI over EI (we fine-tuned EI with optimal algorithmic parameters for fastest convergence, but we did not fine-tune our own methods, so there is likely even more room for acceleration).
• Figure 2: Test on another unseen CT100 sample with Radon scan number is 50.
• Figure 3: Test on unseen CT100 sample with Radon scan number is 40 (a test on out-of-distribution generalization).
• Figure 4: PSNR and MSE curves with respect to iteration during the training of inpainting reconstruction. We observe a significant acceleration can be achieved by FEI over EI (we fine-tuned EI with optimal algorithmic parameters for fastest convergence, but we did not fine-tune our own methods, so there is likely even more room for acceleration).
• Figure 5: Test on another unseen Urban100 sample with mask rate is 0.6.