模型异构联邦学习这几年逐渐从论文里的炫技变成了落地时的刚需。做联邦的人应该都撞过同一堵墙客户端设备五花八门算力参差不齐你想要性能好的设备用大模型又想老旧设备跟它一起训练经典联邦平均FedAvg那一套要求所有客户端用同一个模型结构结果只能让高配设备陪着低配跑小模型非常浪费。pFedES 这个方向就是在解决这个问题全称很直观Generalized Proxy Feature Extractor Sharing for Model Heterogeneous Personalized Federated Learning。简单说它让结构不同的客户端共享一个轻量的代理特征提取器来交换知识同时保留各自的本地模型和个性化头部既绕开了模型异构又照顾到了数据分布不同带来的个性化需求。这篇文章我会从背景动机、核心机制、代码复现思路、实验评估到调参避坑一条线拆完。适合正在研究模型异构联邦、准备复现个性化联邦算法的同学也适合被端侧模型不统一搞得头疼的工程朋友。读完你至少能回答三个问题代理特征提取器到底在解决什么它在服务端和客户端之间怎么流转落地时最容易踩哪些坑1. 背景与动机为什么一定要做模型异构个性化联邦1.1 同构联邦的尴尬算力差异被强制拉平传统联邦学习的出发点很纯粹数据不动模型动客户端在本地算梯度服务端汇总更新。FedAvg 能成立有一个隐藏前提——所有参与者使用结构完全一样的模型这样权重才能逐层做平均。但真实场景里这个前提几乎不成立。某公司要在一批手机设备上做联邦学习旗舰机的内存和算力足够跑一个 ResNet 级别的模型老机型可能只带得动轻量 CNN。强行走同构要么所有人都用最弱模型性能天花板被拉低要么老设备根本无法参与训练可用数据量下降。模型异构带来的第一个问题是权重无法对齐。两个网络中间的卷积核个数、隐藏层维度、归一化参数都不一样参数张量形状不匹配逐层平均根本没有数学意义。就算强行用一些小 trick 把维度补全特征空间的语义也完全对不上。A 模型第 3 层学到的边缘纹理和 B 模型第 3 层学到的语义概念很可能不是一个东西平均之后既不是 A 也不是 B。1.2 更麻烦的是数据分布Non-IID 下的个性化需求模型异构之外还有第二重压力每个客户端的数据分布不是独立同分布的。用术语讲叫 Non-IID现实中表现为标签偏移、特征偏移、数量不平衡。不同客户端可能是不同地区门店的设备记录的数据风格不同也可能是不同个人使用的手机产生的行为数据差异很大。这种场景下全局最优解往往不是每个客户端的最优解。你希望学到的是大家都认可的公共知识同时每个客户端在公共知识之上再做本地适应。这就是个性化联邦学习的意义。于是问题变成了模型异构和个性化这两个需求同时放在一起应该怎么设计方法pFedES 给出的答案是公共知识可以放在一个代理特征提取器上个性化知识留在本地模型里。1.3 现有技术路线的三条岔路与共同短板过去几年模型异构联邦的基本思路大致分成三路。第一路是正则化与局部微调路线。FedProx 在 FedAvg 基础上加了客户端本地模型与全局模型的距离惩罚Per-FedAvg 引入元学习的思想。问题在于它们大多还是假设或者弱化模型同构真正完全异构时不好用。第二路是知识蒸馏路线。FedMD、FedDistill 这类方法通过公共数据集或者代理数据在客户端之间传递软标签知识。它们可以突破结构限制因为蒸馏的是模型输出而不是参数。缺点是要么需要一个大家都能访问的公共数据集这在隐私敏感场景里很难办到要么软标签质量高度依赖公共数据与实际数据分布的相似程度公共数据和真实业务数据差太远蒸馏出来的知识就歪了。第三路是表征空间对齐路线。想方设法让不同结构的模型将样本映射到一个共享向量空间比如固定某一层的宽度或者设定一个公共投影头。这类方法对模型结构的自由度限制比较多不算真正的任意异构。这几路都各自有用但共同短板是没有把共享什么这件事做得既轻又通用。要么需要额外数据集要么要求模型部分结构对齐要么只适合窄范围的异构。pFedES 的思路是从共享一个代理特征提取器这个角度切入的。1.4 pFedES 的切入点共享代理保留个性pFedES 的核心词汇是 Proxy Feature Extractor代理特征提取器。你可以把它理解成一个轻量级的公共翻译官网络。它本身不需要与任何客户端模型结构一致因为它的任务不是替客户端干活而是承接所有客户端想要共享的知识。具体到机制上每个客户端在本地训练时同时跑自己的本地模型和一个代理特征提取器。真实标签用来监督本地模型代理特征提取器通过知识蒸馏的方式学习本地模型的输出分布。客户端之间不共享原始模型参数只共享代理特征提取器的参数。服务端把各个客户端更新过的代理特征提取器做加权平均再分发回去。这样一轮一轮下来公共知识在代理提取器里越来越丰富本地模型又始终保留个性化结构。标题里的Generalized我理解就是在强调这个方法不对客户端结构做限定不要求公共标签也不依赖集中式代理数据集扩展性比早期蒸馏方案强不少。还有一个容易被忽略的细节用户没有直接上传原始特征向量或者原始模型权重而是上传一个经过蒸馏约束的轻量提取器这在隐私表现上比裸传梯度要好一截。当然说完全没泄露是骗人的这一点后面我会专门谈。2. 核心机制拆解代理特征提取器是怎么共享的2.1 先理解特征空间对齐的必要性模型异构之所以难核心在于特征空间不对齐。我说一个特别简单的类比两个人分别用中文和英文写同一篇新闻内容是一样的但你不能把他们的句子逐字相加成一篇新文章。神经网络里的特征表示就好比语言每个模型的隐藏层有自己的词汇表和语法直接平均参数等于把中文句子和英文句子按字硬拼。pFedES 的办法是引入一个统一编码的翻译官。所有客户端在本地都加载同一个代理特征提取器输入同一批样本输出一个固定维度的特征向量和对应的软标签。这个代理特征提取器就像约定好的通用语言客户端之间通过它来对齐语义而不是强行对齐网络权重。共享的不是结构而是对不同样本应该如何提取信息的能力。2.2 客户端本地更新的三段式结构每一个参与训练的客户端本地实际同时维护两个网络任务模型local model和代理特征提取器proxy。任务模型的结构任意可以是一个三层小 CNN也可以是带残差连接的模型服务端完全不关心。代理特征提取器是全局统一的所有客户端从服务端拿到的都是同一份参数。一次本地前向计算会走两条路径。第一条路径样本输入本地模型得到本地特征再进本地分类头得到本地预测 logits这个 logits 用真实标签做交叉熵监督。第二条路径同样一批样本输入代理特征提取器得到代理特征和代理 logits代理 logits 的监督信号不是真实标签而是本地模型对这批样本的预测分布。换句话说代理特征提取器学的是任务模型是怎么看这批数据的。这里面有一个很容易踩坑的细节当代理 logits 在做知识蒸馏学习时本地模型的 logits 应该被视为固定目标stop-gradient只让代理特征提取器的参数更新。如果不这样做两个模型会互相追逐训练目标一直在移动很容易发散。实现上就是用 detach() 把本地 logits 从计算图里摘出来。2.3 服务端聚合与分发异构模型的关键解耦点服务端做的事情非常干净把各客户端返回的代理特征提取器做加权平均。权重一般用客户端本地样本数量这和 FedAvg 的加权逻辑一致。因为唯一被聚合的结构是代理特征提取器它的结构是全局统一的所以平均操作是完全合法的。本地任务模型从头到尾不上传参数服务端根本不知道也不关心客户端用的是什么结构。这个设计把问题和方案做了一个漂亮的解耦异构问题交给了各管各的本地结构来解决协作问题交给了公共代理提取器来解决。这也是它在工程上很吸引人的原因——不需要客户端把自己的模型结构上报更不需要为了对齐而改模型设计老设备和强设备各跑各的只在代理提取器层面汇合。2.4 为什么用蒸馏而不是共享特征向量有的方法会让客户端直接上传特征向量到服务端做全局对齐乍看更直接但有三个问题。第一特征向量里可能残留身份信息和敏感信息中间层的表征经常能反推出原始数据的一些属性隐私风险不可控。第二不同客户端的特征向量分布天然不同服务端拿到一堆分布漂移很大的向量很难统一处理。第三原始特征向量是稠密的通信开销比只传一个轻量模型还大。蒸馏方案传递的是类别关系这种抽象知识。真实标签是硬目标本地模型输出的软分布是软目标软目标里包含了类间相似度的信息。比如一张猫的图片模型可能认为它 80% 像猫、15% 像狗、5% 像兔子这 15% 和 5% 恰恰是知识迁移的关键。代理提取器学到的就是这个细粒度的类间关系而不是某个客户端的具体特征向量。蒸馏温度 T 在这里起放大或压缩类间关系的作用。温度越高输出分布越平滑类间相对差异保留得越完整温度越低分布越接近 one-hot知识越接近真实标签。联邦场景里我建议从 T3 开始试太低不行太高会把噪声也放大。2.5 关于广义的理解边界标题里 Generalized 给我的信息是它试图覆盖三种常见的异构限制。一是客户端模型结构异构从卷积网络到多层感知机都能参与不影响知识传递。二是标签空间异构每个客户端即使只有部分类别的样本代理提取器依然可以借助本地软标签把这些类别的判别知识带进公共体系。三是不依赖公共数据集传统蒸馏联邦常常需要一个公共代理数据集来对齐pFedES 把公共数据集这个外部依赖干掉改成在客户端本地数据上做蒸馏实际落地时省去很多隐私合规上的麻烦。但广义不是无限广义。如果客户端之间的数据完全分布在不同的域比如一个是医疗影像一个是自然图像类别又完全不重叠任何知识蒸馏方法都很难传递有效信息。理解它的边界比吹捧它的能力更重要。3. 从论文到实现我复现 pFedES 的方案与代码骨架3.1 环境、数据切分与异构模型模拟先说环境。我自己做复现时用的是 PyTorch因为动态图和 detach 操作在实现蒸馏逻辑时很方便。数据方面经典实验通常用 EMNIST 或 CIFAR-10/100。模拟联邦 Non-IID 最常见的手段是狄利克雷分布采样为每个类别生成一个分布每个客户端按这个分布抽取样本参数 alpha 越小每个客户端持有的类别越集中Non-IID 程度越高。建议同时测 alpha0.5 和 alpha0.1 两个档位。异构模型模拟要克制。不要一上来就给 20 个客户端各配一种结构实验解释性会变得很差。我习惯把客户端分成三组一组用两层小 CNN一组用带残差的稍大 CNN一组用 MLP。三组结构的特征维度可以完全不同但最后分类的类别数必须一致因为分类头输出的 logits 要参与蒸馏。分组的好处是能单独看每组模型的个性化提升幅度而不是只看一个平均数字。3.2 代理特征提取器的结构选择与初始化代理特征提取器的参数量是第一个要平衡的点。它太大低算力客户端每轮要多算一份前向和反向负担重它太小承载不了异构模型蒸馏出来的知识成为信息瓶颈。我用过的最省心的配置是一个轻量 CNN两到三个卷积块加一个全局池化然后接一个固定维度比如 128 维的嵌入层和一个全连接分类头。下面是示意代码只保留关键逻辑class ProxyExtractor(nn.Module): def __init__(self, in_channels, embed_dim128, num_classes10): super().__init__() self.encoder nn.Sequential( nn.Conv2d(in_channels, 32, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.AdaptiveAvgPool2d(1), ) self.embed nn.Linear(64, embed_dim) self.head nn.Linear(embed_dim, num_classes) def forward(self, x): feat self.encoder(x) feat feat.view(feat.size(0), -1) embed self.embed(feat) logits self.head(embed) return logits, embed代理特征提取器要和客户端任务模型的输入大小匹配数据是图像就用卷积是表格或者文本就用 MLP 或 Transformer 的编码器这个按场景换即可核心结构不变。客户端本地模型就不强制统一只要能输出 logits 就行class LocalModel(nn.Module): def __init__(self, backbone, classifier): super().__init__() self.backbone backbone self.classifier classifier def forward(self, x): feat self.backbone(x) return self.classifier(feat)这里的 backbone 不同客户端各自定义代码里不需要任何特殊处理。3.3 一次通信轮次内的完整训练流程一轮联邦训练的流程大概是五步。第一步服务端向本轮选中的客户端分发当前全局代理特征提取器参数。第二步客户端加载本地模型和代理提取器在本地数据上训练若干个 epoch。第三步客户端把更新后的代理提取器参数返回服务端本地模型参数留在本地。第四步服务端按样本量加权平均所有返回的代理提取器。第五步重复下一轮。客户端内部的优化目标是最关键的地方。我用一个总的 losslocal_logits local_model(x) proxy_logits, _ proxy_model(x) soft_target F.softmax(local_logits.detach() / T, dim1) kd_loss F.kl_div( F.log_softmax(proxy_logits / T, dim1), soft_target, reductionbatchmean, ) * (T ** 2) ce_loss F.cross_entropy(local_logits, y) total_loss ce_loss alpha * kd_loss注意local_logits.detach()这一步它保证了交叉熵只更新本地模型KL 散度只更新代理提取器。两个模型虽然在同一条数据上训练但监督信号的流向是完全分开的。这样设计的目的前面说过是为了避免两个模型互相追逐导致的不收敛。本地训练 epoch 数建议 1 到 5。epoch 太多代理提取器在单个客户端上过拟合严重服务端聚合后震荡epoch 太少软标签还没学好就交回去了蒸馏质量差。我测下来 3 个 epoch 在多数基准上是个稳妥值。3.4 服务端聚合、EMA 平滑与通信开销控制服务端聚合不只是简单平均我做了一个小改进对代理特征提取器做指数滑动平均让全局代理参数的更新更平滑。global_proxy 0.9 * global_proxy 0.1 * aggregated_proxy这个 EMA 阻尼在客户端数据分布差异很大的时候很有用。直白点说如果这一轮碰巧都是 Non-IID 程度极高的用户它们的代理提取器更新方向会很偏EMA 能把这种偶然偏差稀释掉不让全局参数剧烈抖动。通信开销方面因为上传下载的始终是同一个轻量代理提取器而不是各种不同结构的客户端模型所以开销是恒定可控的。你甚至可以提前算出来代理提取器参数量乘以 4 字节再乘 2上下行就是单客户端单轮的理论通信量。对比一下 FedAvg 要传整个全局模型在异构场景下这个优势特别明显。3.5 评测方案别只看一个平均准确率评估个性化联邦学习最忌讳只报一个全局平均测试准确率。我建议至少看三个指标。第一是个性化准确率每个客户端在本地保留的测试集上分别评估本地任务模型然后取加权平均。这是衡量个性化的核心指标。第二是全局代理提取器准确率拿代理提取器在服务端汇总测试集上评估。它可以反映公共知识的质量。如果代理提取器精度一直很低说明蒸馏没有传递出有效信息这时候去查软标签是否失真、温度是否合理。第三是异构友好度对比各个模型组小 CNN 组、大 CNN 组、MLP 组的准确率差距。理想情况下大模型组能保留容量优势小模型组通过蒸馏从公共知识里受益而不是被同构基线拖到一个水平线上。还要加一个稳定性指标多次随机种子实验的方差。模型异构 Non-IID 的联邦实验方差普遍比同构实验大如果只跑一次就下结论很容易被随机性欺骗。4. 为什么有效结果解读、收益边界与陷阱4.1 三类典型实验结论的解读从公开结果和我自己跑过的模拟实验看pFedES 这种架构在三个维度上都有比较明显的收益。第一个维度是个性化表现。由于每个客户端保留了自己的本地模型和本地分类头模型可以针对本地数据分布做专门的调整平均个性化准确率比强制同构 FedAvg 高是正常的。这个收益本质上是模型结构自由 头部个性化带来的代理提取器在中间保证了大家不至于各学各的学到完全分裂。第二个维度是通信效率。传统 FedAvg 每轮传输整个模型模型异构场景下甚至要传输多个副本。pFedES 只传一个轻量代理提取器通信量和客户端模型结构无关预算容易控制。在端侧网络环境不稳定的场景里这个优势往往是能不能落地的关键。第三个维度是收敛稳定性。相比直接在异构特征空间上做匹配蒸馏带来的软目标更平滑梯度方向更稳定。我观察到的现象是早期轮次代理提取器精度快速爬升本地模型也随之受益后期即使数据分布偏斜也很少出现剧烈的精度来回震荡前提是超参没有设坏。4.2 收益来源的机理分析为什么一个看起来简单的代理提取器能承担这么大的作用我认为有三个机制在同时起作用。第一隐式公共特征空间的形成。所有客户端共享同一个代理提取器它不断吸收各个客户端模型对数据的看法实际上是在异构结构之间建立了一个统一的中间表征空间。本地模型虽然结构各异但都在通过蒸馏约束向这个空间靠拢等于有了一条软性的对齐通道。第二蒸馏本身的正则效应。软标签提供了比 one-hot 标签更丰富的监督信息它对本地模型的训练也有一定的正则作用。在客户端本地数据很少的时候这种正则能明显减轻过拟合。这可以理解为代理提取器把其他客户端的知识以平滑约束的形式传给了本地模型。第三容量错配下的知识补偿。大模型能从数据里学到更细腻的特征这些特征通过软标签传给代理提取器代理提取器再间接影响小模型。小模型因此获得了超出它自身容量限制的部分学习经验虽然不能完全复现大模型的能力但比闭门造车强得多。4.3 必须警惕的收益边界再好的方法也有边界。代理提取器的容量是第一条边界它太小时所有异构模型的知识被压进一个狭窄的瓶颈蒸馏出来的软标签信息严重降级它太大时低算力客户端训练负担加重甚至不如直接跑完整模型。第二条边界是数据分布的偏移程度如果客户端之间不仅标签分布不同连输入域都不同代理提取器很难同时讨好所有域最终可能两头不讨好。第三条边界在隐私上。虽然上传的是代理提取器而非原始特征或本地模型但代理提取器是在真实本地数据上训练的理论上仍然可能编码部分数据分布信息。真要用于强隐私场景需要配合差分隐私噪声、梯度裁剪或者安全聚合。实践中最常见的三个误解我也列一下有人把代理提取器当成全局模型来用拿它做服务端推理这其实是误用代理提取器的本职是知识传递不是最终推理模型有人以为实现起来就是普通蒸馏加联邦忽略了 stop-gradient 和聚合频率这些细节结果训练不稳定还有人只在 CIFAR-10 这种简单数据集上看到提升就下结论真实场景里模型异构往往伴随域偏移必须单独做域偏移实验验证。5. 实战中的坑与调参经验给你一份避坑清单5.1 精度不升反降时的排查顺序模型异构联邦实验出现负优化的情况不少见而且原因往往不是单一的。我把自己常用的排查顺序整理成了一张表。现象可能原因排查与解决办法本地模型精度和单机训练差不多代理提取器没起作用检查 kd_loss 是否真的在下降去掉 detach 后是否发散alpha 是否设得太小代理提取器不收敛本地 epoch 过多或 T 过低减少本地 epoch把温度升到 3 到 5确认 soft_target 没有退化成 one-hot全局聚合后精度震荡客户端数据 Non-IID 太重使用 EMA 平滑提高参与率对代理更新做梯度裁剪小模型组没受益代理提取器容量不足增加 embed_dim 或卷积通道数检查小模型输入分辨率和代理提取器是否匹配大模型组反而掉点alpha 太大、蒸馏约束过强降低 alpha 到 0.1 档次只让代理提取器学本地模型不要让本地模型被反向影响这里有一个我自己的经验先让代理提取器在一个客户端上单独训练几个 epoch如果它连单个客户端的数据分布都没法学出有区分度的特征那问题很可能出在代理提取器结构或输入预处理上先修这个再上联邦。5.2 聚合与通信的精细化控制服务端聚合策略可以做得更细。参与率方面建议每轮至少采样 20% 的客户端样本量太少的轮次加权平均的方差会非常大。聚合权重方面直接用样本量加权简单有效但要注意如果某个客户端本地样本数量特别大它的代理提取器更新会在聚合里占主导这时候可以做个权重上限截断比如单个客户端权重不超过 0.3。梯度裁剪也值得加。客户端本地训练结束后计算代理提取器参数的更新量如果更新量的 L2 范数超过阈值我常用 5.0按比例缩放。理由是极端 Non-IID 客户端可能产生非常大的更新不加裁剪会把全局代理提取器带偏。本地步长和多轮聚合的配合也要留意。客户端本地训练 1 个 epoch 时每轮通信可以很频繁但服务端聚合不能太激进本地训练 3 到 5 个 epoch 时聚合间隔可以适当拉大。我通常的做法是固定每轮客户端本地跑 3 个 epoch然后每 5 轮做一次全局 EMA 更新相当于给公共知识一个沉淀期。5.3 隐私增强的方向前面反复强调代理提取器不是绝对隐私安全。如果需要在敏感场景部署我推荐从三个方向叠加防护。第一梯度裁剪配合噪声注入。客户端在返回代理提取器更新前先裁剪更新范数再加一定量的高斯噪声这相当于在联邦学习里引入差分隐私的标准做法。代价是蒸馏软标签的信噪比下降精度会有所损失需要根据隐私预算重新调 alpha。第二安全聚合。服务端通过安全多方计算协议聚合所有客户端的代理更新服务端只能拿到加总后的结果看不到单个客户端的代理提取器。这解决的是聚合过程中的隐私泄露。第三只上传蒸馏统计量而不是完整代理参数。某些场景可以只传每个类别的平均 soft prediction服务端用这些统计量做聚合和下发。这种做法的特征是更粗糙知识传递效率会打折扣但隐私保护更直观。实际选哪种取决于合规要求甚至强于精度要求。5.4 场景迁移的扩展思考pFedES 这类架构扩展性很强我在自己的模拟项目里试过几个方向简单分享。一个是跨域联邦。某医疗场景中不同参与方的成像设备、扫描参数、病例分布都不一样而且各方往往各有一套自己的诊断模型。把 pFedES 的代理提取器换成它们能统一加载的轻量编码器再配合本地诊断头和蒸馏是可以做到资源不出本地知识共同进步的。这个方向我特别看好因为它精准命中了结构异构 数据私有 数据分布偏移三个痛点。另一个是端侧小模型与大模型协同。如果把云端大模型看作一个特殊客户端它的软标签可以下发给代理提取器端侧小模型又通过代理提取器间接受益这等于建立了一条从大模型向小模型的知识传送带。在边缘设备算力受限、又要享受大模型表达力的场景里这个思路比直接给端侧部署大模型可行性高得多。还有一个偏前瞻的方向把代理提取器和联邦半监督结合起来。现实中很多客户端只有少量有标签数据、大量无标签数据代理提取器可以先用有标签数据学出一个基础表示再利用无标签数据的伪标签参与蒸馏。这属于扩展研究不是 pFedES 的原生能力但设计空间是打开的做研究的朋友可以往这个方向挖。最后说说我复现下来的真实体会pFedES 不是那种惊艳到让你拍桌子的方法但它是一个思路非常干净、工程上可落地的架构。我复现过不少模型异构联邦的方案pFedES 最让我欣赏的一点是它把问题拆得足够清晰协作的事交给代理提取器个性的事交给本地模型各司其职。实际跑下来我最深的体会是超参之间是耦合的。温度 T 影响软标签平滑度alpha 控制蒸馏强度代理提取器容量决定信息瓶颈这三者必须一起调。我建议的起点是T4alpha0.1代理提取器嵌入维度 128本地 epoch3服务端 EMA 衰减 0.9。先跑通再往两个方向探索alpha 往上走观察大模型组是否被压制代理容量往上走观察小模型组是否继续受益。另外每轮记录代理提取器在服务端测试集上的精度把它当作训练健康的晴雨表——它连续多轮不涨说明知识传递出问题了这时候优先查软标签和蒸馏梯度。最后再分享一个小技巧无论你最终的项目是什么先在白盒环境里把每个客户端本地模型对同一批测试样本的 logits 分布打印出来对比它们和代理提取器 logits 分布的 KL 距离。距离一直在减小说明共享是真的发生了距离纹丝不动说明你的代理提取器只是在跟着本地模型瞎转蒸馏没有真正建立联系。这个观察比任何指标都更能帮你判断架构是否搭对。