约束优化 ADMM 算法：对偶上升算法、ADMM 算法

为了做论文当中的优化问题，特地学了一下约束优化相关的教材之外的方法，最近才有精力研读了一下 ADMM 算法相关的内容，发现很适合目前正在做的问题. 本文章的主要内容来源于 Stephen Boyd 的 ADMM 小册子 ¹ .

如果你需要求解非光滑的等式约束优化问题：

$$ \begin{aligned} & \min \quad f(\boldsymbol{x}) \\ & \mathrm{s.t.} \quad \boldsymbol{Ax} = \boldsymbol{b} \end{aligned} $$

或许本文有所帮助。

对偶上升算法 (Dual Ascent)

等式约束的对偶问题

现在考虑如下的等式约束优化问题：

$$ \begin{aligned} & \min \quad f(\boldsymbol{x}) \\ & \mathrm{s.t.} \quad \boldsymbol{Ax} = \boldsymbol{b} \end{aligned} $$

其中 $\boldsymbol{x}\in \mathbb{R}^{n}$ ，$\boldsymbol{A}\in \mathbb{R}^{m\times n}$ 以及 $f: \mathbb{R}^{n}\to \mathbb{R}$ 为凸函数. 此时可知其 Lagrange 函数为：

$$ L(\boldsymbol{x},\boldsymbol{y}) = f(\boldsymbol{x}) + \boldsymbol{y}^{\top} (\boldsymbol{Ax}-\boldsymbol{b}) $$

其中 $y\in \mathbb{R}^{m}$ 为 Lagrange 乘子，因此对偶函数为：

$$ g(\boldsymbol{y}) = \inf_{\boldsymbol{x}} L(\boldsymbol{x},\boldsymbol{y}) = - f^{*}(-\boldsymbol{A}^{\top} \boldsymbol{y}) - \boldsymbol{b}^{\top} \boldsymbol{y} $$

这里 $f^{*}$ 是 $f$ 的凸共轭函数，其定义为：

定义：共轭函数

$$ f^{*} (\boldsymbol{y}) = \sup_{\boldsymbol{x}\in \mathrm{dom}(f)} (\boldsymbol{y}^{\top} \boldsymbol{x} - f(\boldsymbol{x})) $$

因此对偶问题就是 $\max\, g(\boldsymbol{y})$ ，也就是最大化对偶函数，假设强对偶条件成立，此时对偶问题和原问题的解是一样的，$g(\boldsymbol{y})$ 的极大值点记为 $\boldsymbol{y}^{\star}$ ，那么可以从

$$ \boldsymbol{x}^{\star} = \arg\min_{\boldsymbol{x}} L(\boldsymbol{x},\boldsymbol{y}^{\star}) $$

得到 $\boldsymbol{x}^{\star}$ ，当 $f$ 是强凸函数，则有唯一的 $\boldsymbol{x}^{\star}$ . 当没有唯一值时，我们用上述的记号表示最小值点的集合.

对偶上升法的流程

我们使用梯度上升来求解对偶问题，此时可知 $\nabla g(\boldsymbol{y}) = \boldsymbol{Ax}^{+}-\boldsymbol{b}$ ，其中 $\boldsymbol{x} = \arg\min_{\boldsymbol{x}} \nabla L(\boldsymbol{x},\boldsymbol{y})$ ，因此更新的规则如下：

$$ \begin{aligned} \boldsymbol{x}^{(k+1)} & := \arg\min_{x} L(\boldsymbol{x},\boldsymbol{y}^{(k)}) \\ \boldsymbol{y}^{(k+1)} & := \boldsymbol{y}^{(k)} + \alpha^{(k)} (\boldsymbol{Ax}^{(k+1)} -\boldsymbol{b}) \end{aligned} $$

其中步长 $\alpha_{k}>0$ 通过线搜索获得，上标表示迭代的次数，这就是对偶上升法，这个算法在 $g$ 不可微时也能进行，此时 $\boldsymbol{Ax}^{+}-\boldsymbol{y}$ 则变成次梯度，上升的稳定性会稍微差一些（会因为函数性质使得上升可能不是单调的）.

如果 $\alpha_{k}$ 是精确的，那么对偶上升法最后会得到

$$ \boldsymbol{x}^{(k)}\to \boldsymbol{x}^{\star}, \quad \boldsymbol{y}^{(k)}\to \boldsymbol{y}^{\star} $$

也就是分别得到原问题最优解和对偶问题最优解.

此外，对偶上升在部分问题可能无法使用，哪怕 $g$ 可微，因为 Lagrange 函数可能出现无法求最小值的情况，且看下例.

例：求解：

$$ \begin{aligned} & \min \quad f(x_{1}, x_{2}) = x_{1} + x_{2}^{2} \\ & \mathrm{s.t.} \quad x_{1}+x_{2}=0 \end{aligned} $$

