关于Proximal Methods，近端梯度下降的理解

本文介绍了两种Proximal Methods的证明方法以及实现。内容主要来源于王然老师的《Proximal Methods》一文以及网络，加入了部分个人理解。由于水平有限，如有不妥之处，敬请指正。

为什么会有Proximal methods这个东东？

在机器学习的损失函数求解过程中，通过计算梯度然后迭代寻找最小值是一个常用的方法。而对于一些函数，是无法求导的，这时就无法用梯度下降等方法求解了。比如加了$L1$正则的损失函数。

$\underset{\beta }{\mathrm{argmin}}\frac{1}{N}\ast {\sum }_i(y_i-x_i\ast {\beta }^t)+\lambda \ast \Vert \beta {\Vert }_1$

proximal methods主要就是解决这个问题的。

proximal methods证明前的铺垫

主要介绍sub-differential和proximal operator这两个概念，后面证明时会用到。

sub-differential 子梯度

先介绍一个概念，sub-differential 子梯度，也叫：subderivative, subgradient, and subdifferential，是对于不可导的凸函数的导数的一种推广。
比如，对于绝对值函数$f(x)=|x|$，当$x=0$时，函数是不可导的。
如下图，对于$x_0$点不可导（类似绝对值函数），但是我们可以在点$(x_0,f(x_0))$上画一条线，这条线经过$x_0$点，并且在曲线的下方，像这样的曲线的斜率就是sub-differential中的一个。

子梯度的严格定义：
对于凸函数$f:I\to \mathbb{R}$在$x_0$的子梯度是一个实数$c$，$c$满足以下条件：
$f(x)-f(x_0)\ge c(x-x_0)$
对于所有在$I$内的$x$，在$x_0$的子梯度是一个非空的闭区间集合$[a,b]$，其中：
$a={\lim }_{x\to x_0^{-}}\frac{f(x)-f(x_0)}{x-x_0}$

$b={\lim }_{x\to x_0^{+}}\frac{f(x)-f(x_0)}{x-x_0}$

sub-differential记为$\partial f$，有：

∂ f = { y ∣ f ( x ) − f ( x 0 ) ≥ y T ( x − x 0 ) , f o r   a l l   x ∈ d o m   f } \partial f = \{ y | f(x)-f(x_0)\geq y^T(x-x_0), for \ all \ x \in dom \ f\} ∂f={y∣f(x)−f(x0​)≥yT(x−x0​),for all x∈dom f}

性质：
1、当一个凸函数在$x_0$处的子梯度只有一个值，即$a=b$时，函数在这个点可导。
2、如果一个凸函数在$x_0$处的子梯度集合为$[a,b]$，当$0\in [a,b]$时，函数在$x_0$处取得最小值。
3、如果$f,g$两个函数都是凸函数，则:
$\partial (f+g)=\partial f+\partial g$

另外，维基百科上说，国内的部分机构认为的凸函数的定义与国外的正好相反，不过本文并不想纠结于这个问题。

详见：
https://en.wikipedia.org/wiki/Subderivative

Proximal Operator

还要介绍一个概念。Proximal操作算子：

$pro{v}_f(v)=\underset{x}{\mathrm{argmin}}(f(x)+\frac{1}{2}\ast \Vert x-v{\Vert }_2^2)$

Proximal Operator有两个神奇的性质，一是不动点，二是proximal operator和sub-differential之间有一定的关系。

• 性质一：不动点

当$x^{\ast }$是$f(x)$的最小值时，等价于：
$x^{\ast }=pro{v}_f(x^{\ast })$
证明：
首先证明：$x^{\ast }$是$f(x)$的最小值时，$x^{\ast }=pro{v}_f(x^{\ast })$
$$\begin{array}{rl}f(x)+\frac{1}{2}\ast \Vert x-x^{\ast }{\Vert }_2^2& \ge f(x^{\ast })\\ & =f(x^{\ast })+\frac{1}{2}\Vert x^{\ast }-x^{\ast }{\Vert }_2^2\end{array}$$
即$f(x)+\frac{1}{2}\ast \Vert x-x^{\ast }{\Vert }_2^2$在$x=x^{\ast }$处取得最小值，即：$x^{\ast }=\underset{x}{\mathrm{argmin}}(f(x)+\frac{1}{2}\ast \Vert x-x^{\ast }{\Vert }_2^2)$，也就是$x^{\ast }=pro{v}_f(x^{\ast })$啦。

