联邦学习优化难题：近似ADMM算法原理与Matlab仿真实践

1. 先看问题：联邦优化的数学形式为什么难解

前阵子帮某实验室调一套联邦学习基线，数据从独立同分布换成非独立同分布之后，FedAvg（联邦平均）的损失曲线立刻开始折腾人——前几十轮还在下降，后面直接卡在一个高损耗区震荡，测试准确率也一直上不去。我一开始以为只是学习率没调好，把网格搜索翻来覆去跑了几轮，最后发现根本不是步长的问题，而是算法骨架在这里天生吃不开。后来换了基于近似交替方向乘子法的求解器，同样一批数据、同样的本地梯度下降流程，收敛速度和对最终模型质量都明显不一样了。

这篇文章就把那次重构的完整思路、数学推导、以及在Matlab里落地仿真的过程摊开来讲。流程很直接：先看联邦优化到底在解一个什么问题，再看标准ADMM为什么在联邦场景里跑不动，然后讨论"近似"在哪里、舍掉了什么、换来了什么，最后是整套可以复制运行的代码和调参避坑记录。适合已经接触过联邦学习、想摆脱"只会调FedAvg参"困境的读者；如果你正准备用Matlab做联邦学习算法仿真，这篇可以直接当参照模板用。

1.1 目标函数写成优化问题

联邦学习的本质，是让分布在多个节点上的本地数据"不出域"地共同训练一个全局模型。假设有 (N) 个客户端，第 (i) 个客户端持有本地数据集，拥有自己的经验损失函数 (f_i(w))，那么理想的全局优化目标是：

[ \min_{w} \quad \frac{1}{N} \sum_{i=1}^{N} f_i(w) ]

这个形式看起来简单，但里面藏了两个核心约束。第一，通信代价极高，客户端和服务器之间只能偶尔交换模型参数，不能频繁传递梯度或原始数据。第二，每个客户端本地数据往往高度异构，有的节点全是某种类别的样本，有的节点偏向某个特征区间，这会让"本地最优方向"和"全局最优方向"严重不一致。

很多论文会把联邦学习描述成"分布式优化的一种特例"，但和经典分布式优化不同的是，联邦场景不允许每次迭代都传一次梯度，也不允许服务器直接获取各个节点的数据分布信息。所有算法设计都被迫围绕一个非常朴素的流程：本地算几步、上传模型、服务器聚合、再下发。这个约束直接决定了哪些优化工具能用、哪些不能用。经典的牛顿法、拟牛顿法这类强二阶方法在这里基本没法用，因为二阶信息要么算不起，要么传不起。大家自然都把目光放到一阶方法、随机梯度类方法上，FedAvg就是其中最典型的代表。

1.2 FedAvg为什么在非独立同分布数据下收敛慢

FedAvg的迭代逻辑非常简单：每个客户端在本地数据上跑几轮梯度下降，然后把更新后模型发给服务器，服务器拿所有客户端模型做平均，作为下一轮的全局模型。当客户端数据是同分布时，本地梯度是全局梯度的无偏估计，平均操作等价于在全局目标上做了一次步长偏大的协调优化，收敛路径比较干净。

但数据变成非独立同分布之后，本地梯度就不再是全局梯度的好估计了。举个例子：客户端A拿到的几乎全是正样本，客户端B拿到的几乎全是负样本。A在本地把决策边界推向"尽量把东西判成正类"，B则相反。服务器把两个模型一平均，相当于把两个互相拉扯的优化方向硬生生折中，结果每一步都在来回摇摆。更麻烦的是，每个客户端本地更新时，根本不考虑别的客户端这轮会怎么动，也没有任何记忆机制来记录"我上一轮偏离了全局多少"。于是偏差被持续累加，形成所谓的客户端漂移，损失曲线平缓得让人怀疑人生。