尽管这个很简单，可以人为手算，但是此时假如使用对偶上升法：

$$ L(x_{1}, x_{2}, y) = x_{1}+x_{2}^{2} + y(x_{1}+x_{2}) $$

那么进行对偶上升时，不管第一步 $y_{0}$ 如何，对于 $x_{1}$ ，它都是一个线性的函数，即

$$ L(x_{1}, x_{2}, y) = (1+y)x_{1} + x_{2}^{2} + yx_{2} $$

Lagrange 函数的最小值在 $1+y\neq 0$ 时是无法计算出来的，如果设定初始值为 $y^{(0)}=-1$ ，那么此时 $x_{1}$ 也将无法确定，只能随机选取. 当然，为了解决问题可以引入原问题约束，这样也就直接得到了 $x_{1} = \frac{1}{2} = -x_{2}$ . 这也是原问题的解. $\square$

大规模/矩阵方案：对偶分解

如果对于 $f(\boldsymbol{x})$ ，我们能将其分解为

$$ f(\boldsymbol{x}) = \sum\limits_{i=1}^{N} f_{i}(\boldsymbol{x}_{i}) $$

其中有 $\boldsymbol{x} = (\boldsymbol{x}_{1}^{\top}, \cdots, \boldsymbol{x}^{\top}_{N})$ ，这里 $\boldsymbol{x}_{i}\in \mathbb{R}^{n_{i}}$，仍然可以使用对偶上升法来求解. 注意这里 $\boldsymbol{x}$ 是维度为 $\sum\limits_{i=1}^{N} n_{i}$ 的向量. 不过如果建模得当，其实可以适用于很多场景，例如矩阵优化：

$$ \boldsymbol{X} = (\boldsymbol{x}_{1},\cdots, \boldsymbol{x}_{N}) \in \mathbb{R}^{d\times N}, \quad f(\boldsymbol{X}) = \sum\limits_{i=1}^{N} \|\boldsymbol{x}_{i}\|_{1} $$

也可以对大规模的向量进行并行求解. 将原有的 $\boldsymbol{A}$ 也进行列分块拆分：

$$ \boldsymbol{A} = \begin{pmatrix}\boldsymbol{A}_{1} & \boldsymbol{A}_{2} & \cdots & \boldsymbol{A}_{N}\end{pmatrix} $$

因此有

$$ \sum\limits_{i=1}^{N} \boldsymbol{A}_{i} \boldsymbol{x}_{i} = \boldsymbol{Ax} $$

Lagrange 函数也分解称为

$$ L(\boldsymbol{x}, \boldsymbol{y}) = \sum\limits_{i=1}^{N} L_{i}(\boldsymbol{x}_{i}, \boldsymbol{y}) = \sum\limits_{i=1}^{N} \left(f_{i} (\boldsymbol{x}_{i}) + \boldsymbol{y}^{\top} \boldsymbol{A}_{i} \boldsymbol{x}_{i} - \frac{1}{N}\boldsymbol{y}^{\top} \boldsymbol{b} \right) $$

此时的算法流程还是类似的：

$$ \begin{aligned} \boldsymbol{x}^{(k+1)}_{i} & := \arg\min_{x_{i}} L_{i}(\boldsymbol{x}_{i},\boldsymbol{y}^{(k)}) \\ \boldsymbol{y}^{(k+1)} & := \boldsymbol{y}^{(k)} + \alpha^{(k)} (\boldsymbol{Ax}^{(k+1)} -\boldsymbol{b}) \end{aligned} $$

可以看出，对偶上升法是具有天然的并行性的，这也说明对偶上升法很适合大规模矩阵、向量优化问题.

对偶上升法的改进：增广 Lagrange 函数

对偶上升法的缺陷刚刚说明了：Lagrange 函数很有可能无法求最小值，从而让方法不适用，这个缺陷比较致命，但是幸运的是，对偶上升法经过改良后可以克服这种困难.

现在将原问题稍微改写：

$$ \begin{aligned} & \min \quad f(\boldsymbol{x}) + \frac{\rho}{2} \|\boldsymbol{Ax} - \boldsymbol{b}\|_{2}^{2} \\ & \mathrm{s.t.} \quad \boldsymbol{Ax} = \boldsymbol{b} \end{aligned} $$

其中 $\rho>0$ 为实数，很显然，这是和原问题等价的，但是目标函数这么一改之后，Lagrange 函数就性质变好了：

$$ L_{\rho} (\boldsymbol{x}, \boldsymbol{y}) = f(\boldsymbol{x}) + \boldsymbol{y}^{\top} (\boldsymbol{Ax}-\boldsymbol{b}) + \frac{\rho}{2} \|\boldsymbol{Ax} - \boldsymbol{b}\|_{2}^{2} $$

因此对偶上升的步骤变为

$$ \begin{aligned} \boldsymbol{x}^{(k+1)} & := \arg\min_{\boldsymbol{x}} L_{\rho}(\boldsymbol{x},\boldsymbol{y}^{(k)}) \\ \boldsymbol{y}^{(k+1)} & := \boldsymbol{y}^{(k)} + \rho (\boldsymbol{Ax}^{(k+1)} -\boldsymbol{b}) \end{aligned} $$

接下来我们要说明两个问题：

• 为什么 Lagrange 函数此刻性质更好了；
• 为什么 $\rho$ 用来当步长了.

先回答第一个问题，之前的 $L$ 不方便最小化的情况来源于 $\boldsymbol{Ax}$ 是一个线性函数，很容易找到趋于无穷的方向，但是现在，由于 $\boldsymbol{x}$ 此时至少是二次的，从而在求解最小值的问题上，它不再容易出现这类无法求最值的问题.

第二个问题，$\rho$ 是惩罚参数，但是其实也是可以调节的，既然如何调节都不影响原目标函数的值，那么也可以将其视为步长进行调节.

因此，增广的对偶上升法可以适用于非强凸的 $f$ ，甚至可以适用于广义实函数，这点是比较大的改进.

ADMM 算法

算法流程

到这里，我们可以算是正式进入 ADMM 的正题，为什么我们要在开始之前讲对偶上升讲这么多呢？对偶上升法我们刚刚说过有两种情形：

• 增广的 Lagrange 函数可以使得适用范围更广；
• 原始的对偶下降可以轻松并行化.

但是问题在于：如果增广了 Lagrange 函数，那么并行化就不再容易（$L_{\rho}$ 并不好拆分），因此，ADMM 就是在这个基础上做了折中.

ADMM 算法主要适用于求解下列 2-block 的凸优化问题：

$$ \begin{aligned} & \min \quad f(\boldsymbol{x})+g(\boldsymbol{z}) \\ & \mathrm{s.t.} \quad \boldsymbol{Ax}+\boldsymbol{Bz}=\boldsymbol{c} \end{aligned} $$

这里面 $\boldsymbol{x}\in \mathbb{R}^{n}, \boldsymbol{z}\in \mathbb{R}^{m}$ ，同时

$$ \boldsymbol{A}\in \mathbb{R}^{p\times n}, \boldsymbol{B} \in \mathbb{R}^{p\times m} , \boldsymbol{c}\in \mathbb{R}^{p} . $$

而 $f: \mathbb{R}^{n}\to \mathbb{R}$ 和 $g: \mathbb{R}^{m}\to \mathbb{R}$ 都是凸函数.

实际当中一般 $3$-block 或者 $n$-block 都行……就是分析起来有点麻烦.

到这里可以写出增广 Lagrange 函数：

$$ L_{\rho}(\boldsymbol{x},\boldsymbol{z},\boldsymbol{y}) = f(\boldsymbol{x}) + g(\boldsymbol{z}) + \boldsymbol{y}^{\top} (\boldsymbol{Ax}+\boldsymbol{Bz}-\boldsymbol{c}) + \frac{\rho}{2} \|\boldsymbol{Ax}+\boldsymbol{Bz}-\boldsymbol{c}\|_{2}^{2} $$

正则项的作用是让 $f$ 的假设不必一定是严格凸，$\boldsymbol{y}$ 其实就是 Lagrange 乘子，而只要是一般的凸函数即可. 利用 Lagrange 乘子法有

$$ (\boldsymbol{x}^{(k+1)}, \boldsymbol{z}^{(k+1)}) = \arg\min_{\boldsymbol{x},\boldsymbol{z}} L_{\rho}(\boldsymbol{x},\boldsymbol{z},\boldsymbol{y}^{(k)}) , \quad \boldsymbol{y}^{(k+1)} = \boldsymbol{y}^{(k)} + \rho(\boldsymbol{Ax}^{(k+1)}+ \boldsymbol{Bz}^{(k+1)}-\boldsymbol{c}) $$

那么对于 ADMM ，实际上就是在原基础上让 $x,z$ 交替迭代，即

$$ \begin{aligned} \boldsymbol{x}^{(k+1)} & = \arg\min_{\boldsymbol{x}} L_{\rho}(\boldsymbol{x},\boldsymbol{z}^{(k)},\boldsymbol{y}^{(k)}) ,\\ \boldsymbol{z}^{(k+1)} & = \arg\min_{\boldsymbol{z}} L_{\rho}(\boldsymbol{x}^{(k+1)},\boldsymbol{z},\boldsymbol{y}^{(k)}) \\ \boldsymbol{y}^{(k+1)} & = \boldsymbol{y}^{(k)} + \rho(\boldsymbol{Ax}^{(k+1)}+ \boldsymbol{Bz}^{(k+1)}-\boldsymbol{c}) \end{aligned} $$

