最近在整理联邦学习模拟实验时我一开始用的还是最常规的FedAvg套路本地做几轮SGD服务端做加权平均。数据集是切到10个客户端上做逻辑回归整体跑下来精度勉强够用但训练曲线总是抖得让人不踏实——只要客户端数据分布一偏Non-IID全局模型的收敛路径就开始摇摆。后来我把求解器换成了“近似交替方向乘子法”近似ADMM效果明显不一样通信轮数几乎没增加收敛曲线平滑了很多对数据分布偏移也没有那么敏感。这个题目里被问得最多的一句话就是“联邦学习不是有FedAvg吗为什么还要搬出ADMM”这篇文章就从那个问题出发把近似ADMM的原理、MATLAB实现和调参经验一次讲透适合还没真正上手写联邦训练代码的入门者也适合正在FedAvg里来回踩坑、想换一条求解思路的人。1. 联邦学习里的优化难题先认清问题长什么样1.1 联邦场景的核心矛盾数据不出门模型却要全局一致联邦学习与普通分布式训练最大的差异在于原始数据不离开本地设备服务端能够拿到的只有客户端上计算出的模型参数或梯度。这个约束直接带来三件事通信带宽有限、客户端异构算力与数据分布均不同、隐私不可回传。矛盾自然出现了——你希望得到一个全局模型这个模型要适配所有客户端的本地数据但你又无法真正把数据集中到一处来训练。FedAvg的思路是“本地训练加权平均”它把问题压缩成一个非常朴素的启发式本地算几步梯度下降然后按数据量占比把模型参数平均起来。这个方法在独立同分布IID数据下工作得不错但一旦数据是Non-IID的本地客户端就会各自朝着自己的最优方向走平均出来的全局模型就会两边不讨好。专业一点的说法叫“客户端漂移”Client Drift每个本地模型偏离全局模型越来越远服务端聚合后模型质量明显下降。1.2 把联邦训练写成一个带约束的优化问题要更系统地解决这个问题需要先换一种数学视角来看联邦训练。经典的联邦学习目标函数是min F(w) Σ (n_k / n) * F_k(w)其中F_k(w)是客户端k在本地数据上的经验损失n_k是它的样本数n是总样本数。这个写法只是一个普通的经验风险最小化没有体现出“联邦”的特殊之处。要做联邦约束就得把它改写成带一致性约束的形式min Σ F_k(x_k) g(z) s.t. x_k z, k 1,...,K这里的x_k是客户端k的本地模型z是服务器维护的全局模型g(z)是全局正则项可以是0。约束x_k z的直观含义是每个客户端自己的模型必须与全局模型完全一致。这个改写很关键因为把耦合在全局模型上的优化问题解耦了。每个客户端在给定z的前提下单独优化自己的x_k服务端负责把x_k们“拉齐”成z同时利用对偶变量去感知本地与全局之间的差异。这正是ADMM的习惯形态约束被拆散但通过增广拉格朗日项重新拧成一股绳。1.3 为什么ADMM在这个问题上有天然优势ADMM适合处理“目标函数可分离 线性一致性约束”的结构而联邦学习恰恰长这个样子。与FedAvg相比ADMM解决的不再是“如何平均参数”这个工程问题而是“如何让本地模型与全局模型在数学上一致”的优化问题。做法是把x_k z约束放入增广拉格朗日函数中L_ρ Σ F_k(x_k) g(z) Σ λ_k^T (x_k - z) (ρ/2) Σ ||x_k - z||^2其中λ_k是对偶变量ρ是惩罚参数。看到没对偶变量λ_k出现在这里作用是自动检测“本地模型是否偏离全局模型”如果x_k和z不一致λ_k就会积累梯度下一轮迭代把它拉回来。这个机制就是FedAvg没有的也是ADMM在处理Non-IID数据时更稳的根本原因——它把“不一致”本身作为优化的对象而不是仅仅作为结果去平均掉。2. 从经典ADMM到“近似”ADMM不是偷懒是被逼出来的2.1 经典ADMM的三步更新到底在做什么经典ADMM每次迭代做三件事模型表达为x更新、z更新、对偶变量更新。写成缩放形式令u_k λ_k / ρ会更简洁x_k ← argmin_x F_k(x) (ρ/2) ||x - z u_k||^2 z ← argmin_z g(z) (ρ/2) Σ ||x_k - z u_k||^2 u_k ← u_k x_k - z第一步每个客户端在固定全局模型z和对偶变量u_k的前提下求解自己的本地子问题。这个子问题的目标函数里除了原有的损失项F_k(x)还多了一个“距离惩罚项”要求x不要偏离z太远——但u_k会偏移这个中心点所以它知道该往哪个方向补偏差。第二步服务端收集所有x_k后把z更新为所有“补偿后的本地模型”的一个近端平均。如果没有g(z)那z就是mean(x_k u_k)非常轻量。第三步对偶变量u_k沿残差方向增加x_k - z。如果x_k和z一致u_k不再变化迭代自然停止如果不一致u_k会逐步编码这个偏差在下一轮反馈给本地更新。2.2 经典ADMM在联邦场景跑不动的两个硬伤理想很丰满但我在实际实验里踩到的第一个坑就是经典ADMM在这里的更新成本太高了。第一步要求“精确求解”本地子问题也就是客户端需要对每个本地数据做完整迭代直到收敛。这意味着单轮通信内每个客户端要跑几十甚至上百次梯度下降这对真实联邦场景根本不可行。通信成本是省下来了计算成本却爆炸了。第二个问题更隐性本地子问题里如果F_k(x)是非凸的比如神经网络精确求解本身就很难收敛你还要保证它的二阶信息可用否则增广拉格朗日函数的二次项在导函数意义上不是单调的。经典ADMM的收敛性证明大多依赖凸假设在联邦深度学习的非凸场景下保证少了许多。这两个硬伤正是“近似ADMM”登上舞台的原因——不是要偷懒省那几步计算而是经典的“精确子问题求解”在联邦学习里本身就跑不动。2.3 近似ADMM的核心思想本地不精确求解照样能收敛近似ADMM的做法非常直接不要求客户端精确解出子问题只要求它做有限几步的近端梯度下降或近端SGD。数学上这叫“不精确ADMM”它在满足一定条件的前提下仍然有收敛保证——条件是每一步引入的近似误差不能累积通常要求误差序列满足可和条件。具体到我实现的版本本地更新写成这样x_k ← x_k - α * (∇F_k(x_k) ρ (x_k - z u_k))再做一个近端投影如果有正则项。这个公式看着简单但它融合了两层含义第一项是本地损失梯度让模型继续沿本地数据的方向下降第二项是增广拉格朗日的惩罚梯度让模型不要偏离全局补偿目标。u_k在这里就像一杆“校准秤”——本地数据分布偏了u_k就把中心点拉偏使本地更新不敢肆无忌惮地漂移走。这个方案与FedAvg的区别在于FedAvg本地做多步普通SGD然后平均近似ADMM本地做多步近端SGD并且多了一个对偶变量参与校准。它不是简单地对模型取平均而是让模型为“一致性约束”付出优化代价。这带来的直接收益是在Non-IID越严重的情况下ADMM的收敛越显得稳。3. 整套算法的MATLAB落地实现3.1 实验准备数据生成与Non-IID划分我实验里用的是合成二分类数据特征维度d20总样本数N500客户端数量K10。数据生成用两簇高斯分布正负样本中心分别偏移0.8再随机翻转10%的标签制造噪声这样既简单又能确保损失函数可收敛。真正需要小心的是Non-IID划分的方式。一种常见做法是按标签排序后分段分配给不同客户端这会以极悬殊的方式制造“标签偏移”% 按标签排序然后交错分给K个客户端制造Non-IID(标签分布偏移) [sortedY, idx] sort(y); clientIdx cell(K,1); for k 1:K clientIdx{k} idx(k:K:N); end这种划分方式会让某些客户端几乎只看到正样本另一些几乎只看到负样本。如果想调节Non-IID的严重程度也可以改用狄利克雷分布按标签比例采样浓度参数越小数据分布越偏。但在做第一版实验时用上面的简单排序划分就够了。3.2 核心更新公式的代码化落地本地梯度的计算采用逻辑回归损失function grad computeGradient(Xk, yk, x) p 1 ./ (1 exp(-Xk * x)); grad Xk * (p - (yk 1) / 2) / size(Xk, 1); end这里一个容易犯的错是把标签写成0/1而不是-1/1导致梯度的符号反向。我最初就在这个点上栽过调试了整整一个下午——模型不收敛所有参数检查都正常最后发现是标签映射反了。逻辑回归的常用标签约定是{-1, 1}p的公式对应的是二分类模型P(y1|x)的sigmoid输出。3.3 近似ADMM主程序骨架function [z, history] fedadmm(X, y, K, R, rho, alpha, locSteps) % fedadmm 联邦学习近似ADMM主程序 % X: 特征矩阵 y: 标签 K: 客户端数 R: 通信轮数 % rho: ADMM惩罚参数 alpha: 本地学习率 locSteps: 本地近端SGD步数 [N, d] size(X); clientIdx cell(K,1); for k 1:K clientIdx{k} k:K:N; % 按样本轮转分组适配上面Non-IID划分 end z zeros(d, 1); % 全局模型 u zeros(d, K); % 缩放对偶变量 x zeros(d, K); % 本地模型 history.primalRes zeros(R,1); history.dualRes zeros(R,1); history.loss zeros(R,1); for r 1:R % 第一步所有客户端本地更新 for k 1:K Xk X(clientIdx{k}, :); yk y(clientIdx{k}); xk x(:,k); for t 1:locSteps grad computeGradient(Xk, yk, xk); xk xk - alpha * (grad rho * (xk - z u(:,k))); end x(:,k) xk; end % 第二步服务端更新全局模型 z_new mean(x u, 2); % 第三步对偶变量更新 u u (x - z_new); % 残差记录用于收敛判断 history.primalRes(r) norm(x - z_new, fro) / sqrt(K); history.dualRes(r) norm(z_new - z, 2) / (1 norm(z, 2)); z z_new; history.loss(r) computeLoss(X, y, z); if history.primalRes(r) 1e-4 history.dualRes(r) 1e-4 break; end end end代码里每行的讲究都不少。第一步中的本地更新在梯度里多了一项rho * (xk - z u)这就是增广拉格朗日对x的导数也是近似ADMM区别于FedAvg本地更新的核心。第三步对偶更新使用残差x - z_new直接加到u上对应的是u_{k}^{t1}u_k^tx_k^{t1}-z^{t1}数学上等价于按残差累积对偶变量。3.4 收敛判据怎么设计才不“骗自己”很多人在做联邦学习实验时只看训练精度这是非常危险的。我见过好几次损失降了但模型崩了的情况——因为训练精度本身在非凸问题上有很大的波动性且容易过拟合到聚合前的本地形态。正确的做法是监控两个残差原始残差primal residual和对偶残差dual residual。原始残差是Σ||x_k - z||^2它衡量本地模型和全局模型之间的差距对偶残差是||z_new - z_old||衡量全局模型本身是否还在大幅漂移。只有这两个量都收敛到阈值以下才说明迭代真正进入了稳定区。上面的代码已经顺手把这两个残差记录了画出来就是两张非常有信息量的曲线。这里有一点经验值得强调不要用固定数字作为绝对阈值。数据规模大时残差的绝对值自然会大合理的做法是采用相对容差(1 norm(z))的缩放代码中dualRes的分母就是为了抵消量纲影响。primalRes也可以用norm(x)-ish来归一化实际工程中我会把两个残差做相对化后再与阈值比较。4. 实验对比与参数调优记录4.1 我跑的对比实验配置我的基线设置是特征维度20客户端数10总样本500本地迭代步数locSteps5通信轮数R30逻辑回归模型。对比三组方法方法本地更新方式全局聚合方式FedAvg普通SGD多步按样本数加权平均近似ADMM近端SGD多步含对偶补偿mean(x u) 对偶更新精确ADMM本地L-BFGS精确求解子问题mean(x u) 对偶更新4.2 三种方法的收敛行为观察实验做下来最直观的感受是三种方法在收敛轨迹上的性格差异很大。FedAvg在早期轮次冲得很快但后期会在最优值附近来回震荡——尤其在Non-IID划分下客户端之间的目标函数方向差异大平均出来的模型方向不稳定。精确ADMM则是最“稳”的30轮内几乎看不出抖动但代价是每轮通信前本地要跑数十次L-BFGS迭代单轮时间成本大约是近似ADMM的10倍以上。近似ADMM的表现介于两者之间它的收敛路径比FedAvg平滑得多而且单轮成本与FedAvg几乎一样——只多做了一次对偶变量更新和一次惩罚梯度项计算开销可以忽略。如果只看达到相同测试精度的通信轮数近似ADMM通常会比FedAvg少用三分之一的轮次。这和我一开始的感受互相印证通信轮数没有增加但训练曲线的稳定性确实改善了。4.3 ρ和本地步长怎么配才能又稳又快调参是ADMM类方法绕不开的环节ρ几乎决定了整个算法的收敛性格。我扫过0.01到10之间对数刻度的若干值记录下的规律非常明确ρ太小primal残差下降快但对偶变量容易震荡ρ太大本地模型被“钉”在全局模型附近模型灵活性下降收敛速度变慢甚至停滞。一个相对可靠的初始值是ρ在0.1到1之间然后观察primal残差和dual残差的比例——如果primal残差比dual残差大一个数量级以上说明ρ偏小要增大反过来则减小。本地步长locSteps同样重要。增广拉格朗日里的二次惩罚项天然带有收缩效应本地更新步数越多客户端越接近子问题的“精确解”但也会越多地消耗本地算力。我建议locSteps不要设得太大第一次跑用5步左右就够。步长α则要与ρ联动调整α偏大时惩罚项的梯度项会被放大容易出现梯度爆炸α偏小则本地更新太慢。实测下来α0.01配ρ0.1是一组相当稳的组合。参数之间不是孤立的。一开始我分别调α和ρ效果很糟糕后来才意识到这两个参数是协同的本地梯度是∇F ρ(x - z u)有效学习率大概是α * (局部Hessian ρI)。所以把α和ρ绑定调整行为更容易预测也更好复现。5. 踩坑记录与问题排查实录5.1 常见问题速查表现象可能原因解决办法训练损失下降但精度不高过拟合本地数据全局模型未跟上增大ρ减少locStepsprimal残差长期不降ρ太小对偶变量无法有效补偿调大ρ同时检查数据划分是否有单客户端占绝对多数对偶残差剧烈震荡α太大增广拉格朗日项权重波动降低α或用较小ρ起步所有客户端收敛到同一局部最优无全局正则约束过松在g(z)中加入L2正则项精确ADMM每轮成本过高本地子问题求解太精确换近似ADMMlocSteps5-10即可逻辑回归梯度符号反向标签映射错误确认y取值为{-1, 1}而非{0, 1}5.2 调试方法论先用小规模跑通再放大实验我在调试这套代码时走了不少弯路最值得讲的经验是“永远先从单客户端单步验证”。第一步只保留一个客户端locSteps设为1ρ设0这时近似ADMM应该退化成普通的SGD如果曲线不吻合说明更新公式里有问题。第二步加入惩罚项ρ(x - z u)检查primal残差是否按预期下降。第三步增加第二个客户端制造一个简单的不一致场景看对偶变量能不能自动把两个客户端拉到同一个z上。这样一步步拆解比在完整实验上一行行调参高效得多。另外实验过程中一定要固定随机种子。Non-IID划分的随机性对结果影响很大如果不固定很难区分参数调整带来的提升和数据集划分带来的波动。我一般会在代码开头固定rng并跑多次取平均结果。关于对偶变量初始化直接全零就行不推荐随机初始化。原因在于对偶变量的物理意义是“累积的历史偏差”从零开始对应的是“初始时刻客户端之间无偏差”的假设这是符合语义的。随机初始化会对前几轮迭代引入不必要的偏置尽管理论上能收敛但对早停策略不友好。收尾前的一个小建议我在实际应用中最后还有一点体会尽管这套近似ADMM逻辑清晰收敛稳定但它毕竟需要调ρ和α两个耦合参数上手门槛比FedAvg高。建议第一次接触的人不要急着上大数据集先用几百个样本的合成数据把primal残差和对偶残差的曲线画出来理解“ρ↔本地步长↔收敛速度”这三者的关系后再放大。另一个经验是如果后续打算迁移到神经网络模型非凸场景本地更新中一定要对梯度做裁剪或归一化否则惩罚项的梯度在深层网络中很容易让更新步越过合理范围。等这套流程跑通你会发现在Non-IID的联邦模拟里近似ADMM的稳定程度确实比FedAvg高不少而代码复杂度却只增加了一点点。