从优化角度看，可以这样理解问题：FedAvg本质上是在交替执行"局部目标下降"和"全局变量聚合"两步，但这两个步骤没有用一个可量化的指标去约束彼此之间的偏差。聚合时缺少了对偶方向的信息，本地更新时又缺少聚合目标对当前状态的"拉力反馈"，所以收敛性质只能依赖数据同质这个前提。如果数据异构，算法骨架就撑不住了。解决思路是在每一步本地更新里，显式引入全局目标对它的牵引，再加一个累积偏差的记忆变量，让优化过程知道"我前几轮偏了多少，这轮得往回收一点"。这正是交替方向乘子法能提供的东西。

2. 从标准ADMM到"近似"ADMM：我们舍掉了什么、换来了什么

交替方向乘子法（ADMM）在分布式优化里成名已久，它最大的特点是把一个复杂的联合优化问题拆成本地子问题和全局协调问题，再用乘子记录每个子问题与全局目标之间的偏差。说句实话，第一次接触ADMM的人往往会被它那套增广拉格朗日形式吓退，但真的一步步拆开看，它讲的就是一个朴素的道理：让每个本地变量向共同目标靠拢，同时用一个"账本"记住谁之前欠了多少偏离，下次迭代让欠债的人多还一点。

2.1 标准一致性ADMM的三步迭代

把联邦学习套进ADMM框架时，最自然的做法是给每个客户端定义一个本地模型 (w_i)，再定义一个全局模型 (z)，然后约束它们相等：

[ \min_{w_1,\dots,w_N,z} \quad \sum_{i=1}^{N} f_i(w_i) ] [ \text{s.t.} \quad w_i = z,\quad i = 1,\dots,N ]

对约束做增广拉格朗日处理，写成缩放形式：

[ L_\rho(w,z,u) = \sum_{i=1}^{N} f_i(w_i) + \frac{\rho}{2}\sum_{i=1}^{N}|w_i - z + u_i|^2 ]

这里的 (u_i) 是缩放对偶变量，(\rho) 是惩罚参数。标准ADMM会交替执行三步：

1. 本地更新：固定 (z, u)，把每个 (w_i) 单独拿出来最小化；
2. 全局更新：固定 (w, u)，解出 (z) 的闭式解；
3. 对偶更新：更新 (u_i)，把本轮 (w_i) 与 (z) 的偏差记入账本。

全局更新这一步很便宜，因为 (z) 只出现在二次惩罚项里，直接对 (z) 求梯度并令其为零，就能得到：

[ z^{k+1} = \frac{1}{N}\sum_{i=1}^{N}\left(w_i^{k+1} + u_i^{k}\right) ]

对偶更新同样简单：

[ u_i^{k+1} = u_i^{k} + \left(w_i^{k+1} - z^{k+1}\right) ]

真正麻烦出在第一步的本地更新。标准ADMM的假设是，(w_i) 子问题能被精确求解，或者至少求解到很高的精度。在很多经典分布式优化场景里，这个假设确实成立，因为子问题常常是二次型或者带简单结构的凸问题，可以闭式求解。但联邦学习里的 (f_i(w_i)) 通常是复杂的非凸经验损失，本地精确求解要么代价高得离谱，要么根本做不到。直接把标准ADMM搬到联邦场景，本地更新一步就会卡死。

2.2 线性化近似：把子问题压缩成几步梯度下降

严格求解子问题行不通，那能不能放宽一点，只对子问题做几步梯度下降就算"差不多解完了"？这是近似ADMM的核心思想，也叫线性化ADMM或不精确ADMM。它把每个本地子问题里 (f_i(w_i)) 的部分做一阶泰勒展开，然后只取展开点附近的一小步下降，不再追求一步到位的精确解。

具体来说，本地更新想最小化的是：

[ G_i(w_i) = f_i(w_i) + \frac{\rho}{2}|w_i - z^k + u_i^k|^2 ]

直接对这个目标求梯度：

[ \nabla G_i(w_i) = \nabla f_i(w_i) + \rho \left(w_i - z^k + u_i^k\right) ]

用梯度下降逼近：

