该研究的主要作者包括 Jianqing Zhang(上海交通大学)、Yang Liu(清华大学人工智能产业研究院及上海人工智能实验室)、Yang Hua(Queen’s University Belfast)、Jian Cao(上海交通大学及上海市可信Web3数据流通与治理重点实验室)。这项研究发表在 2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) 的会议论文集中,DOI 为10.1109/CVPR52733.2024.01151。
随着个性化需求的增长,许多公司开发了符合自身需求的定制化模型。然而,由于某些领域(如医学)的数据限制和模型复杂性,模型训练往往面临数据不足和通信代价过高的问题。联邦学习(Federated Learning, FL)通过允许多方协作训练模型并保持数据隐私性提供了解决这一问题的新途径。
传统的联邦学习(TFL)主要关注为所有客户端训练一个统一的全局模型,无法满足数据分布非独立同分布(Non-IID)的情况下每个客户端的个性化需求。个性化联邦学习(Personalized Federated Learning, PFL)虽然能够生成定制化模型,但大多仍依赖于客户端之间的同构模型架构。然而,随着模型规模的增大,这种方法在通信成本和隐私保护上面临巨大挑战。
针对这一问题,研究提出了异构联邦学习(Heterogeneous Federated Learning, HTFL),允许多种模型架构协同训练。但HTFL在数据和模型多样性条件下共享知识的效率仍然较低,现有方法大多依赖于知识蒸馏(Knowledge Distillation, KD)技术和全局数据集,但这些方法在通信开销、知识传递质量和隐私保护上仍存在不足。为了应对这些问题,本文提出了一种新的高效上传知识迁移方案,称为“联邦知识迁移循环”(Federated Knowledge-Transfer Loop, FedKTL)。
研究目标是利用服务器端公开的预训练生成器,将其已存储的知识高效地迁移到客户端异构模型中,实现隐私保护、低通信成本和高精度的联合训练。
研究设计了 FedKTL 框架,在每次通信迭代中包括以下六个主要步骤: 1. 客户端经过本地训练生成类别原型(Prototypes)。 2. 客户端将生成的类别原型上传至服务器。 3. 服务器使用一个轻量化的特征变换器(Feature Transformer, F)将客户端上传的类别原型对齐到生成器的潜在输入空间(Valid Latent Domain)。 4. 基于对齐的潜在向量,服务器通过预训练生成器生成类特定的图像。 5. 服务器将生成的类别图像及其潜在向量对(Image-Vector Pairs)下发到客户端。 6. 客户端利用这些图像-向量对进行附加的监督学习任务,从而增强本地模型的特征提取能力。
为了消除非均匀数据分布对类别原型的偏倚,FedKTL 引入了等角紧框架(Equiangular Tight Frame, ETF)分类器。ETF 分类器作为理想分类器,利用固定的原型向量实现类间最大间隔,从而生成具有全局一致性和无偏性的类别原型。
在客户端的本地模型中,特征提取器和ETF分类器通过联合训练以最小化客户端本地数据的分类交叉熵损失。训练完成后,客户端即可生成类别原型并上传至服务器。
上传的类别原型位于ETF空间中,与生成器的有效输入空间不直接对齐。因此,研究设计了一个轻量化的特征变换器 F,通过最小化均方误差(Mean-Squared Error, MSE)及最大均值差异(Maximum Mean Discrepancy, MMD)损失函数对齐类别原型和生成器潜在空间。
一旦域对齐完成,服务器可以通过输入类中心(Class-Wise Latent Centroids)到生成器,生成清晰、具有类间判别性的图像。
客户端下载图像-向量对后,基于本地模型执行附加的监督任务。具体来说,客户端将生成器知识迁移到特征提取器,以提升提取潜在模式的能力,同时减少生成图像的语义与本地数据完全一致的要求,进一步实现隐私保护。
在CIFAR10、CIFAR100、Tiny-ImageNet和Flowers102等四个图像数据集上的实验表明,FedKTL相比7种主流HTFL方法(如LG-FedAvg、FedGen、FedProto等)在两种数据异质性场景下(Pathological Setting和Practical Setting)均取得了更高的测试精度,其中在CIFAR100的Practical Setting下精度提升最高达5.40%。
在不同模型异质性场景(如HTFE2到HTFE9,从2种到9种异构模型)下,FedKTL展现了更高的适应性,性能甚至随着模型异质性增加而继续提升,最高提高7.31%的精度。
FedKTL 仅需每类上传一个类别原型,即每类仅k(等于类别数)个向量元素,上传开销显著低于大多数其他方法(如FedGen、FedProto)。这种上传效率特别适合于现实场景中上行带宽较低的设备。
研究不仅在最新的StyleGAN-XL生成器上进行了实验,还验证了FedKTL对其他风格化GAN(StyleGAN3)生成器的适配性,如预训练于不同数据集(AFHQv2、Benches等)的模型,结果表明其知识迁移循环仍能在不依赖生成器语义一致性的前提下产生区分度高的类特定图像。
FedKTL 不需要共享客户端的私有参数或本地数据,上传的类别原型和生成图像均不包含敏感信息。此外,对生成图像的语义相关性要求较低,进一步减少了隐私暴露的风险。
研究为HTFL提供了一种全新的、高效的知识共享方式,克服了数据和模型异质性对知识传递的限制。这种框架打破了使用全局数据集或全局模型的依赖,在无数据共享的情况下实现了高效协作。
FedKTL 在数据稀缺、带宽受限、隐私保护要求高的场景中具有重要应用意义,特别是在医学、生物信息、金融等数据敏感性较高的领域,为企业开发定制化模型提供了新的途径。
研究为解决HTFL长期以来的多重难题提供了全新思路,推动了隐私友好型分布式学习领域的发展。