论文笔记 | LDMNet: Low Dimensional Manifold Regularized Neural Networks(2)

在上一篇文章中，我们简单地介绍了LDMNet: Low Dimensional Manifold Regularized Neural Networks这篇文章的总体思路，而跳过了文章中采用的优化方法。本文将会更加注重数学上推导的内容。

ADMM

文章中，作者需要求解这样一个优化问题：

$$\begin{align} \small &\min_{\theta, \alpha}\quad J(\theta)+\frac{\lambda}{|\mathcal{M}^{(k)}|}\sum_{j=1}^d||\nabla_{\mathcal{M}^{(k)}}\alpha_j||^2_{L^2(\mathcal{M}^{(k)})} \\ &\ s.t.\quad\alpha(x_i, f_{\theta^{(k)}})=(\alpha_i, f_\theta(x_i)), \quad\forall i=1,...,N\tag{1} \end{align}$$

对于这样的问题，作者使用ADMM (alternation direction method of multipliers)来求解。那么什么是ADMM呢？

Dual Ascent

首先考虑一个等式约束的凸优化问题

$$\begin{align} &\min\ f(x) \\ &\ s.t.\ Ax=b\tag{2} \end{align}$$

其中$x\in\mathbb{R}^n$,$A\in\mathbb{R}^{m\times n}$,$f:\mathbb{R}^n\rightarrow\mathbb{R}$为凸函数 。这一问题对应的Lagrangian为

$$\begin{align} L(x,y)=f(x)+y^T(Ax-b) \end{align}$$

其中$y$为拉格朗日乘子。对偶函数为

$$\begin{align} g(y)=\inf_xL(x,y) \end{align}$$

对偶问题为

$$\begin{align} \max\ g(y) \end{align}$$

原问题与对偶问题的最优解应该是相同的，因此可以通过对偶问题的最优解$y^\ast$获得原问题的最优解$x^\ast$.

$$\begin{align} x^*=\mathrm{arg}\min_xL(x,y^*) \end{align}$$

首先通过固定$y$，求解此时的最优解$x^+$，然后再利用梯度上升更新$y$，$\nabla g(y)=Ax^+-b$。这样，就可以通过以下两步不断迭代求解：

$$\begin{align} &x^{k+1}:=\mathrm{arg}\min L(x,y^k) \\ &y^{k+1}:=y^k+\alpha^k(Ax^{k+1}-b)\tag{3} \end{align}$$

之所以称第二步为梯度上升，是因为对于合适的$\alpha^k$，每一步对偶函数都是在增长的。

Augmented Lagrangians and the Method of Multipliers

有时函数 $f$ 不能满足严格的凸性或者有限性，为了提高这一方法的强健性，(2)式的augmented Lagrangian可以写成

$$\begin{align} L_\rho(x,y)=f(x)+y^T(Ax-b)+(\rho/2)||Ax-b||_2^2 \end{align}$$

其中$\rho>0$被称为penalty parameter. 增加惩罚项的好处在于$g(y)$可以在大多数情况下变的可导。迭代的过程与(3)式非常类似

$$\begin{align} &x^{k+1}:=\mathrm{arg}\min_xL_\rho(x,y^k) \\ &y^{k+1}:=y^k+\rho(Ax^{k+1}-b) \end{align}$$

这种方法被称为method of multipliers. 这一方法比dual ascent更容易收敛，即使$f$会取到$\infty$，或是不满足严格的凸函数的要求。

Alternating Direction Method of Multipliers

类似于上文所讲，这一算法可以用于求解以下形式的问题

$$\begin{align} &\min\ f(x)+g(z) \\ &\ s.t.\ Ax+Bz=c\tag{4} \end{align}$$

其中$x\in\mathbb{R}^n$，$z\in\mathbb{R}^m$，$A\in\mathbb{R}^{p\times n}$，$B\in\mathbb{R}^{p\times m}$，$c\in\mathbb{R}^p$. 这里假定$f $和$g$都是凸函数。这一问题对应的 augmented Lagrangian为

$$\begin{align} L_\rho(x,z,y)=f(x)+g(z)+y^T(Ax+Bz-c)+(\rho/2)||Ax+Bz-c||_2^2 \end{align}$$

那么 ADMM 的迭代过程为

$$\begin{align} &x^{k+1}:=\mathrm{arg}\min_xL_\rho(x,z^k,y^k) \\ &z^{k+1}:=\mathrm{arg}\min_zL_\rho(x^{k+1},z,y^k) \\ &y^{k+1}:=y^k+\rho(Ax^{k+1}+Bz^{k+1}-c) \end{align}$$