[ w_i^{k+1} = w_i^{k} - \eta \left[ \nabla f_i(w_i^{k}) + \rho \left(w_i^{k} - z^k + u_i^k\right) \right] ]

把这个式子多重复几次，也就是让每个客户端在本地做 (L) 步这样的更新，再统一把 (w_i) 发给服务器去做 (z) 和 (u) 的更新。当 (L) 足够大时，本地 (w_i) 会逐渐逼近精确子问题的解，近似ADMM就退化成标准ADMM；当 (L) 很小时，它仍然保留了ADMM的乘子反馈机制，比FedAvg多了一个纠正偏差的方向。

这其实就是标题里"近似"二字的落点所在。这里想强调一点：这种不精确求解并不是工程上的无奈妥协，而是有理论支撑的做法。ADMM的收敛性分析里，只要子问题求解的误差序列是可控的（通常要求误差之和有限），算法仍然能收敛到稳定点。所以"近似"换来的是每轮通信代价的大幅下降，付出的是每轮迭代质量的轻微折扣，这个交换在联邦场景里非常划算。

2.3 与FedAvg的算法骨架对比

把近似ADMM和FedAvg放在一起逐行对比，差异其实就两处。FedAvg的本地更新只做 (w_i \leftarrow w_i - \eta \nabla f_i(w_i))，聚合时简单平均 (z \leftarrow \frac{1}{N}\sum w_i)；近似ADMM则在本地梯度里多了一个 (\rho(w_i - z + u_i)) 项，服务器聚合时用的是 (w_i + u_i) 的平均，并且额外维护一个对偶变量。下表把骨架差异说得更直白一些：

从成本和收益看，近似ADMM给每个节点增加的计算量几乎可以忽略，只是多算了一项 (\rho(w_i - z + u_i)) 的梯度成分。但这多出来的两个惩罚项，恰好同时解决了FedAvg面临的两个问题：本地更新时缺少全局牵引力、聚合时缺少历史偏差反馈。用最朴素的话说，FedAvg平均完就散了，近似ADMM平均完还要回头检查一下"谁偏了、偏了多少、下次怎么纠正"。

在思考优化方向的时候有一个体会很关键：不要一上来就去复现复杂的联邦学习论文算法，先在一个能跑通的最小例子里把ADMM和FedAvg的骨架差异看清楚，后面再遇到别的联邦场景，心里会特别有数。

3. Matlab仿真：手写一个带对偶补偿的联邦求解器

理论部分聊完，直接上代码。这里用Matlab手写一个完整的近似ADMM求解器，场景设定为10个客户端协作训练一个逻辑回归二分类模型。为了让大家看清楚效果，我特意让数据非独立同分布——每个客户端使用不同的偏移方向生成标签，模拟真实异构环境。同时实现FedAvg作为对照，最后画出收敛曲线和残差曲线。完整脚本可以在MATLAB R2020a及以上版本运行。

3.1 数据生成：模拟非IID的二分类逻辑回归

数据仿真这一步很多人会忽略，但它直接决定算法对比的公平性。我采用的做法是：先生成一个全局真实权重 (\theta_{true})，然后让每个客户端在它基础上叠一个随机偏移，偏移幅度随客户端序号逐渐变大。这样既保证了各客户端有共性规律，又制造了明显的数据异构。别小看这个偏移量，它几乎就是整个实验里最能控制"难度"的旋钮，偏移越大，FedAvg漂移越严重，近似ADMM的优势越突出。

rng(42); N = 10; % 客户端数量 d = 20; % 特征维度 m = 500; % 每个客户端本地样本数 C = 0.01; % L2正则系数 X = cell(N,1); y = cell(N,1); theta_true = randn(d,1); for i = 1:N X{i} = randn(m, d); % 偏移幅度随客户端序号增大，模拟Non-IID theta_local = theta_true + 0.3 * (i-1) / N * randn(d,1); p = 1 ./ (1 + exp(-X{i} * theta_local)); y{i} = (p > rand(m, 1)) * 2 - 1; % 标签取 +1 / -1 end

