Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge

Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge

• 背景
  • 存在矛盾：联邦学习(FL)可能会对边缘节点的计算能力造成不必要的负担，尽管它由于其隐私和机密性特性对FL有很强的实际需求
• 引言
  • FedGKT设计了一种交替最小化方法的变体，在边缘节点上训练小型cnn，并通过知识蒸馏周期性地将它们的知识转移到大型服务器端CNN
  • 普通CNN的缺陷：众所周知，扩大 CNN 大小（例如宽度、深度等）是提高模型精度的有效方法。不幸的是，训练大型 CNN 对资源受限的边缘设备具有挑战性
  • 利用SL还剩存在缺陷：为了解决边缘节点的计算限制，基于模型并行的分割学习(SL)[11,12]对大型模型进行分区，并将神经体系结构的某些部分卸载到云端，但SL存在严重的掉队问题，因为单个小批迭代需要服务器和边缘之间的多轮通信。
  • 本文方法，取其精华，去其糟粕：FedGKT旨在通过使用FedAvg中的本地SGD进行训练来合并FedAvg[6]和SL[11,12]的好处，但也将低计算需求放在边缘，如SL
    • FedGKT的本质是它将FL重新表述为交替最小化(AM)方法[13,14,15,16,17,18]，它通过交替固定一个和优化另一个随机变量(边缘模型和服务器模型)来优化两个随机变量。
    • FedGKT 将多个优势统一到一个框架中： 1. FedGKT 是内存和计算效率，类似于 SL； 2. FedGKT 可以以像 FedAvg 这样的本地 SGD 方式进行训练，以减少通信频率； 3. 与像 FedAvg 中那样交换整个模型相比，交换 SL 中的隐藏特征减少了通信带宽要求。4. FedGKT自然支持异步训练，避免了SL中严重的同步问题。
• 相关工作
  • 与现有的FL（FedAvg [6]、FedOpt [26] 和 FedMA [8]）不一样的地方在，现有的方法在训练大型CNN存在障碍。
  • 最近的工作FedNAS [27, 3] 和 [28] 在大型 CNN 上工作，但它们依靠 GPU 训练来完成评估，而本方法不需要
  • 其他人[29,30,31,32,33,34,35,36,37]在不考虑边缘计算限制的情况下优化通信成本。
  • 基于模型并行的分割学习[11,12]试图打破计算约束，但它需要与服务器频繁通信
  • 与普通KD蒸馏不同的有在于：
    • 以前的工作只考虑将知识从大型网络转移到较小的网络[19，20，21，22]，或者他们从组中转移知识，但组中的每个成员共享相同的大型模型架构或具有特定尾部或头部层的神经架构的很大一部分
    • 蒸馏中的所有教师和学生共享相同的数据集 [50, 52, 53, 54]，而在我们的设置中，每个成员（客户端）只能访问自己的独立数据集
    • 以前的方法使用集中训练，但我们使用了交替训练方法
  • 我们的工作还涉及到边缘设备上的高效深度学习，如模型压缩[55,56,57]、手动设计的体系结构(MobileNets[58]、ShuffeNets[59]、SqueezeNets[60])甚至高效的神经架构搜索
• Group Knowledge Transfer技术介绍
  • 初步
    • 定义
      • 整个数据集：D
      • C个类别的监督学习
      • 第K个节点有自己的数据集

       

 

• • • • • Xi：第i个训练样本
        • yi：Xi的对应标签
        • yi ∈ {1, 2, . . . , C}：（多分类任务）
        • N (k)： 是Dk数据集中的样本数
    • 一般来说，我们可以将基于cnn的联邦学习表述为一个分布式优化问题

 

