【论文学习】SOLVING INVERSE PROBLEMS IN MEDICAL IMAGING WITH SCORE-BASED GENERATIVE MODELS

前言

好不容易写完了这么长的一篇，整体看来，这篇文章更像是对去噪过程的一个改进。通过在不同时间步引入对应的条件，来控制图像的生成。

文章中所有的译者注中举例都按照从X光复原CT图像这个问题出发，因为导师给的任务是这个。

相关概念

线性逆问题

对于一个未知的信号$x\in R^n$，我们已知$y\in R^m$，则建模为$y=Ax+ϵ$，其中$A$为线性算子，$ϵ$为未知噪音。线性逆问题的求解就是在已知$y$的情况下求解$x$的值。如果没有任何更进一步的假设，那么当$m<n$的时候，这个问题就是一个不明确的问题。于是我们假设$x$的采样服从一个先验的概率分布$p(x)$，在这个假设下，该问题的$x$与$y$可以由概率分布联系起来$p(y|x)=q_{ϵ}(y-Ax)$，其中的$q_{ϵ}$代表了噪音。那么在这个假设下，如果我们已知$p(x)$和$p(y|x)$我们就可以通过后验分布的方式求解$p(x|y)$。

在医学图像中，我们往往重建CT或者MRI。

贝叶斯公式：$$p(x|y)=\frac{p(y|x)p(x)}{\int p(y|x)p(x)dx}$$

译者注：因为我们假设有一个先验分布，所以对于任意的CT图像，我们都可以认为是从一个大分布中采样而来，而若这真的是一个线性逆问题，那么A应当是固定的线性算子。于是对于一张X光，我们只需要找到概率最大的CT图像即可。

基于分数的生成模型

对于当我们处理医学图像中的线性逆问题时，我们可以获得的是$y$和$p(y|x)$，而目标是获取$p(x|y)$，我们知道，先验分布$p(x)$显然是未知的，因此我们可以利用一个已有的数据集 { x 1 , x 2 , . . . , x n } ∼ p ( x ) \{x_1,x_2,...,x_n\} \sim p(x) {x1​,x2​,...,xn​}∼p(x)训练一个生成模型，用以估计$p(x)$，如果我们能估计$p(x)$那么根据贝叶斯公式，就可以获得后验概率了。

我们引入最近提出的基于分数的生成模型，其迭代采样过程使得以观测值为条件的可控生成特别容易。尤其是我们使用马尔可夫扩散过程生成噪音并将其采样为目标数据的时候。

扰动过程

对于一个给定的数据集，我们利用一个时域为[0,1]的随机过程扰动这些数据，这个扰动受到这样一个随机方程控制。$$d{x}_t=f(t)x_{t}dt+g(t)d{w}_t,t\in [0,1]\text{\hspace{1em}(1)}$$其中 f ( t ) : [ 0 , 1 ] → R , g : [ 0 , 1 ] → R , { w t ∈ R n } f(t):[0,1]\rightarrow R,g:[0,1]\rightarrow R, \{w_t\in R^n \} f(t):[0,1]→R,g:[0,1]→R,{wt​∈Rn}，这是一个标准维纳过程。而$x_t$代表了在这个随机过程中随机变量的轨迹。我们可以进一步表示$x_t$的边缘概率为$p(x_t)$，由此表示从$x_0$和$x_t$的分布转换。根据定义，我们明确有$p_0(x)\equiv p(x)$，进一步的，我们设置两个函数$f(t)$和$g(t)$来使得对任意初始化的分布$p_0(t)$在经过变换后，最终分布$p_1(t)$都会接近于我们预先定义的噪声分布$\pi (x)$。我们认为转换密度$p_{0t}(x_t|x_0)$总是一个条件线性高斯分布，其形式为$p_{0t}(x_t|x_0)=\aleph (x_t|\alpha (t)x_0,{\beta }^2(t)\mathbf{I})$，其中$\alpha :[0,1]\to R,\beta :[0,1]\to R$能够由$f(t)$和$g(t)$解析导出。

我们最终发现VE SDEs在我们的实验中表现最好。

逆过程

通过对（1）式求逆，我们从一个噪声取样开始$x_1\sim p_1(x)$，通过逐渐移除噪声来获得一个数据样本$x_0\sim p_0(x)\equiv p(x)$。下面我们给到下面公式来实现逆时间SDE$$d{x}_t=[f(t)x_t-g(t{)}^2{\nabla }_{x_t}log{p}_t(x_t)]dt+g(t)d\widehat{w_t}\text{\hspace{1em}(2)}$$其中 { w t ^ } \{\hat{w_t}\} {wt​^​}是逆世界方向的标准维纳过程。$dt$代表了无穷小的负时间步。这个式子中${\nabla }_{x_t}log{p}_t(x_t)$就是分数函数$p_t(x_t)$。通过逆向时间的定义，我们就得知可以由等式（2）的逆随机过程的轨迹。