这里把每个客户端的数据量都设为相同的500条，目的是先排除数据量不均衡这个干扰变量，让对比聚焦在数据分布差异对算法的影响上。实际项目里数据量往往不均衡，但在算法验证阶段先把变量控制住，能少走很多弯路。

3.2 近似ADMM核心循环实现

主循环严格按前面推导的三步走：本地近端梯度下降、服务器更新全局变量、对偶变量更新。每一步的物理含义我都写在注释里，方便对照公式看。需要特别留神的是 (z = mean(w + u, 2)) 这一行，很多从标准ADMM推导过来的人会在这里写错，如果漏了加 (u)，对偶补偿机制就彻底失效了，算法会退化成一个带正则的FedAvg，效果大打折扣。

%% 近似ADMM: 全局一致性约束 + 对偶变量补偿 rho = 0.1; % 惩罚参数 L = 5; % 本地梯度下降步数 eta = 0.02; % 本地学习率 K = 200; % 通信轮次 w = zeros(d, N); % 各客户端本地模型 z = zeros(d, 1); % 全局模型 u = zeros(d, N); % 缩放对偶变量 hist_loss = zeros(K, 1); hist_prim = zeros(K, 1); % 原残差 || w - z || hist_dual = zeros(K, 1); % 对偶残差 rho * || z_new - z_old || for k = 1:K % ---- 客户端本地更新: 近端梯度下降 ---- for i = 1:N Xi = X{i}; yi = y{i}; for l = 1:L g = logistic_grad(w(:, i), Xi, yi, C); g = g + rho * (w(:, i) - z + u(:, i)); w(:, i) = w(:, i) - eta * g; end end % ---- 服务器更新全局变量 z ---- z_old = z; z = mean(w + u, 2); % ---- 对偶变量更新: 把本轮偏差记入"账本" ---- u = u + (w - z); % ---- 记录指标 ---- total_loss = 0; for i = 1:N total_loss = total_loss + logistic_loss(w(:, i), X{i}, y{i}, C) / N; end hist_loss(k) = total_loss; hist_prim(k) = norm(w - z, 'fro') / sqrt(N * d); hist_dual(k) = rho * norm(z - z_old); end

辅助函数是两个常规的逻辑回归梯度与损失计算：

function g = logistic_grad(w, X, y, C) p = 1 ./ (1 + exp(-X * w)); g = X' * (p - (y + 1) / 2) / size(X, 1) + C * w; end function loss = logistic_loss(w, X, y, C) p = 1 ./ (1 + exp(-X * w)); loss = -mean( (y == 1) .* log(p) + (y == -1) .* log(1 - p) ) + C / 2 * sum(w.^2); end

这里把你可能见过的一个常见错误直接规避掉：逻辑回归的标签不能是0/1然后直接算 ((p-y))，因为逻辑回归的输出概率对应的是0/1标签，如果想要梯度表达式 (X^T(p-y)) 成立，得先把 (+1/-1) 标签转成0/1，也就是代码里的 ((y+1)/2)。另一种做法是用对数损失的原始形式直接求梯度，但表达式会绕一些。我习惯用这种转换写法，简洁且不容易出错。

3.3 FedAvg对照组实现与评估

为了公平对照，我给FedAvg也加了一个近端正则项 (\lambda|w_i - z|^2)。很多文献里的原生FedAvg是不带这个正则的，但这里我的目的是区分"乘子补偿"带来的增益，而不是"正则化"带来的增益，所以让两个算法拥有相同的正则结构，唯一区别就是有没有 (u_i) 这个历史偏差记忆。这样对比出来的结果更能说明ADMM骨架本身的价值，而不是被正则项的副产物干扰。