再证明：当$x^{\ast }=pro{v}_f(x^{\ast })$时，$x^{\ast }$是$f(x)$的最小值。

$x^{\ast }=pro{v}_f(x^{\ast })$，根据sub-differential的性质，有:
$$\begin{array}{rl}0& \in \partial (pro{v}_f(x^{\ast }))\\ 0& \in \partial (f(x)+\frac{1}{2}\Vert x-x^{\ast }{\Vert }_2^2)\\ 0& \in \partial f(x)+(x-x^{\ast })\\ 令x& =x^{\ast }，则有：\\ 0& \in \partial f(x^{\ast })\end{array}$$
即$x^{\ast }$是$f(x)$的最小值。

• 性质二：proximal operator实际上是sub-differential的一种解析形式，有：
  $$\begin{array}{r}pro{v}_{\lambda f}=(I+\lambda \partial f{)}^{-1}\end{array}$$
  说明：$pro{v}_{\lambda f}$和$(I+\lambda \partial f{)}^{-1}$都是操作算子，$pro{v}_{\lambda f}(v)=\underset{x}{\mathrm{argmin}}(f(x)+\frac{1}{2\lambda }\Vert x-v{\Vert }_2^2)$,$(I+\lambda \partial f{)}^{-1}$是$(I+\lambda \partial f)$的反函数。
  证明：
  如果：
  $$\begin{array}{rl}z& \in (I+\lambda \partial f{)}^{-1}(x)\\ (I+\lambda \partial f)(z)& \ni x\\ z+\lambda \partial f(z)& \ni x\\ 0& \in \lambda \partial f(z)+(z-x)\\ 0& \in \partial (\lambda f(z)+\frac{1}{2}\Vert z-x{\Vert }_2^2)\\ 0& \in \partial (f(z)+\frac{1}{2\lambda }\Vert z-x{\Vert }_2^2)\\ 即：& \\ z& =\underset{u}{\mathrm{argmin}}(f(u)+\frac{1}{2\lambda }\Vert u-x{\Vert }_2^2)\end{array}$$
  即：$(f(u)+\frac{1}{2\lambda }\Vert u-x{\Vert }_2^2)$在$z$处取得最小值，$z=pro{v}_{\lambda f}(x)$，注意这里的$x$其实是前面的$v$。
  这里有点儿神奇，当$z\in (I+\lambda \partial f{)}^{-1}(x)$时，$z=pro{v}_{\lambda f}(x)$
  两个看起来没什么关系的东西竟然也能联系在一起。。。

Proximal Methods的求解证明

文章的开头，我们就提出了一个问题：对于两个函数$f+g$，当$f$可导，但$g$不可导时，如何求解最小值呢？

我们先给出答案，再对其进行证明。
通过以下迭代，能够计算出$f+g$的最小值。

$x^{k+1}=pro{v}_{{\lambda }^{k}g}(x^k-{\lambda }^k\nabla f(x^k))$

• 证明方法一
  如果$x^{\ast }$是$f+g$的最小值，则有$0\in \nabla f(x^{\ast })+\partial g(x^{\ast })$，
  $$\begin{array}{rl}0& \in \lambda \nabla f(x^{\ast })+\lambda \partial g(x^{\ast })\\ 0& \in \lambda \nabla f(x^{\ast })-x^{\ast }+x^{\ast }+\lambda \partial g(x^{\ast })\\ 0& \in \lambda \nabla f(x^{\ast })-x^{\ast }+(I+\lambda \partial g)(x^{\ast })\\ (I+\lambda \partial g)(x^{\ast })& \ni x^{\ast }-\lambda \nabla f(x^{\ast })\\ x^{\ast }& =(I+\lambda \partial g{)}^{-1}(x^{\ast }-\lambda \nabla f(x^{\ast }))\\ x^{\ast }& =pro{v}_{\lambda g}(x^{\ast }-\lambda \nabla f(x^{\ast }))\end{array}$$
  这个证明过程也是很神奇的。。。

