关于论文《ISTA-Net》的研究心得

关于论文《ISTA-Net》的研究

本篇内容为关于2018年CVPR论文《ISTA-Net: Interpretable Optimization-Inspired Deep Network for Image Compressive Sensing》的个人理解。该篇论文第一作者为北京大学副教授张建。该论文为deep unfolding方法在图像压缩感知方面的成功应用，搭建出了在“层”一级可解释的神经网络ISTA-Net。该网络的每一层对应于迭代软阈值算法(ISTA-Net)的一次迭代运算。该方法很好的结合了传统迭代算法和深度网络算法的优势，既大大提高了计算效率，又赋予了网络明确的可解读性。该文行文流畅、用词精准而不晦涩，实为一篇佳作。

要想很好地理解这篇论文，首先要从“迭代软阈值算法”的概念入手。

1. 迭代软阈值算法(ISTA)

我们常见的优化问题的目标函数可以表示为：

$$\hat{X}=argmin||X-B|{|}_2^2+\lambda ||X|{|}_1\text{\hspace{1em}(1)}$$
这里$\hat{X}$只最优解，$||\cdot |{|}_2$代表二范数，$||\cdot |{|}_1$代表一范数, $\lambda$为一个常数。(1)式中等号的右边两项，前者$argmin||X-B|{|}_2^2$为保真项, 用来衡量$X$和$A$的相似程度，后者$\lambda ||X|{|}_1$为正则项。保真项的意义很好理解，因为我们的目的就要找到和$B$最接近的$X$，所以要是这一项尽可能得小。而正则项的意义对于初学者往往不太容易理解，在这个问题里可以这样思考：

正则项相当于是给最小化保真项这一过程中添加了额外一个约束，即仅仅“满足保真项足够小”还不够，还要在"正则项不能太大"的约束下完成，即如果找到了一组解$\hat{X}$, 使得$||\hat{X}-B|{|}_2^2$很小，但$\lambda ||\hat{X}|{|}_1$很大，这样的$\hat{X}$是不可取的，必须二者的和一起约束到最小。
那为什么要这么干呢？答：有好几个作用。 主要有(1)缓解过拟合; (2) 满足$X$的稀疏性 （这一点在压缩感知里尤为重要）

再回到优化问题本身，究竟如何求解满足(1)式的$\hat{X}$呢？ 迭代软阈值算法(ISTA)就是一个迭代求解$\hat{X}$的算法。迭代很好理解，那么“软阈值”是什么意思呢？下面首先解释软阈值函数的概念