%% FedAvg 对照: 本地近端梯度下降 + 简单平均 lambda = 1; % 一致性正则权重 w_fed = zeros(d, N); z_fed = zeros(d, 1); hist_fed = zeros(K, 1); for k = 1:K for i = 1:N Xi = X{i}; yi = y{i}; for l = 1:L g = logistic_grad(w_fed(:, i), Xi, yi, C); g = g + lambda * (w_fed(:, i) - z_fed); w_fed(:, i) = w_fed(:, i) - eta * g; end end z_fed = mean(w_fed, 2); total_loss = 0; for i = 1:N total_loss = total_loss + logistic_loss(z_fed, X{i}, y{i}, C) / N; end hist_fed(k) = total_loss; end

最后画图的时候，把两个损失曲线叠在同一张图上，再加一张原残差和对偶残差的子图。原残差应当随轮次逐渐变小，对偶残差通常会出现先升后降的形态，这是ADMM类算法的正常表现，不必惊慌。

figure('Color', 'w'); subplot(1, 2, 1); plot(hist_loss, 'LineWidth', 1.5); hold on; plot(hist_fed, 'LineWidth', 1.5); legend({'approximate ADMM', 'FedAvg'}, 'FontSize', 10); xlabel('Communication round'); ylabel('Global loss'); title('Loss comparison'); subplot(1, 2, 2); plot(hist_prim, 'LineWidth', 1.5); hold on; plot(hist_dual, 'LineWidth', 1.5); legend({'primal residual', 'dual residual'}, 'FontSize', 10); xlabel('Communication round'); ylabel('Residual'); title('ADMM residuals');

在默认参数下跑出来的典型曲线是：前20轮两个算法都快速下降，看起来差距不大；大概从30轮开始，FedAvg的曲线进入高频低幅震荡状态，下降速度明显放缓；而近似ADMM由于对偶变量的持续补偿，仍然保持稳定的下降趋势。到了200轮，两者的全局损失通常能差出0.1到0.2，这在逻辑回归任务里已经是很可观的差距了。

4. 收敛曲线实测：参数如何影响迭代行为

参数调优这个环节是整篇文章里最容易让人浮躁的部分，你一定会在某个瞬间产生"直接把rho拉到100看看会不会更快"的冲动。我的建议是别急着暴力搜索，先把每个参数影响的是"哪个数学量级"搞清楚，再动手调。

4.1 默认参数下的收敛表现

把 (\rho=0.1)、(L=5)、(\eta=0.02) 这组默认参数跑完200轮，得到的典型曲线特征有两点值得注意。一是全局损失的下降不是单调的，中间会有小幅毛刺，这是因为服务器聚合只发生在一轮本地更新之后，而本地更新并不能保证每一步都让所有客户端的一致性偏差同步下降。二是原残差在头几十轮快速下降，对偶残差会先升后降，说明算法前期主要在做"校准方向"，后期才进入精细收敛阶段。

如果看到损失一路平稳下降、没有任何抖动，反而要警惕：那通常说明问题太简单，简单到任何算法都能轻松收敛，参数和算法优劣根本体现不出来。所以做仿真时别把数据生成得太平滑，适当提高异构偏移，才能看出各种算法在真实压力下的差异。

4.2 惩罚参数rho的取舍

(\rho) 是近似ADMM里最重要的参数，它控制的是"本地模型与全局模型偏差的惩罚强度"。(\rho) 太小，乘子的反馈作用微弱，算法行为逼近FedAvg，客户端漂移问题重新出现；(\rho) 太大，惩罚项占据本地梯度的主导地位，模型会被强拉向当前全局变量，本地数据里的有效信息反而没被充分利用，而且容易在聚合后产生振荡。实际操作里，(\rho) 的取值范围通常落在0.01到1之间，具体取多少和数据特征尺度关系很大。

一个比较实用的调参方法：先把 (\rho) 设成0.01，记录收敛轮数；再调成0.1、0.5、1，分别看损失曲线形态。如果 (\rho) 增大后收敛明显变快且曲线稳定，说明原问题确实需要更强的全局牵引；如果曲线出现高频振荡，说明惩罚过猛，需要回调。也可以采用动态调度策略，训练前期用较大的 (\rho) 把各客户端快速拉到一个相似的区域，后期减小 (\rho)，让本地模型在全局目标附近做更精细的调整。