采样

给定一个来自噪声分布$p_1$的初始样例以及中间时间步的分数${\nabla }_{x}log{p}_t(x)$，我们可以模拟（2）式来过去从数据分布$p_0(x)\equiv p(x)$中获取样例数据。在实践中，初始样例大概从$\pi (x)$抽取，且认为$\pi (x)\approx p_1(x)$。并且这些分数是可以由在数据集 { x ( 1 ) , x ( 2 ) . . . , x ( N ) } \{x^{(1)},x^{(2)}...,x^{(N)}\} {x(1),x(2)...,x(N)}训练得到的神经网络$s_{\theta }(x,t)$估计（这个神经网络名为 score model）。通过去噪得分匹配可以求解如下目标$${\theta }^{\ast }=\underset{\theta }{\mathrm{argmin}}\frac{1}{N}\sum _{i=1}^N{\mathbb{E}}_{t\sim \mathcal{U}[0,1]}{\mathbb{E}}_{x_t^{(i)}\sim p_{0t}(x_t^{(i)}|x^{(i)})}[||s_{\theta }(x_t^{(i)},t)-{\nabla }_{x_t^{(i)}}log{p}_{0t}(x_t^{(i)}|x^{(i)})|{|}_2^2]$$其中$\mathcal{U}[0,1]$代表了[0,1]上的均匀分布，去噪分数匹配理论保证了$s_{{\theta }^{\ast }}(x,t)\approx {\nabla }_{x}log{p}_t(x)$。在训练这个分数模型后，我们将其插入（2）式，求解之得$$d{x}_t=[f(t)x_t-g(t{)}^2{s}_{{\theta }^{\ast }}(x_t,t)]dt+g(t)d\widehat{w_t}\text{\hspace{1em}(3)}$$一种采样方式为利用欧拉-丸山数值解法来求解（3）式。

利用基于分数的生成模型求解逆问题

通过利用基于分数的生成模型，我们可以训练一个打分模型$s_{{\theta }^{\ast }}(x,t)$来从医学图像的先验分布$p(x)$生成无条件采样。然而为了求解逆问题，我们需要从后验概率$p(x|y)$才采样。这一操作可以通过给定条件来实现，给定原始随机过程 { x t } \{x_t\} {xt​}以条件$y$，产生条件随机过程 { x t ∣ y } \{x_t|y\} {xt​∣y}，我们给出在t时刻的边缘分布$p_t(x_t|y)$，我们的目标是从$p_0(x_0|y)$采样。我们可以通过求解下式得到$$d{x}_t=[f(t)x_t-g(t{)}^2{\nabla }_{x_t}log{p}_t(x_t|y)]dt+g(t)d\widehat{w_t}\text{\hspace{1em}(4)}$$条件分数函数${\nabla }_{x_t}log{p}_t(x_t|y)$是（4）式的重要部分。然而这计算并不容易。一个解决方案是通过训练一个新的显式依赖于$y$的得分函数$s_{{\theta }^{\ast }}(x_t,y,t)$来估计分数函数，这样由$s_{{\theta }^{\ast }}(x_t,y,t)\approx {\nabla }_{x_t}log{p}_t(x_t|y)$，然而这需要成对的训练数据$\{(x_i,y_i){\}}_{i=1}^N$并且具有与其它监督学习方法相同的缺点。本文中不考虑这种方式。

另一种无监督的替代方式是通过使用一个非条件的打分函数$s_{{\theta }^{\ast }}(x_t,t)\approx {\nabla }_{x_t}log{p}_t(x_t|y)$和测量分布$p(y|x)$来估计条件分数函数。已有的工作依赖于求解SVD，这对处理医学图像十分困难。而在另外的工作里，一种名叫ALD的采样方式被证明不如更新的采样方式。

因此我们提出一种新的采样方式来求解（吹牛部分省略）。总体来说，我们首先训练一个无条件打分模型$s_{{\theta }^{\ast }}(x_t,t)$，这一训练过程并不假设任何测量过程。在实验阶段，每给定一个观测$y$ ，我们通过在$y$上添加合适的噪音构造一个随机过程 { y t } \{y_t\} {yt​}，此后我们离散化（4）式这个逆时间SDE问题，这一过程利用我们之前提到的$s_{{\theta }^{\ast }}(x_t,t)$，与此同时将$y$的信息与附件优化步联合生成。

