联邦学习客户端数据量不均衡问题及训练表现观测方法咨询

客户端数据量不一致是联邦学习落地时的常见情况，并非必然导致训练失效，但如果不做针对性调整，会带来两类明确的负面影响：

• 聚合结果偏移：如果使用默认等权聚合，数据量极小的客户端很容易在本地训练时过拟合，带噪声的更新会拉低全局模型效果；如果按样本数加权聚合，数据量极大的客户端会完全主导模型更新方向，全局模型在小数据量客户端对应的数据分布上泛化性会明显下降。
• 收敛稳定性下降：不同客户端数据量差过大时，本地更新的梯度范数差距会达到几个数量级，聚合后的全局更新方向抖动明显，会拖慢收敛速度，极端情况下甚至会导致训练发散。

针对上述问题，以下方案均可在TensorFlow Federated框架内直接落地，无需魔改底层逻辑：

• 调整聚合权重逻辑：不要直接用默认的等权或纯样本数线性加权，可以通过tff.learning.aggregators自定义加权规则：比如用截断加权，将样本量超过总体95分位值的客户端权重截断到分位值，避免超大数据量客户端主导更新；或者用平方根加权，权重与客户端样本数的平方根成正比，平衡大小客户端的贡献占比。
• 对齐本地训练步数：不要给所有客户端设置固定的本地训练epoch数，改为固定每轮每个客户端的本地训练batch数，比如统一设置为每轮本地跑20个batch：小数据量客户端可以通过重复采样凑够batch数，避免过拟合；大数据量客户端只采样对应数量的batch参与本轮训练，避免更新幅度过大。这个逻辑只需要调整客户端数据预处理pipeline即可实现。
• 替换鲁棒聚合规则：如果数据量差异同时伴随数据分布的非独立同分布，可以将原生FedAvg替换为中位数聚合、修剪均值聚合，同时给客户端更新加范数裁剪，压制异常更新的影响，TFF的鲁棒聚合模块已经提供了上述方法的现成实现，直接替换训练流程的聚合器即可。
• 优化客户端采样逻辑：每轮选择参与训练的客户端时，不要使用完全随机采样，改为按客户端数据量分层抽样，保证每轮参与训练的大、中、小数据量客户端占比和全局分布一致，避免单轮采样偏差导致的更新抖动。

TensorFlow Federated提供了成熟的原生接口支持逐客户端训练表现观测，不需要修改框架核心代码：

• 自定义逐客户端指标：在定义模型的model_fn时，除了全局聚合需要的损失、准确率等指标，额外将需要观测的本地训练损失、本地验证集准确率、梯度范数、本地实际训练步数等指标加入metrics字典，配置对应指标为逐客户端返回、不做全局聚合即可。
• 抓取迭代过程的逐客户端输出：TFF的训练迭代过程基于tff.templates.IterativeProcess实现，每轮执行next()方法后返回的结果结构中，会保留本轮所有参与客户端的全量逐客户端指标，直接将这部分指标写入日志，关联客户端ID、样本量属性，就可以做后续分析。
• 配套分桶统计分析：拿到逐客户端指标后，可以按客户端样本量分区间统计对应区间的本地损失、准确率分布，直接定位数据量差异对训练效果的具体影响程度。

内容的提问来源于stack exchange，提问作者Alwani