Self-Supervised Deep Equilibrium Models for Inverse Problems with Theoretical Guarantees

2025-04-15 2 0 6.42MB 16 页 10玖币

侵权投诉

Self-Supervised Deep Equilibrium Models for Inverse Problems

with Theoretical Guarantees

Weijie Gan1, Chunwei Ying3, Parna Eshraghi2, Tongyao Wang2, Cihat Eldeniz3, Yuyang Hu4,

Jiaming Liu4, Yasheng Chen5, Hongyu An2,3,4,5, Ulugbek S. Kamilov1,4

1Department of Computer Science & Engineering, Washington University in St. Louis, St. Louis

2Department of Biomedical Engineering, Washington University in St. Louis, St. Louis

3Mallinckrodt Institute of Radiology, Washington University in St. Louis, St. Louis

4Department of Electrical & System Engineering, Washington University in St. Louis, St. Louis

5Department of Neurology, Washington University in St. Louis, St. Louis

{weijie.gan,chunwei.ying,p.eshrag,tongyaow,cihat.eldeniz,h.yuyang,

jiaming.liu,yasheng.chen,hongyuan,kamilov}@wustl.edu

Abstract

Deep equilibrium models (DEQ) have emerged as a powerful alternative to deep unfolding (DU) for image recon-

struction. DEQ models—implicit neural networks with effectively inﬁnite number of layers—were shown to achieve

state-of-the-art image reconstruction without the memory complexity associated with DU. While the performance of

DEQ has been widely investigated, the existing work has primarily focused on the settings where groundtruth data is

available for training. We present self-supervised deep equilibrium model (SelfDEQ) as the ﬁrst self-supervised recon-

struction framework for training model-based implicit networks from undersampled and noisy MRI measurements.

Our theoretical results show that SelfDEQ can compensate for unbalanced sampling across multiple acquisitions and

match the performance of fully supervised DEQ. Our numerical results on in-vivo MRI data show that SelfDEQ leads

to state-of-the-art performance using only undersampled and noisy training data.

1 Introduction

We consider an inverse problem where one seeks to recover an unknown image x∈Cnfrom its undersampled and

noisy measurements y∈Cm. Inverse problems are ubiquitous across medical imaging, bio-microscopy, and computa-

tional photography. In particular, compressed sensing magnetic resonance imaging (CS-MRI) is a well known inverse

problem that aims to recover diagnostic quality images from undersampled and noisy k-space measurements [26].

Deep learning (DL) has recently gained popularity in inverse problems due to its state-of-the-art performance [25,29].

Traditional DL methods train convolutional neural networks (CNNs) to map acquired measurements to the desired im-

ages [17, 43]. Recent work has shown that deep unfolding (DU) can perform better than generic CNNs by accounting

for the physics of the imaging system [1,35]. DU models are often obtained from optimization methods by interpreting

aﬁxed number of iterations as layers of a deep architecture and training it end-to-end. Despite the empirical success

of DU in some applications, the high memory complexity of training DU models limits its use in large-scale imaging

applications (e.g., 3D/4D MRI).

Recently, neural ODEs [7,19] and deep equilibrium models (DEQ) [4, 9] have emerged as frameworks for training

deep models with effectively inﬁnite number of layers without the associated memory cost. The potential of DEQ

to address imaging inverse problems was recently shown in [11]. Training a DEQ model for inverse problems is

analogous to training an inﬁnite-depth DU model with constant memory complexity. However, DEQ is traditionally

trained using supervised learning, which limits its applicability to problems with no groundtruth training data. While

there has been substantial interest in developing self-supervised learning methods that use undersampled and noisy

measurements for training [2, 36, 41], the potential of self-supervised learning has never been explored in the context

arXiv:2210.03837v1 [eess.IV] 7 Oct 2022

of DEQ. This work bridges this gap by proposing self-supervised deep equilibrium model (SelfDEQ) as a framework

for training implicit neural networks for MRI without groundtruth data. Our contributions are as follows:

• We introduce SelfDEQ as an image reconstruction framework for CS-MRI based on training a model-based

implicit neural network directly on undersampled and noisy measurements. SelfDEQ extends the line of work

based on Noise2Noise (N2N) [22] by introducing a model-based implicit architecture, a specialized loss function

that accounts for unbalanced sampling, and a memory-efﬁcient training method using Jacobian-Free Backprop-

agation (JFB) [9].

• We present new theoretical results showing that for certain measurement operators SelfDEQ computes updates

that match those obtained by fully-supervised DEQ. In the context of CS-MRI, our results imply that under a

set of explicitly speciﬁed assumptions, SelfDEQ can provably match the performance of DEQ trained using