• 证明方法二

$$\begin{array}{rl}{x}^{k+1}& =pro{v}_{{\lambda }^{k}g}(x^k-{\lambda }^k\nabla f(x^k))\\ x^{k+1}& =\underset{x}{\mathrm{argmin}}(g(x)+\frac{1}{2{\lambda }^k}\Vert x-(x^k-{\lambda }^k\nabla f(x^k)){\Vert }_2^2)\\ x^{k+1}& =\underset{x}{\mathrm{argmin}}(g(x)+\frac{{\lambda }^k}{2}\Vert \nabla f(x^k){\Vert }_2^2+\nabla f(x^k{)}^T(x-x^k)+\frac{1}{2{\lambda }^k}\Vert x-x^k{\Vert }_2^2)\end{array}$$
由于上式是对于$x$求最小值，而$\frac{{\lambda }^k}{2}\Vert \nabla f(x^k){\Vert }_2^2$是一个与$x$无关的常量，则可将其替换为$f(x^k)$，则上式等价于：
$$\begin{array}{r}{x}^{k+1}=\underset{x}{\mathrm{argmin}}(g(x)+f(x^k)+\nabla f(x^k{)}^T(x-x^k)+\frac{1}{2{\lambda }^k}\Vert x-x^k{\Vert }_2^2)\end{array}$$
根据泰勒级数展开：
$$\begin{array}{r}f(x)=f(x^k)+\nabla f(x^k{)}^T(x-x^k)+\frac{1}{2{\lambda }^k}\Vert x-x^k{\Vert }_2^2\end{array}$$
则有：
$$\begin{array}{r}{x}^{k+1}=\underset{x}{\mathrm{argmin}}(g(x)+f(x))\end{array}$$
说句实在话，对于上面这种方式，个人表示还能凑合着理解，第一种证明的思路实在是难以想象。

根据前文不动点的性质，$x^{\ast }=pro{v}_f(x^{\ast })$，类似$x^{k+1}=pro{v}_{{\lambda }^{k}g}(x^k-{\lambda }^k\nabla f(x^k))$这种形式迭代方式也称为不动点迭代，

对于Proximal Method的理解

这是我在网上找到的比较能够理解的说法：
对于函数$f+g$，给定起点$x^k$，首先可微函数$f(x)$沿着起点的负梯度方向，作步长为${\lambda }^k$的梯度下降得到一个预更新值$x^k-{\lambda }^k\nabla f(x)$，然后使用近端映射寻找一个$x$ ，这个$x$ 能使得不可微函数$g(x)$足够小，且接近这个预更新值$x^k-{\lambda }^k\nabla f(x)$，就用这个$x$作为本次迭代的更新值$x^{k+1}$ 。

还有一个问题

$x^{k+1}=pro{v}_{{\lambda }^{k}g}(x^k-{\lambda }^k\nabla f(x^k))$，这个迭代算法为什么会成立？
除了不动点迭代外，还有一种解释这里只简单提一下，我也没深入研究（其实是水平不够，看文章太累了。。。），只是看了个皮毛。
当$\nabla f$是 Lipschitz continuous的，并且Lipshitz constant是$L$的情况下，当${\lambda }^k\in (0,1/L]$时，这是一个majorization-minimization method，具体可以查一下这个算法相关的资料。当${\lambda }^k>1/L$时，是另外一个问题。

关于不动点迭代的问题，继续解释可以了解：Forward-backward integration of gradient flow。

Proximal Methods的应用

设$f_{\beta }(X)$是负对数似然函数，其中$\beta$是需要求解的参数，$X$是样本数据，我们希望得到下面式子的最小值：
$$\begin{array}{r}{f}_{\beta }(X)+\lambda \Vert \beta {\Vert }_1，其中\lambda >0\end{array}$$
怎么求解$\beta$呢？
我们直接用$x^{k+1}=pro{v}_{{\lambda }^{k}g}(x^k-{\lambda }^k\nabla f(x^k))$这个迭代来搞定。
为了计算方便，我们令$\omega ={\beta }^k-{\lambda }^k\nabla f_{{\beta }^k}(x^k)$，其中${\lambda }^k$中在第$k$步迭代的步长，${\beta }^k$是在第$k$步迭代的$\beta$。
则有：
$$\begin{array}{rl}{\beta }^{k+1}& =pro{v}_{\lambda g}(\omega )\\ =& \underset{{\beta }_k}{\mathrm{argmin}}({\lambda }^k\lambda g({\beta }^k)+\frac{1}{2}\Vert {\beta }^k-\omega {\Vert }_2^2)\\ =& \underset{{\beta }_k}{\mathrm{argmin}}(\lambda \Vert {\beta }^k{\Vert }_1+\frac{1}{2{\lambda }^k}\Vert {\beta }^k-\omega {\Vert }_2^2)\end{array}$$
而：
$\Vert {\beta }^k{\Vert }_1={\sum }_i|{\beta }_i|,\Vert {\beta }^k-\omega {\Vert }_2^2={\sum }_i({\beta }_i-{\omega }_i{)}^2$
要计算$\lambda \Vert {\beta }^k{\Vert }_1+\frac{1}{2{\lambda }^k}\Vert {\beta }^k-\omega {\Vert }_2^2)$的最小值，我们只要找到每个$\lambda |{\beta }_i|+\frac{1}{2{\lambda }_k}({\beta }_i-{\omega }_i{)}^2$的最小值，然后求和就是总体的最小值了。

对于$\lambda |{\beta }_i|+\frac{1}{2{\lambda }_k}({\beta }_i-{\omega }_i{)}^2$的最小值，因为有绝对值，需要分类讨论：

• 当${\beta }_i\ge 0$时

$$\begin{array}{rl}& \lambda |{\beta }_i|+\frac{1}{2{\lambda }_k}({\beta }_i-{\omega }_i{)}^2\\ =& \frac{1}{2{\lambda }_k}({\beta }_i^2+2({\lambda }_k\lambda -{\omega }_i){\beta }_i+{\omega }^2)\end{array}$$
此时，当${\beta }_i={\omega }_i-{\lambda }_k\lambda$时，取得最小值，由于${\beta }_i\ge 0$，要求：${\omega }_i-{\lambda }_k\lambda \ge 0$。
但如果：${\omega }_i-{\lambda }_k\lambda <0$，${\beta }_i$无法取到${\omega }_i-{\lambda }_k\lambda$，当${\beta }_i=0$时，取到最小值。

• 当${\beta }_i<0$时
  $$\begin{array}{rl}& \lambda |{\beta }_i|+\frac{1}{2{\lambda }_k}({\beta }_i-{\omega }_i{)}^2\\ =& \frac{1}{2{\lambda }_k}({\beta }_i^2-2({\lambda }_k\lambda +{\omega }_i){\beta }_i+{\omega }^2)\end{array}$$
  此时，当${\beta }_i={\omega }_i+{\lambda }_k\lambda$时，取得最小值，由于${\beta }_i<0$，要求：${\omega }_i+{\lambda }_k\lambda <0$。
  但如果：${\omega }_i+{\lambda }_k\lambda >0$，${\beta }_i$无法取到${\omega }_i+{\lambda }_k\lambda$，当${\beta }_i=0$时，取到最小值。