• 软阈值函数(soft-threshold)
  软阈值函数的表达式为
  $$soft(x,T)=\{\begin{array}{c}x+T,x\le -T\\ 0,|x|<T\\ x-T,x\ge T\end{array}\text{\hspace{1em}(2)}$$
  这个长相略微奇怪的函数跟求解(1)式有什么关系呢？我们再回到(1)式，把这些向量都写开, 设$X=[x_1,x_2,...,x_N{]}^T$, $B=[b_1,b_2,...,b_N{]}^T$,优化函数$F(X)$, 则
  $$\begin{gathered} F(X)=||X-B|{|}_2^2+\lambda ||X|{|}_1 \\ =\sum _{n=1}^N(x_n-b_n{)}^2+\lambda |x_n|\text{\hspace{1em}(3)} \end{gathered}$$
  即求$N$个形如 $f(x)=(x-b{)}^2+\lambda |x|$的函数的极小值。求这个函数的极值是不难的，直接对$x$求导，可得
  $$\frac{\partial f(x)}{\partial x}=2(x-b)+\lambda sgn(x)\text{\hspace{1em}(4)}$$
  令其为0，即
  $$2(x-b)+\lambda sgn(x)=0\text{\hspace{1em}(5)}$$
  解得，$$x=b-\frac{\lambda }{2}sgn(x)\text{\hspace{1em}(6)}$$
  $sgn(x)$的值取决于$x$和0的大小关系，因此要分情况讨论求极值点。
  a). 当$x>0时$,$sgn(x)=1,\to$ $x=b-\frac{\lambda }{2}>0,\to b>\frac{\lambda }{2}$
  b). 当$x<0时$,$sgn(x)=-1,\to$ $x=b+\frac{\lambda }{2}<0,\to b<-\frac{\lambda }{2}$
  因此，$f(x)$的极值点为
  $$\hat{x}=\{\begin{array}{c}b+\frac{\lambda }{2},b\le -\frac{\lambda }{2}\\ 0,|x|<\frac{\lambda }{2}\\ b+\frac{\lambda }{2},b\ge \frac{\lambda }{2}\end{array}\text{\hspace{1em}(7)}$$

现在再来仔细看一眼式(7), 是不是觉得这个形式有点眼熟？ 对啦！ 式(7)和软阈值函数的表达式的形式(式2)是相同的。
所以，(1)式的解$\hat{X}$可以写成:
$$\hat{X}=soft(B,\frac{\lambda }{2}).\text{\hspace{1em}(8)}$$
到这里，迭代软阈值算法就介绍完毕了。不过这里面貌似没有迭代的操作，因此还需要结合压缩感知信号恢复的具体应用背景来研究。下面将介绍压缩感知恢复的概念，并将ISTA贯穿其中。

2. 压缩感知回复(Compressed Sensing Reconstruction)

压缩感知理论的基础知识这里不再赘述，仅从恢复压缩后的信号开始介绍。设一个线性量测$\mathbf{y}$, 传统CS算法通过以下优化问题来恢复原始信号$\mathbf{x}$:
$$argmin\frac{1}{2}||\mathbf{\Phi }\mathbf{x}-\mathbf{y}|{|}_2^2+\lambda ||\mathbf{\Psi }x|{|}_1\text{\hspace{1em}(9)}$$
$\mathbf{\Psi }$为稀疏基矩阵，$\mathbf{\Phi }$为传感矩阵, 即$\mathbf{\Psi }$和观测矩阵(文中未显式表示)的矩阵乘积。注意，这里的符号表示和一般的压缩感知文献有区别, 一般的压缩感知文献$\mathbf{\Phi }$往往单纯指观测矩阵, 而论文中明确提到$\mathbf{x}$为original image, 而观测矩阵并不直接作用于原数据，而是原数据在稀疏域的映射，因此这里$\mathbf{\Phi }$为传感矩阵。同样，(9)式的右边第一项为保真项，第二项为正则项。

下面的论文直接给出了两个迭代公式:
$${\mathbf{r}}^{(k)}={\mathbf{x}}^{(k-1)}-\rho {\mathbf{\Phi }}^T(\mathbf{\Phi }{\mathbf{x}}^{(k-1)}-\mathbf{y})\text{\hspace{1em}(10)}$$

$${\mathbf{x}}^{(k)}=ar{g}_{\mathbf{x}}min\frac{1}{2}||\mathbf{x}-{\mathbf{r}}^{(k)}|{|}_2^2+\lambda ||\mathbf{\Psi }\mathbf{x}|{|}_1\text{\hspace{1em}(11)}$$
这下坏了，对于优化问题不是非常了解的读者读到这里开始二和尚摸不着头了。这都什么鬼，怎么这么突兀？式(11)的形式跟式(9)很像, 但是式(10)中的$\rho {\mathbf{\Phi }}^T(\mathbf{\Phi }{\mathbf{x}}^{(k-1)}-\mathbf{y})$这一项直接冒出来。论文由于篇幅所限没有详细解释这两个最重要的式子是怎么来的（默认读者是懂的），下面我来解释一下：

这里实际上用到了梯度下降的思想。式(10)中的${\mathbf{\Phi }}^T(\mathbf{\Phi }{\mathbf{x}}^{(k-1)}-\mathbf{y})$这个令人感到最突兀的一项其实就是保真项$\frac{1}{2}||\mathbf{x}-{\mathbf{r}}^{(k)}|{|}_2^2$在${\mathbf{x}}^{(k-1)}$处对$\mathbf{x}$的梯度（后有证明）。式(10)的含义是：在第$k-1$次迭代中, ${\mathbf{x}}^{(k-1)}$向保真项的负梯度方向移动，步长为$\rho$，得到的结果命名为${\mathbf{r}}^{(k)}$。这样操作后，${\mathbf{r}}^{(k)}$比${\mathbf{x}}^{(k-1)}$向着保真项极小值点的方向更近了一步。式(11)的含义是： 寻找一个新的${\mathbf{x}}^{(k)}$，使其逼近于上一次用式(9)算出来的${\mathbf{r}}^{(k)}$。算法的具体操作步骤为：首先初始化${\mathbf{x}}^{(0)}$, 带入式(10), 计算处${\mathbf{r}}^{(1)}$, 再将${\mathbf{r}}^{(1)}$带入式(11), 计算出${\mathbf{x}}^{(1)}$, 以此类推。

最后，这里附上保真项梯度的计算证明：
设$\mathbf{x}=[x_1,x_2,...,x_N]$, $\mathbf{y}=[y_1,y_2,...,y_M]$, $\mathbf{\Phi }={\varphi }_{ij},i=1,2,...,M,j=1,2,...,N,$ $M\ll N$.
$$\begin{gathered} f(\mathbf{x})=\frac{1}{2}||\mathbf{\Phi }\mathbf{x}-\mathbf{y}|{|}_2^2 \\ =\frac{1}{2}\sum _{i=1}^M[(\sum _{j=1}^N{\varphi }_{ij}{x}_j)-y_i{]}^2\text{\hspace{1em}(12)} \end{gathered}$$
设$x_t\in \mathbf{x}$, 则保真项式(12)对$x_t$的导数为：
$$\begin{gathered} {\nabla }_{x_t}f(\mathbf{x})=\frac{\partial f(\mathbf{x})}{\partial x_t}=\frac{1}{2}\frac{{\sum }_{i=1}^M[({\sum }_{j=1}^N{\varphi }_{ij}{x}_j)-y_i{]}^2}{\partial x_t} \\ =\frac{1}{2}\cdot 2\sum _{i=1}^M[(\sum _{j=1}^N{\varphi }_{ij}{x}_j)-y_i]\cdot \frac{\partial [{\sum }_{j=1}^N{\varphi }_{ij}{x}_j-y_i]}{\partial x_t} \\ =\sum _{i=1}^M[(\sum _{j=1}^N{\varphi }_{ij}{x}_j)-y_i]{\varphi }_{it} \\ =\sum _{i=1}^M{\varphi }_{it}[\mathbf{\Phi }\mathbf{x}-\mathbf{y}{]}_i \\ =\sum _{i=1}^M{\mathbf{\Phi }}_{t,:}^T[\mathbf{\Phi }\mathbf{x}-\mathbf{y}{]}_i\text{\hspace{1em}(13)} \end{gathered}$$
${\nabla }_{x_t}f(\mathbf{x})$是一个标量数字，为保真项关于$\mathbf{x}$中的任意一项的梯度。而保真项对于整个$\mathbf{x}$的梯度是一个$N$维向量，即
$${\nabla }_{\mathbf{x}}f(\mathbf{x})=[{\nabla }_{x_1}f(\mathbf{x}),{\nabla }_{x_2}f(\mathbf{x}),...{\nabla }_{x_N}f(\mathbf{x}){]}^T={\mathbf{\Phi }}^T(\mathbf{\Phi }\mathbf{x}-\mathbf{y})\text{\hspace{1em}(14)}$$

3. 软阈值函数在迭代算法中的作用

讲到这里，相信大部分读者对于这些概念已经明晰。我们在回过头来看一下上面两部分，有的读者可能会有这样的疑问： 第一部分解释了软阈值函数，第二部分提出了迭代算法，也没看出来这个迭代算法里用到了软阈值函数啊？这两部分有什么关系？这里就来解释一下这个问题。其实是在迭代算法中是用到了软阈值函数，但是没有写得那么直白。

式(11)其实还是一个优化问题，我们可以看到这个优化问题的形式其实跟式(1)是类似的，式(11)的解就是靠第一部分提到的软阈值函数来求解的。 注意啊，这只是类似，但是不一样！不能直接套用。所以作者在论文中的Equation.5-10, 都是通过各种手段把式(11)变成跟式(1)完全一样的形式，即：
$${\mathbf{x}}^{(k)}=ar{g}_{\mathbf{x}}min\frac{1}{2}||ϝ(\mathbf{x})-ϝ({\mathbf{r}}^{(k)})|{|}_2^2+\theta ||ϝ(\mathbf{x})|{|}_1\text{\hspace{1em}(15)}$$
的解为：
$$ϝ({\mathbf{x}}^{(k)})=soft(ϝ(\mathbf{x}),\theta )\text{\hspace{1em}(16)}$$
这里$ϝ(\cdot )$为两层卷积神经网络，作用等效于稀疏变换$\mathbf{\Psi }\mathbf{x}$的作用。

4. 对称结构$\widetilde{ϝ}({\mathbf{x}}^{(k)})$的作用

下面再简要解释一下原文中为什么要有用一个对称结构$\widetilde{ϝ}({\mathbf{x}}^{(k)})$?

这是因为式(16)得到的式$ϝ({\mathbf{x}}^{(k)})$, 而我们真正需要的是$({\mathbf{x}}^{(k)})$,因此需要一个$\widetilde{ϝ}({\mathbf{x}}^{(k)})$起到$ϝ({\mathbf{x}}^{(k)})$逆变换的作用，从$ϝ({\mathbf{x}}^{(k)})$中反解出$({\mathbf{x}}^{(k)})$。