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

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. In addition, the proposed scheme enables efficient test-time adaptation of a pretrained model to individual samples to secure further performance improvements. Extensive experiments show that the proposed approach provides a noticeable efficiency and performance gain over existing unsupervised methods and model adaptation techniques.

💡 Research Summary

The paper introduces Fast Equivariant Imaging (FEI), an efficient unsupervised learning framework that dramatically speeds up the training of deep imaging networks without ground‑truth data. Traditional Equivariant Imaging (EI) enforces measurement consistency together with an equivariance loss that relies on the symmetry of the image distribution under geometric transformations. While EI achieves state‑of‑the‑art results, its training is slow because the equivariance term provides useful gradients only after the network already produces near‑perfect reconstructions, leading to weak supervision in early epochs and multiple costly forward passes per iteration.

FEI addresses this bottleneck by reformulating the EI objective through variable splitting and the method of Lagrange multipliers. The original problem is decomposed into two alternating sub‑steps: (1) a Latent‑Reconstruction step that refines an auxiliary image estimate using measurement fidelity, a regularizer, and optionally a pretrained plug‑and‑play denoiser (PnP‑FEI), and (2) a Pseudo‑Supervision step that treats the refined latent image as a pseudo‑ground‑truth to enforce the equivariance constraint on the network parameters. This decoupling eliminates the need to compute equivariance gradients with respect to the latent variable, reducing both runtime and memory consumption.

Two concrete algorithmic instantiations are presented. The first adapts Half‑Quadratic Splitting (HQS) with an adaptive penalty parameter and integrates Adam’s momentum to accelerate convergence. The second employs a Linearized ADMM scheme, using a linearized primal update and a closed‑form dual update, while allowing inexact sub‑problem solutions under bounded error conditions (as proven by inexact optimization theory). Both schemes reuse gradient history across iterations, achieving up to a ten‑fold speedup over standard EI.

The framework also supports efficient test‑time adaptation (TTA). After pretraining, a short FEI cycle (latent reconstruction followed by pseudo‑supervision) is run on each test sample, quickly fine‑tuning the model to the specific measurement distribution without accessing the original training data. Experiments on X‑ray computed tomography reconstruction and image inpainting demonstrate that FEI attains comparable or superior PSNR/SSIM while reducing training time from hours to minutes. Incorporating a pretrained denoiser in the latent‑reconstruction step further improves reconstruction quality, confirming the benefit of leveraging both image‑domain and measurement‑domain priors.

In summary, FEI combines variable splitting, inexact HQS/ADMM optimization, and plug‑and‑play priors to make unsupervised equivariant imaging practical for high‑dimensional inverse problems, delivering substantial gains in speed, accuracy, and adaptability. Future work may extend FEI to recent EI variants such as Robust EI, Multi‑Operator Imaging, and Sketched EI, broadening its impact across medical imaging and low‑level vision tasks.