the groundtruth MRI images. It is worth highlighting that the theoretical guarantees provided by our analysis

leverage the proposed correction for unbalanced sampling.

• We present new numerical results on experimentally-collected in-vivo brain MRI data. Our results show that

SelfDEQ can (a) outperform recent self-supervised DU methods; (b) match the performance of fully-supervised

DEQ, corroborating our theoretical analysis; and (c) enable highly-accelerated data-collection in parallel MRI.

2 Background

2.1 Imaging Inverse Problems

We consider inverse problems where the measurements yare speciﬁed by a linear system

y=MAx +e,(1)

where xis the unknown image, e∈Cmis additive white Gaussian noise (AWGN),A∈Cn×nis a measurement

matrix, and M∈ {0,1}m×nis a diagonal sampling matrix. A well-known application of (1) is CS-MRI [26], where

the measurements correspond to the noisy samples in the Fourier domain (referred to as k-space).

Inverse problems are generally ill-posed. Traditional methods recover xby solving a regularized optimization

x= arg min

f(x)with f(x) = g(x) + h(x),(2)

where gis the data-ﬁdelity term that quantiﬁes the discrepancy between the measurements and the solution, and his

a regularizer that imposes prior knowledge on the unknown image. Well-known examples in the context of imaging

inverse problems are the least-squares and total variation (TV)

g(x) = (1/2) ky−M Axk2

2and h(x) = τkDxk1,(3)

where Dis an image gradient and τ > 0is the regularization parameter.

2.2 Deep Learning

The focus in the area has recently moved to DL (see recent reviews in [25, 29]). A widely-used DL approach is to train

a CNN to learn a mapping from the measurements to the corresponding groundtruth images [17, 43]. There is also

a growing interest in deep model-based architectures (DMBAs) that can combine physical measurement models and

learned image priors speciﬁed using CNNs. Well known examples of DMBAs are plug-and-play priors (PnP) [18,39],

Regularized by Denoiser (RED) [32], and deep unfolding (DU) [1, 14, 35]. In particular, DU has gained notoriety due

to its ability to achieve the state-of-the-art performance, while providing robustness to changes in data acquisition.

DU architectures are typically obtained by unfolding iterations of an image reconstruction algorithm as layers, rep-

resenting the regularizer within image reconstruction as a CNN, and training the resulting network end-to-end. DU

architectures, however, are usually limited to a small number of unfolded iterations due to the high memory complexity

of training [35].

Figure 1: Illustration of SelfDEQ for CS-MRI. The forward pass of SelfDEQ computes a ﬁxed-point of an operator

consisting of data consistency layer and a CNN prior. The backward pass of SelfDEQ computes a descent direction

using the Jacobian-free update that can be used to optimize the training parameters. SelfDEQ is trained using the

proposed weighted loss that directly maps pairs of undersampled and noisy measurements of the same object to each

other without fully-sampled groundtruth.

2.3 Deep Equilibrium Models

DEQ has emerged as a framework for training recursive networks that have inﬁnitely many layers without storing

intermediate latent variables [4, 9, 11, 24, 30, 31, 42]. It is implemented by running two consecutive steps in each

training iteration, namely the forward pass and the backward pass. The forward pass computes a ﬁxed point ¯

xof an

operator Tθparameterized by weights θ

x=Tθ(¯

x,y),(4)

where yis the measurement vector. The ﬁxed point ¯

xis often computed by running a ﬁxed-point iteration with an

acceleration algorithm (such as Anderson acceleration [3]). It is worth noting that, when Tθdenotes a step of DMBA,

the DEQ forward pass is equivalent to DU with inﬁnitely many unfolded layers. Given a loss function, the backward

pass produces gradients with respect to θby implicitly differentiating through the ﬁxed points without the knowledge

of how they are estimated (see Sec. 4.2 in [11] for more details). DEQ does not require storing the intermediate

variables for computing the gradient, which dramatically reduces the memory complexity of training. There have been

several applications of DEQ in imaging, including applications to MRI [11, 30, 31], computed tomography (CT) [24]

and video snapshot imaging [42].

2.4 Self-Supervised Deep Image Reconstruction

There is a growing interests in developing DL methods that reduce the dependence on the groundtruth training data (see

recent reviews [2, 36,41]). Some well-known strategies include Noise2Noise (N2N) [22], Noise2Void (N2V) [21], deep

image prior (DIP) [38], Compressive Sensing using Generative Models (CSGM) [5,13], and equivariant imaging [6]. In

particular, N2N is one of the most widely-used self-supervised DL frameworks for image restoration that directly uses