• • • • W 表示每个客户端中全局 CNN 的网络权重
      • f (k)(W) 是第 k 个客户端的局部目标函数（第 k 个客户端的损失函数），用于衡量异构数据集 Dk 上的局部经验风险。
      • l是全局 CNN 模型的损失函数
      • F(W)就是全局，通过N(k)/N进行权重平均化的总体损失函数
      • 加上min (W) 就是最小化这个损失函数 作为此函数的目标
    • 大多数现成的联合优化方法（例如 FedAvg [6]、FedProx [62]、FedNova [63] 和 FedOpt [26]）提出使用变体局部 SGD [10] 优化方法解决目标函数 (1)（上面的函数） 用于通信高效训练，并通过线性模型（逻辑回归）或浅层神经网络（2 个卷积层）的实验来展示它们的特征
    • 问题
      • CNN无法在边缘进行高强度计算
      • SL：有一个通病会导致落后者 (Straggler) 问题
        • All-Reduce虽然解决了中心节点通信瓶颈问题，但却因其流水线式的同步方式而遇到新的挑战：落后者 (Straggler)问题。Straggler问题也被称之为异质性 (Heterogenity)问题，指的是当一个平台上不同的Worker具有不同的运行性能时，整个系统的性能受限于运行效率最低的Worker。Straggler可能是确定性的，比如某台机器的计算能力或者网络带宽低于其他机器；也可能是随机性的，比如低性能是因为资源共享，后台操作系统任务，缓存，功耗限制等引起的。
        • 解决Straggler问题的一种常用方式是进行异步训练（所以本文用了异步去改进）。然而All-Reduce流水线式的通信结构对异步训练带来了很大的挑战。上节中介绍过的去中心化的结构更有利于异步训练的实施。发表在PMLR 2018上的AD-PSGD [4] 就是去中心化异步训练的一个代表算法。
  • 重新优化！
    • 非凸优化
      • 为了解决现有 FL 中的资源受限的问题，我们重新考虑了解决 FL 优化问题的另一种方法
      • 我们将Eq（上面的公式1）.(1)中的全局CNN W分为两个分区:一个小的特征提取器模型We和一个大规模的服务器端模型Ws，We放在边缘，Ws放在服务器上。
      • 我们还为 We 添加了一个分类器 Wc，以在边缘创建一个小型但完全可训练的模型
      • 因此，我们将单个全局模型优化重新表述为非凸优化问题，该问题要求我们同时解决服务器模型Fs和边模型Fc
      • argmin 求函数最小值
      • min求 集合最小值
      • 将公式（1）拆分

       

• • • • • 都是将损失函数最小化
      • ls服务服务器模型的损失函数， lc边模型的损失函数
      • fs是服务器模型，fk是边缘侧模型，fe(k)是特征提取器，fc(k)是分类器
        • fk由fe(k)和fc(k)组成
        • Ws We(k) Wc(k) 分别对应 fs fe(k) fc(k)的网络权重
      • 公式3：Hi(k)是特征提取器fe(k)输出的第i个样本的特征图（隐藏层向量或张量）
      • 公式5：这个公式可以在每个客户端独立计算
      • 第k个客户端模型f(k)在其局部数据集（等式5）上训练
      • 服务器fs使用Hi(k)作为输入特征进行训练
      • 客户端k的最终训练模型架构由特征提取器fe(k)的架构和服务器模型fs的架构堆叠
      • 优势和挑战
        • 当我们假设 f (k) 的模型大小比 fs 小几个数量级时，边缘训练是负担得起的
        • 对于大型 CNN 训练，将 H(k)i 传输到服务器的通信带宽大大小于传达传统联邦学习中所做的所有模型参数。
  • 使用知识转移扩展边缘数据限制
    • 将知识蒸馏损失纳入优化方程，以降低优化难度
    • 从服务器模型转移的知识可以提高边缘的优化
    • 我们建议双向转移Group Knowledge
      • 双向：服务器 CNN 从许多边缘吸收知识，单个边缘 CNN 从服务器 CNN 中获得增强的知识。
    • 实现双向知识转移的公式