一种简便的线性测量过程形式

已有的医学图像处理的流程大多包含相同的计算步骤，例如基于Radon变换或基于类似的空间域上的傅里叶变换。为了严格描述这一运算过程，我提出了一种特别的形式的$A$，使得可以有效的在医疗图像应用中使用。不妨假设线性算子$A$是满秩的（$rank(\mathbf{A})=\min (n,m)=m$），下面的结果将会给出$A$的另一个形式。

观点1 如果$rank(\mathbf{A})=m$，那么存在一个可逆矩阵$T\in R^{n\times n}$和对角矩阵 Λ ∈ { 0 , 1 } n × n \Lambda\in \{0,1\}^{n\times n} Λ∈{0,1}n×n其迹$tr(\Lambda )=m$，所以有$\mathbf{A}=P(\Lambda )\mathbf{T}$，其中 P ( Λ ) ∈ { 0 , 1 } m × n P(\Lambda)\in \{0,1\}^{m\times n} P(Λ)∈{0,1}m×n是一个维数压缩算子，对任意n维的向量$a\in {\mathbb{R}}^n$，将其维度压缩到m维，这一性质由某些$i$的${\Lambda }_{ii}=0$保证。

如图，其实在大多数医学图像处理的过程中，算子$T$都是相同的，哪怕这些算子都与线性算子$A$有关。根据前文提到的方法，这个算子有可能是Radon变换，也有可能只是傅里叶变换。直观来看，对角矩阵$diag(\Lambda )$可以被视作一种对正弦图或k空间次采样（subsampling）的mask，这使得我们可以将原始的正弦图或k空间变成尺度更小的样本$y$。除此之外，由于方法已经确定，我们可以很容易求得$T^{-1}$，例如直接采用逆Radon变换或逆傅里叶变换即可。

将给定的观测结果融合进无条件采样过程

接下来，我们将证明，能够从条件随机过程 { x t ∣ y } \{x_t|y\} {xt​∣y}中利用一个无条件的打分模型$s_{{\theta }^{\ast }}(x,t)$产生近似样本。基本思想就是通过“劫持”无条件采样过程来使之融合观测条件$y$。

我们已经讨论过，直接求解 { x t ∣ y } \{x_t|y\} {xt​∣y}用于样本生成是十分困难的。为了绕过这个问题，我们先考虑一个简单的随机过程情况。让我们回到前文提到的$p_{0t}(x_t|x_0)=\aleph (x_t|\alpha (t)x_0,{\beta }^2(t)\mathbf{I})$，其中$\alpha (t)$和$\beta (t)$可以由$f(t)$和$g(t)$表示。任给定一个无条件的随机过程 { x t } \{x_t\} {xt​}，我们定义 { y t } t ∈ [ 0 , 1 ] \{y_t\}_{t\in [0,1]} {yt​}t∈[0,1]​，其中$y_t=A{x}_t+\alpha (t)ϵ$。则随机过程 { y t ∣ y } \{y_t|y\} {yt​∣y}是不同于 { x t ∣ y } \{x_t|y\} {xt​∣y}的，这个随机过程完全可处理。因为首先我们有$y_0=A{x}_0+\alpha (0)ϵ=A{x}_0+ϵ=y$，根据前文提到的$p_{0t}(x_t|x_0)=\aleph (x_t|\alpha (t)x_0,{\beta }^2(t)\mathbf{I})$，我们可以得到$x_t=\alpha (t)x_0+\beta (t{)}^{2}z,z\in {\mathbb{R}}^n\sim \aleph (0,\mathbf{I})$（p.s.这个性质是其实是高斯分布的独立可加性）。

根据定义，我们有$y_t=A{x}_t+\alpha (t)ϵ=A(\alpha (t)x_0+\beta (t{)}^{2}z)+\alpha (t)ϵ=\alpha (t)(A{x}_0+ϵ)+\beta (t{)}^{2}Az=\alpha (t)y_0+\beta (t{)}^{2}Az=\alpha (t)y+\beta (t{)}^{2}Az$，由此我们可以容易得到生成样本的方式$\widehat{y_t}\sim p_t(y_t|y)$。即我们首先采样一个随机噪音$z\sim \aleph (0,\mathbf{I})$，然后计算$\widehat{y_t}=\alpha (t)y+\beta (t{)}^{2}Az$