noisy observations {b

xi,j =xi+ei,j }of groundtruth images {xi}for training. The N2N training can be formulated

arg min

θX

j6=j0kfθ(b

xi,j )−b

xi,j0k2

2,(5)

where fθdenotes the DL model with trainable parameters θ. There have been many extensions of N2N to different

imaging problems, such as MRI [10, 27, 40], OCT [16], and CT [15]. In particular, SSDU [40] is a recent state-of-the-

art method based on training a DU model without groundtruth by dividing a single k-space MRI acquisition into two

subsets that are used as training targets for each other. The work [27] has provided a theoretical justiﬁcation for SSDU

by extending Noisier2Noise [28] to variable-density subsampled MRI data.

2.5 Our Contributions

While DEQ has been shown to achieve the state-of-the-art imaging performance, the existing work has focused on set-

tings where groundtruth data is available for training. Our work addresses this gap by enabling DEQ training on noisy

and undersampled sensor measurements, which has not been investigated before. The proposed SelfDEQ framework

consists of several synergistic elements: (a) a model-based implicit network that integrates measurement operators and

CNN priors; (b) a self-supervised loss that accounts for sampling imbalances; (c) a Jacobian-free backward pass that

leads to efﬁcient training.

3 Method

3.1 Weighted Self-Supervised Loss

Consider the training set of measurement pairs {yi,y0

i}N

i=1 with each pair yi,y0

icorresponding to the same object xi

yi=MiAxi+eiand y0

i=M0

iAxi+e0

i.(6)

Here, N≥1denotes the number training pairs. One can obtain measurement pairs by physically conducting two

acquisitions or splitting each acquisition into two subsets.

Existing algorithms based on N2N directly map the measurement pairs to each other during training. However,

the measurements in the training dataset often have a signiﬁcant overlap. For example, each acquisition may share

the auto calibration signal (ACS) region [37], thus giving more weight to corresponding regions of the k-space (see

SelfDEQ (unweigted) in Fig. 6). We introduce a diagonal weighted matrix W=diag(w0, w1, ..., wn)∈Rn×nthat

accounts for the oversampled regions in the loss function. We set the diagonal entries of Was follows

wk=





√E[M0TM0]k,k qE[M0TM0]k,k 6= 0

0qE[M0TM0]k,k = 0

,(7)

where, in practice, the expectation over random sampling patterns can be replaced with an empirical average over the

training set. We can then deﬁne the following self-supervised training loss function

`self (θ) = EkM0A0¯

x(θ)−y0k2

W,(8)

where W=M0W(M0W)T∈Rm×mdenotes a subsampled variant of Wgiven M0, and ¯

x=Tθ(¯

x,y)denotes

the ﬁxed-point of Tθfor the input yand weights θ.

3.2 Forward and Backward Passes

The SelfDEQ forward pass is a ﬁxed-point iteration

xk=Tθ(xk−1,y),(9)

where

Tθ(x) = αfθ(s) + (1 −α)s

with s=x−γ∇g(x)(10)

The vector xkdenotes the image at the kth layer of the implicit network, γand αare two hyper-parameters, and fθ

is the CNN prior with trainable parameters θ. The implicit neural network is initialized using the pseudoinverse of

文档加载中……请稍候！
如果长时间未打开，您也可以点击刷新试试。

下载文档到电脑，查找使用更方便

10 玖币 0人已下载

立即下载

摘要：

Self-SupervisedDeepEquilibriumModelsforInverseProblemswithTheoreticalGuaranteesWeijieGan1,ChunweiYing3,ParnaEshraghi2,TongyaoWang2,CihatEldeniz3,YuyangHu4,JiamingLiu4,YashengChen5,HongyuAn2,3,4,5,UlugbekS.Kamilov1,41DepartmentofComputerScience&Engineering,WashingtonUniversityinSt.Louis,St.Louis2Depa...

展开>> 收起<<

Self-Supervised Deep Equilibrium Models for Inverse Problems with Theoretical Guarantees.pdf

共16页,预览4页

还剩页未读，继续阅读

声明：本站为文档C2C交易模式，即用户上传的文档直接被用户下载，本站只是中间服务平台，本站所有文档下载所得的收益归上传人(含作者)所有。玖贝云文库仅提供信息存储空间，仅对用户上传内容的表现方式做保护处理，对上载内容本身不做任何修改或编辑。若文档所含内容侵犯了您的版权或隐私，请立即通知玖贝云文库，我们立即给予删除！

Self-Supervised Deep Equilibrium Models for Inverse Problems with Theoretical Guarantees

相关推荐

开通VIP享超值会员特权

作者详情

相关内容

热门标签

举报选择: