共识优化

假设我们有一个凸优化问题，目标函数中有 \(N\) 个项

\[\begin{split}\begin{array}{ll} \mbox{minimize} & \sum_{i=1}^N f_i(x)\\ \end{array}\end{split}\]

例如，我们可能要将一个模型拟合到数据上，而 \(f_i\) 是第 \(i\) 个训练数据块的损失函数。

我们可以将这个问题转化为共识形式

\[\begin{split}\begin{array}{ll} \mbox{minimize} & \sum_{i=1}^N f_i(x_i)\\ \mbox{subject to} & x_i = z \end{array}\end{split}\]

我们将 \(x_i\) 解释为局部变量，因为它们特定于给定的 \(f_i\)。相比之下，变量 \(z\) 是全局的。约束条件 \(x_i = z\) 强制一致性，即共识。

我们可以使用交替方向乘子法 (ADMM) 解决共识形式的问题。ADMM 的每个迭代可以简化为以下更新步骤：

\[\begin{split}\begin{array}{lll} % xbar, u parameters in prox. % called proximal operator. x^{k+1}_i & := & \mathop{\rm argmin}_{x_i}\left(f_i(x_i) + (\rho/2)\left\|x_i - \overline{x}^k + u^k_i \right\|^2_2 \right) \\ % u running sum of errors. u^{k+1}_i & := & u^{k}_i + x^{k+1}_i - \overline{x}^{k+1} \end{array}\end{split}\]

其中 \(\overline{x}^k = (1/N)\sum_{i=1}^N x^k_i\)。

下面的代码使用 CVXPY 来解决本地子问题，实现共识 ADMM。

我们将 \(x_i\) 变量分拆到 \(N\) 个不同的工作进程中。工作进程并行更新 \(x_i\)。 主进程收集和平均 \(x_i\)，并向工作进程广播 \(\overline x\)。 工作进程在本地更新 \(u_i\)。

from cvxpy import *
import numpy as np
from multiprocessing import Process, Pipe

# Number of terms f_i.
N = ...
# A list of all the f_i.
f_list = ...

def run_worker(f, pipe):
    xbar = Parameter(n, value=np.zeros(n))
    u = Parameter(n, value=np.zeros(n))
    f += (rho/2)*sum_squares(x - xbar + u)
    prox = Problem(Minimize(f))
    # ADMM loop.
    while True:
        prox.solve()
        pipe.send(x.value)
        xbar.value = pipe.recv()
        u.value += x.value - xbar.value

# Setup the workers.
pipes = []
procs = []
for i in range(N):
    local, remote = Pipe()
    pipes += [local]
    procs += [Process(target=run_process, args=(f_list[i], remote))]
    procs[-1].start()

# ADMM loop.
for i in range(MAX_ITER):
    # Gather and average xi
    xbar = sum(pipe.recv() for pipe in pipes)/N
    # Scatter xbar
    for pipe in pipes:
        pipe.send(xbar)

[p.terminate() for p in procs]