Equivariant Test-Time Training with Operator Sketching for Imaging Inverse Problems

TL;DR

This paper introduces a computationally efficient method for test-time training of deep imaging networks using equivariant regularization combined with randomized sketching, significantly reducing training time in high-dimensional inverse problems.

Contribution

The authors propose a sketched equivariant imaging regularization technique that accelerates test-time network adaptation for imaging inverse problems, with a parameter-efficient approach focusing on normalization layers.

Findings

Achieves significant computational acceleration in test-time training.

Effective in X-ray CT and MRI reconstruction tasks.

Reduces redundancy in high-dimensional unsupervised training.

Abstract

Equivariant Imaging (EI) regularization has become the de-facto technique for unsupervised training of deep imaging networks, without any need of ground-truth data. Observing that the EI-based unsupervised training paradigm currently has significant computational redundancy leading to inefficiency in high-dimensional applications, we propose a sketched EI regularization which leverages the randomized sketching techniques for acceleration. We apply our sketched EI regularization to develop an accelerated deep internal learning framework, which can be efficiently applied for test-time network adaptation. Additionally, for network adaptation tasks, we propose a parameter-efficient approach to accelerate both EI and Sketched-EI via optimizing only the normalization layers. Our numerical study on X-ray CT and multicoil magnetic resonance image reconstruction tasks demonstrate that our approach can achieve significant computational acceleration over the standard EI counterpart, especially in test-time training tasks.