Scaled Form of ADMM

通过将augmented Lagrangian中的线性项与平方项合并，并对dual variable进行缩放，ADMM还可以被写成一种更简单更常用的形式。

定义residual$r=Ax+Bz-c$，

$$\begin{align} y^Tr+(\rho/2)||r||^2_2&=(\rho/2)||r+(1/\rho)y||_2^2-(1/2\rho)||y||_2^2 \\ &=(\rho/2)||r+u||_2^2-(\rho/2)||u||^2_2 \end{align}$$

其中$u=(1/\rho)y$，被称为scaled dual variable.ADMM的迭代过程就变成了

$$\begin{align} x^{k+1}&:=\mathrm{arg}\min_x(f(x)+(\rho/2)||Ax+Bz^k-c+u^k||_2^2) \\ z^{k+1}&:=\mathrm{arg}\min_z(g(z)+(\rho/2)||Ax^{k+1}+Bz-c+u^k||^2_2) \tag{5}\\ u^{k+1}&:=u^k+Ax^{k+1}+Bz^{k+1}-c \end{align}$$

回到最初的问题

现在我们回到文章中提出的(1)式，并结合(5)式，就很容易得到论文中给出的结果

$$\begin{align} \alpha_\xi^{(k+1)}=&\mathrm{arg}\min_{\alpha_\xi}\sum_{j=d_1+1}^d||\nabla_{\mathcal{M}^{(k)}}\alpha_j||_{L^2(\mathcal{M}^{(k)})} \\ &+\frac{\mu|\mathcal{M}^{(k)}|}{2\lambda N}\sum\limits_{i=1}^N||\alpha_\xi(x_i, f_{\theta^{(k)}}(x_i))-(f_{\theta^{(k)}}(x_i)-Z_i^{(k)})||_2^2 \\ \theta^{(k+1)}=&\mathrm{arg}\min_\theta J(\theta)+\frac{\mu}{2N}\sum\limits_{i=1}^N||\alpha_\xi^{(k+1)}(x_i, f_{\theta^{(k)}}(x_i))\tag{6} \\ &- (f_{\theta}(x_i)-Z_i^{(k)}) ||_2^2 \\ Z_i^{(k+1)}=&Z_i^{(k)}+\alpha_\xi^{(k+1)}(x_i, f_{\theta^{(k)}}(x_i))-f_{\theta^{(k+1)}}(x_i) \end{align}$$

更新 $\theta$ 的反向传播

(6)式的求解用到了SGD，这个比较简单，就简要地说一下。令

$$\begin{align} E_i(\theta)=\frac{\mu}{2}||\alpha_\xi^{(k+1)}(x_i,f_{\theta^{(k)}}(x_i))-(f_{\theta}(x_i)-Z_i^{(k)})||^2_2 \end{align}$$

并注意到$J(\theta)$的定义

$$\begin{align} J(\theta)=\frac{1}{N}\sum_{i=1}^Nl(f_\theta(x_i), y_i) \end{align}$$

得到目标函数

$$\begin{align} \tilde{J}(\theta)=\frac{1}{N}\sum_{i=1}^Nl(f_\theta(x_i), y_i)+\frac{1}{N}\sum_{i=1}^NE_i(\theta) \end{align}$$

对于一个给定的网络，上式中第一项的反向传播是可以确定的，而第二项的梯度为

$$\begin{align} \frac{\partial E_i}{\partial f_\theta(x_i)}=\mu(f_{\theta}(x_i)-Z_i^{(k)}-\alpha_\xi^{(k+1)}(x_i,f_{\theta^{(k)}}(x_i))) \end{align}$$

更新 $\alpha$ 的 Point Integral Method (PIM)

好吧这里没有完全看懂，具体的数学推导参见以下两篇文章

Point Integral Method for Solving Poisson-type Equations on Manifolds from Point Clouds with Convergence Guarantees

Convergence of the Point Integral method for Poisson equation on point cloud

参考文献

[1] Boyd S , Parikh N , Chu E , et al. Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers[J]. Foundations & Trends in Machine Learning, 2010, 3(1):1-122.

[2] Li Z , Shi Z , Sun J . Point Integral Method for Solving Poisson-Type Equations on Manifolds from Point Clouds with Convergence Guarantees[J]. Communications in Computational Physics, 2017, 22(01):228-258.

[3] Shi Z , Sun J . Convergence of the Point Integral method for the Poisson equation with Dirichlet boundary on point cloud[J]. Mathematics, 2016.