4.3 本地迭代轮数的教训

本地迭代轮数 (L) 在近似ADMM里有个反直觉的现象：不是越大越好。(L) 太小，本地子问题求解不充分，每轮传给服务器的 (w_i) 与全局目标偏差大，收敛慢；但 (L) 过大之后，本地优化过度，把客户端数据里的局部噪声也学进了模型，反而加剧了异构漂移，服务器再如何用对偶补偿也补不回来。这个现象在FedAvg里也有，只是近似ADMM因为有乘子补偿，对 (L) 的耐受度会高一些。

在我这批合成数据上，(L=5) 时效果和 (L=10) 接近，而 (L=20) 时开始出现轻微的性能回落。经验和直觉一致：在联邦场景里，本地迭代的目标不是"把本地数据吃得透透的"，而是"在偏移不太大的前提下提供有信息量的更新方向"。方向准比步子远重要得多。

5. 实操踩坑记录：调这套代码时我犯过的错

很多跑过ADMM的人都有共识：理论推导看着顺顺当当，代码一写全是隐藏陷阱。下面这几个坑我基本都踩过一遍，有些是自己花了一下午盯着一堆发散曲线才反应过来的，有些是后来反复对比实验才总结出来的。

5.1 对偶变量符号约定

ADMM的缩放形式里，(u_i) 的更新符号有不止一种写法。我用的约定是：

[ u \leftarrow u + (w - z) ]

那么本地梯度里的惩罚项就是 (\rho(w - z + u))。如果哪篇文章里用的是 (u \leftarrow u + (z - w))，那它本地梯度里带的就是 (\rho(w - z - u))。两者差一个符号，从数学上等价，但代码里一旦混用，等价性就被破坏，系统会直接发散。判断方法很简单：把 (\rho) 设为0，算法必须退化到与FedAvg完全一致；再把 (\rho) 设为正值，如果损失没有比FedAvg更好，那多半就是符号或者 (z) 更新公式出了问题。

5.2 数据标准化与数值稳定

逻辑回归对特征尺度很敏感。如果某个特征在0.1量级，另一个特征在1000量级，那么梯度里不同维度的尺度会差好几个数量级，(\eta) 选得稍微大一点，特征尺度大的维度就梯度爆炸，(\eta) 选得小一点，特征尺度小的维度又几乎不更新。在联邦仿真里，每个客户端的数据最好在本地各自做标准化，并且把标准化参数固定住。千万不要图省事只在其中一个客户端上标准化，那样等于人为引入额外的数据异构，最后对比出的算法差距会失真。

5.3 收敛判据不能只看损失

只看损失曲线的最大风险是：损失看着在降，但一致性约束压根没满足。换句话说，各客户端停留在同一个损失谷底附近，但彼此之间并没有真正达成共识。此时你从服务器下发一个 (z)，它实际预测效果可能并不好。所以近似ADMM的收敛判断，必须同时观察原残差 (|w - z|) 和对偶残差 (\rho|z^{k+1} - z^k|)。工程上，我会把原残差和对偶残差的归一化值都降到某个阈值以下，才认为算法收敛。只把损失降到平台期就提前收手，很容易得到一个"看起来训练好了但测试效果一般"的模型。

这套代码现在还被实验室拿去当新算法的基准。每次跑新数据之前，我都会先把 (\rho) 从0.01扫到1，看一眼不同异构程度下的曲线形态，再决定要不要动 (L)。如果是你第一次在联邦学习里用近似ADMM，建议也照着这个流程走一遍：先复现FedAvg和近似ADMM在合成非IID数据上的差距，再试着把 (\rho) 和 (L) 分别调大调小，观察残差和损失的变化。真正理解了对偶变量那根"记忆链条"之后，你会发现在很多分布式优化问题里，它都能帮你把不太聪明的平均操作升级成有反馈、有纠正的协调机制。