这个方法的核心在于可以用于修饰现存的用于处理无条件随机过程的迭代采样算法，利用我们定义的随机过程 { y t ∣ y } \{y_t|y\} {yt​∣y}会使得样本更加一致。总体而言，一个基于分数的生成模型会选择一系列时间步作为其迭代过程$${\hat{x}}_{t_{i-1}}=h({\hat{x}}_{t_i},z_i,s_{{\theta }^{\ast }}(x,t)),i=N,N-1,...,1\text{\hspace{1em}(5)}$$
其中${\hat{x}}_{t_N}\sim \pi (x),z_i\sim \aleph (0,I)$这里的${\theta }^{\ast }$是分数模型的参数。下面是该迭代函数的去噪过程，例如对前文提到的算法1，Euler-Maruyama采样的细节为$$h({\hat{x}}_{t_i},z_i,s_{{\theta }^{\ast }}(x,t))={\hat{x}}_{t_i}-f(t_i){\hat{x}}_{t_i}/N+g(t_i{)}^2{s}_{{\theta }^{\ast }}({\hat{x}}_i,t_i)/N+g(t_i)z_i/\sqrt{N}$$

为了体现根据 { y t ∣ y } \{y_t|y\} {yt​∣y}引出的限制，我们在迭代规则（5）式增加了一项$${\hat{x}}_{t_i}^{\prime }=k({\hat{x}}_{t_i},{\hat{y}}_{t_i},\lambda )\text{\hspace{1em}(6)}$$$${\hat{x}}_{t_{i-1}}=h({\hat{x}}_{t_i}^{\prime },z_i,s_{{\theta }^{\ast }}({\hat{x}}_{t_i},t_i))\text{\hspace{1em}(7)}$$
其中${\hat{x}}_{t_N}\sim \pi (x),{\hat{y}}_{t_i}\sim p_{t_i}(y_{t_i}|y)$并且$0\le \lambda \le 1$是超参数。整体的流程如下图3.迭代函数$k(\cdot ,{\hat{y}}_{t_i},\lambda ):R^n\to R^n$保证了数据连续性，通过求解最近邻优化步骤，这同时最小化${\hat{x}}_{t_i}^{\prime }$与${\hat{x}}_{t_i}$的间隔和${\hat{x}}_{t_i}^{\prime }$与超平面 { x ∈ R n ∣ A x = y ^ t i } \{x\in \mathbb{R}^n|Ax=\hat{y}_{t_i}\} {x∈Rn∣Ax=y^​ti​​}的间隔。这一优化过程对两者的权衡是由$\lambda$保证的。$${\hat{x}}_{t_i}^{\prime }=\underset{z\in {\mathbb{R}}^n}{\mathrm{argmin}}\{(1-\lambda )\Vert z-{\hat{x}}_{t_i}{\Vert }_T^2+\underset{u\in {\mathbb{R}}^n}{\min }\lambda \Vert z-u{\Vert }_T^2\},s.t.Au={\hat{y}}_{t_i}\text{\hspace{1em}(8)}$$

这里我们先回看之前的定义$A=P(\Lambda )T$，在等式中我们选择用范数$\Vert a{\Vert }_T^2=\Vert Ta{\Vert }_2^2$来简化理论推导。上面提到的观点1可以让我们得到一个近似形式：

定理1$${\hat{x}}_{t_i}^{\prime }=T^{-1}[\lambda \Lambda P^{-1}(\Lambda ){\hat{y}}_{t_i}+(1-\lambda )\Lambda T{\hat{x}}_{t_i}+(I-\Lambda )T{\hat{x}}_{t_i}]\text{\hspace{1em}(9)}$$

这里需要注意，根据之前的定义$\Lambda$作为一个欠采样的mask，则$P^{-1}(\Lambda )$则是将其恢复为原来的尺寸。

当$\lambda =0$时，${\hat{x}}_{t_i}^{\prime }=k({\hat{x}}_{t_i},{\hat{y}}_{t_i},0)={\hat{x}}_{t_i}$，即完全忽略限制条件$A{\hat{x}}_{t_i}^{\prime }={\hat{y}}_{t_i}$，这种设置下，我们的采样模型实际上就是等式（7）。当我们的测量是有噪声的，我们选择$0<\lambda <1$，来给限制条件一定的松弛度，$\lambda$的值在平衡${\hat{x}}_{t_i}^{\prime }\approx {\hat{x}}_{t_i}$与$A{\hat{x}}_{t_i}^{\prime }\approx {\hat{y}}_{t_i}$十分重要。实际中，我们在验证集上利用贝叶斯优化自动调整$\lambda$的数值。当我们的衡量过程没有噪声，则可以设定$\lambda =1$，此时将严格保证$A{\hat{x}}_{t_0}=y$。