• • • • lCE是预测值和真实标签之间的交叉熵损失。
      • DKL是 Kullback Leibler (KL) Divergence 函数，它作为损失函数 ls 和 lc 中的一个术语，用于将知识从网络转移到另一个。
        • KL 散度损失试图使软标签和地面实况彼此靠近。这样做时，服务器模型吸收从每个边缘模型中获得的知识。类似地，边缘模型试图将它们的预测更接近服务器模型的预测，从而吸收服务器模型知识以提高其特征提取能力。
      • 边缘模型 fk的概率预测

       

• • • • 服务器模型 fs的概率预测
      • 

       

• • • • 以上概率预测都是由calculated with the softmax of logits z
      • Zs和Zc(k)都是最后一个全连接层的输出
      • T是是 softmax 函数的温度系数
    • 改进交替最小化
      • 在将Eq.(6)和(7)代入我们的重构(Eq.(2)和(5))后，我们提出了一种交替最小化(AM)的变体来解决重新表述的优化问题如下:
      • 

       

• • • • W(k)是 We(k) 和 Wc(k) 的组合
      • AM 是凸非凸优化理论和实践中的求解器，交替优化两个随机变量
      • 在等式(8) ，我们固定 W (k) 并优化Ws几个 epoch。然后我们切换到等式(10) 固定Ws优化W(k)几个 epoch。这种优化发生在等式之间的多轮中。 (8) 和 (10) 直到达到收敛状态。
    • Key Insight
      • 交替最小化！！！
        • 本质是交替最小化使用跨所有边缘的知识蒸馏来简化优化，这在联邦学习中扩展了每个边缘的数据集限制
      • 我们使用仅基于基本事实和模型输出计算的局部交叉熵损失来实现这一目标，以及第二个损失使用边缘和服务器之间的 KL 散度，这有效地捕获了来自多个客户端数据集的贡献（知识）
      • 此外，每个最小化子问题都可以用SGD及其变体求解(e.g., SGD with momentum [64], ADAM [65, 66]).
    • 训练算法
      • 为了详细说明，我们在图1(a)中说明了交替训练算法FedGKT，并将其总结为算法1。在每一轮训练中，客户端使用本地SGD训练多个epoch，然后将提取的特征映射和相关logits发送给服务器。当服务器从每个客户端接收提取的特征和 logits 时，它会训练更大的服务器端 CNN。然后服务器将其全局 logits 发送回每个客户端。这个过程迭代多轮，在每一轮中，所有客户端的知识都被转移到服务器模型中，反之亦然。对于 FedGKT 训练框架，剩下的步骤是为客户端模型和服务器模型设计特定的神经架构。为了评估 FedGKT 的有效性，我们设计了基于 ResNet [1] 的 CNN 架构，如图 1(b) 所示。更多细节可以在附录 B.3 中找到。
• 实验
  • 主要是看了实验结果，包括结尾也参考了，对于通信代价的计算公式
    • 两个方法差距就在 软标签和梯度
• 讨论
  • 还有自我批评
• 复现
  • Batch Normalization是google团队在2015年论文《Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift》提出的。通过该方法能够加速网络的收敛并提升准确率。在通常搭建网络的过程中，第一步会将图像集进行预处理，将图像数据调整到满足某一分布规律，从而加速网络训练。但是调整过的图像通过卷积层后得到的feature map就不一定满足我们所期望的某一分布规律，所以Batch Normalization就是调整feature map，让每一层的feature map都满足均值为0，方差为1的分布规律。参考
  • 服务器模型自己的17年电脑跑不了，显存不够，但是又是以项目的形式，自己调了半天以后知道，
    • 客户端模拟的是ResNet8模型进行训练
    • 服务器模拟的是ResNet50模型进行训练
    • 现在的思路是学习ResNet将ResNet50的模型降低其深度
    • 还需要了解知识蒸馏
    • 还需要了解拆分学习 split learning 将网络结构进行拆分