• 综上：
  $$\begin{array}{r}{\beta }_i=\{\begin{array}{l}{\omega }_i-{\lambda }_k\lambda ,{\omega }_i>{\lambda }_k\lambda \\ \mathrm{０}，-\lambda {\lambda }_k<\omega <\lambda {\lambda }_k\\ {\omega }_i+{\lambda }_k\lambda ,{\omega }_i<-{\lambda }_k\lambda \end{array}\end{array}$$

正是由于$-\lambda {\lambda }_k<\omega <\lambda {\lambda }_k$时，${\beta }_i$会出现截断，取值为0时才能取得最小值，才使得损失函数+$L1$正则化时，得到稀疏解。

Proxiaml Methods的实现

这里我就不贴自己写的代码了，直接贴一下王然老师的代码：

1. 构造一个sigmoid函数：

def sigmoid(x):
    return 0.5 * (jnp.tanh(x / 2) + 1)

2. 构建逻辑回归模型：

def predict(beta, x):
    return sigmoid(x.dot(beta))

3. 构造数据

key = random.PRNGKey(0)
x_key, beta_key, beta_test_key = random.split(key,3)
x = random.normal(x_key, (10000, 10))
beta = random.normal(beta_key, (10,))*2.0    #beta是一个列向量
beta_test = random.normal(beta_test_key, (10,))
y = (sigmoid(x.dot(beta))>=0.5).astype(jnp.float32)

4. 建立逻辑回归的对数似然函数

def loss(beta):
    preds = predict(beta,x)
    #下面用了一个trick，进行了计算简化，如果不简化的话，应该是：y*jnp.log(preds) + (1 - y)jnp.log(1 - preds) ，而由于y只能为0或1，所以可以通过简化用以下的步骤实现：
    label_probs = preds * y + (1 - preds) * (1 - y) 
    return -jnp.sum(jnp.log(label_probs))/10000.00 

5. 对损失函数求梯度，有两种方式，两个的结果是一样的：
   一是数学推导，如下：

def custom_grad(beta):
    residual = y - predict(beta, x)
    return jnp.transpose(x).dot(-residual)/10000.00

二是通过jax.grad进行计算：

grad_func = jax.grad(loss)

6. 构造软阈值函数，就是Proximal Method最后那个${\beta }_i$。这里是通过jax.lax.cond来实现的，具体的介绍可以看一下官方文档，这个比较简单。
   前面写了那么那么多，在代码实现的时候，只有最后的结论能用的上。。。

def soft_threshold(x, thres):
    return jax.lax.cond(x > thres,
                        lambda _: x - thres,
                        lambda _: jax.lax.cond(
                            x < -thres,
                            lambda _: x + thres,
                            lambda _:0.0,
                            None
                        ),
                        None)

7. Proximal methods算法的迭代过程，具体我不多介绍了，应该算是一个比较标准的迭代过程。
   特别要说明一下，其实写这些代码的关键在于如何检测每步计算都是正确的，特别是在有向量，矩阵，求导，迭代的过程中，如何验证正确性是很麻烦的，检测的过程是保证结果正确的关键。

另外这里面计算每个${\beta }_i$时，用的是jax.vmap实现的并行计算。
对于jax.vmap，可以参考：https://jiayiwu.me/blog/2021/04/05/learning-about-jax-axes-in-vmap.html

def proximal_methods(beta_init, max_iter, eps, lr, penalty):
    converged = False
    beta_old = beta_init
    beta_new = beta_init
    soft_threshold_partial = lambda x: soft_threshold(x, lr*penalty)
    current_iter = 0
    while not converged and current_iter < max_iter:
        print("Current iteration is %d"% current_iter)
        beta_copy = beta_old 
        current_loss = loss(beta_copy) + penalty*jnp.linalg.norm(beta_copy, 1)
        current_grad = custom_grad(beta_copy)
        w = beta_copy - lr*current_grad
        beta_new = jax.vmap(soft_threshold_partial, 0)(w)
        new_loss = loss(beta_new) + penalty*jnp.linalg.norm(beta_new, 1)
        diff = jnp.abs(new_loss-current_loss)
        print("The difference is %.5f"%diff, "   current_loss%.5f"%current_loss, "   new_loss%.5f"%new_loss,)
        beta_old = beta_new
        if diff <= eps:   
            converged = True
            print("Algorithm converged")
            break
        else:
            current_iter +=1
            if current_iter >= max_iter:
                print("The algorithm have failed to converge.")
                break

    return beta_new, converged

参考资料

1. Proximal Mehtods,Ran Wang
2. Proximal Algorithm，Neal Parikh,Department of Computer Science Stanford University
3. 机器学习 | 近端梯度下降法 (proximal gradient descent)
4. 对近端梯度算法（Proximal Gradient Method）的理解
5. 【凸优化笔记4】-近端梯度下降（Proximal gradient descent）
6. Majorization-Minimization优化框架
7. 浅谈MM优化算法以及CCP算法