最终即可收敛到原问题的最优解.

尺度缩放版本

定义 $\boldsymbol{r} = \boldsymbol{Ax}+ \boldsymbol{Bz}-\boldsymbol{c}$ ，也就是残量，此时 ADMM 增广 Lagrange 函数的后两项变为

$$ \boldsymbol{y}^{\top} \boldsymbol{r} + \frac{\rho}{2} \|\boldsymbol{r}\|_{2}^{2} = \frac{\rho}{2} \|\boldsymbol{r}+ \boldsymbol{u}\|_{2}^{2} - \frac{\rho}{2}\|\boldsymbol{u}\|_{2}^{2} $$

这里 $\boldsymbol{u} = \frac{1}{\rho} \boldsymbol{y}$ ，于是可以得到 ADMM 迭代步骤：

$$ \begin{aligned} \boldsymbol{x}^{(k+1)} & = \arg\min_{\boldsymbol{x}}\left(f(\boldsymbol{x}) + \frac{\rho}{2} \|\boldsymbol{Ax}+\boldsymbol{Bz}^{(k)} - \boldsymbol{c} + \boldsymbol{u}^{(k)}\|_{2}^{2}\right) ,\\ \boldsymbol{z}^{(k+1)} & = \arg\min_{\boldsymbol{z}}\left(g(\boldsymbol{z}) + \frac{\rho}{2} \|\boldsymbol{Ax}^{(k+1)}+\boldsymbol{Bz} - \boldsymbol{c} + \boldsymbol{u}^{(k)}\|_{2}^{2}\right) \\ \boldsymbol{u}^{(k+1)} & = \boldsymbol{u}^{(k)} + \boldsymbol{Ax}^{(k+1)}+ \boldsymbol{Bz}^{(k+1)}-\boldsymbol{c} \end{aligned} $$

可以看到，每一步都是无约束优化问题或者直接的加减运算，至少减少了一个麻烦点，不过内层的优化迭代也会很大程度上影响速率.

ADMM 简单应用

LASSO

我们考虑 LASSO 的优化.

$$ \min_{\beta} \|\boldsymbol{y} - \boldsymbol{X \beta}\|_{2}^{2} + \lambda \|\boldsymbol{\beta}\|_{1} $$

这个优化问题就可以等价为

$$ \begin{aligned} & \min_{\beta} \|\boldsymbol{y} - \boldsymbol{X \beta}\|_{2}^{2} + \lambda \|\boldsymbol{\xi}\|_{1} \\ & \mathrm{s.t.} \quad \boldsymbol{\xi} = \boldsymbol{\beta} \end{aligned} $$

从这个约束优化出发可以写出 Lagrange 函数

$$ L(\boldsymbol{\beta}, \boldsymbol{\xi}, \boldsymbol{u}) = \|\boldsymbol{X \beta} - \boldsymbol{y}\|_{2}^{2} + \lambda \|\boldsymbol{\xi}\|_{1} + \frac{\rho}{2} \|\boldsymbol{\beta} - \boldsymbol{\xi} + \boldsymbol{u}\|_{2}^{2} $$

此时 ADMM 的优化步骤为

1. 更新 $\boldsymbol{\beta}$ ：

   $$ \boldsymbol{\beta}^{(k)} = \arg\min_{\boldsymbol{\beta}} \|\boldsymbol{y} - \boldsymbol{X \beta}\|_{2}^{2} + \frac{\rho}{2} \|\boldsymbol{\beta} - \boldsymbol{\xi}^{(k-1)} + \boldsymbol{u}^{(k-1)}\|_{2}^{2} $$

2. 更新 $\boldsymbol{\xi}$ ：

   $$ \boldsymbol{\xi}^{(k)}= \arg\min_{\boldsymbol{\xi}} \frac{\rho}{2}\|\boldsymbol{\beta}^{(k)} - \boldsymbol{\xi} + \boldsymbol{u}^{(k-1)}\|_{2}^{2} + \lambda \|\boldsymbol{\xi}\|_{1} = S_{\frac{\lambda}{\rho}} (\boldsymbol{\beta}^{(k)} + \boldsymbol{u}^{(k-1)}) $$

3. 更新 $\boldsymbol{u}$ ：

   $$ \boldsymbol{u}^{(k)} = \boldsymbol{u}^{(k-1)} + (\boldsymbol{\beta}^{(k)}- \boldsymbol{\xi}^{(k)}) $$

   这里面，第一步是光滑优化，第二步有显式解，因此迭代起来非常简单易行.