FlexTrain: A Dynamic Training Framework for Heterogeneous Devices Environments

• 作者/机构: Mert Unsal, Ali Maatouk, Antonio De Domenico, Nicola Piovesan, Fadhel Ayed (均为 Huawei)

A1 主要贡献

核心问题与研究目标：
随着深度学习模型变得越来越大，在异构设备环境中部署它们带来了巨大挑战。大型模型难以部署在低功耗或资源受限的设备上，会导致推理时间长和能耗高。在联邦学习（Federated Learning）场景下，这些挑战会加剧，因为大型模型会导致设备与中央服务器之间产生高昂的通信开销，收敛速度变慢，并且由于部分资源有限的设备无法参与训练，最终模型可能存在偏差。

创新点与主要贡献：
为应对这些挑战，本文提出了 FlexTrain，一个专为异构设备环境设计的动态训练框架，它在训练阶段就能适应不同设备的多样化存储和计算资源。FlexTrain 的核心是两项关键技术：

1. 主动层采样 (Active layers sampling)：在训练过程中动态选择残差网络（如 ResNets 和 Transformers）的一个子集层进行训练。这使得训练出的模型能够更好地适应并部署在不同能力的异构设备上。
2. 自动蒸馏 (Auto-distillation)：将在较大模型配置上学到的知识蒸馏到较小的模型配置中，从而提升在低容量设备上的模型性能。

通过这些技术，FlexTrain 旨在实现以下目标：
* 高效部署：训练一个单一的全局模型，该模型可以轻松地部署到具有不同资源限制的设备上，从而节省训练时间和能耗。
* 性能提升：在联邦学习环境中，FlexTrain 通过有效分享从能力较弱的设备上学到的知识给能力更强的设备，从而超越了标准的联邦学习基准。这种优势在设备数量众多和数据非独立同分布（non-i.i.d.）的场景下尤为明显。

本文通过在 CIFAR-100 数据集上的实验证明了 FlexTrain 在集中式训练中的有效性，并将其扩展到联邦学习场景，在 CIFAR-10 和 CIFAR-100 数据集上展示了其相对于标准联邦学习基准的优越性。

A2 方法细节

在本节中，我们首先描述 FlexTrain 中使用的逐层采样和激活技术，这些技术使我们能够实现所需的灵活性。然后，我们介绍我们提出的自动蒸馏方法，该方法促进了各种采样模型学习到的特征之间的知识交换，从而实现了更好的整体性能。借助这些技术，我们可以创建一个高度优化的模型，该模型可以高效部署，同时在对时间敏感的环境中最大限度地降低通信成本，并确保与各种设备无缝集成。

2.1 主动层采样

模型与优化目标。考虑一个数据集 $D = X \times Y \subseteq R^{d_x} \times R^{d_y}$，由 n 个（输入，标签）对 $\{(x_i, y_i)\}_{1 \le i \le n}$ 组成，以及一个属于给定架构族的神经网络模型 $N$

$$\mathcal{N}^\theta = \{y_{out}(\cdot; W) : \mathcal{X} \to \mathcal{Y} | W \in \mathcal{W}^\theta\}$$

该模型由参数空间 $W^{\theta}$ 参数化，其中 $\theta \in \Theta$ 是一个固定的超参数，代表与架构相关的量，例如网络深度 $K$。设 $\ell : R^{d_x} \times R^{d_y} \rightarrow R$ 是一个损失函数，例如二次损失、交叉熵损失等，我们将单个样本 $(x, y) \in D$ 的模型损失定义为

$$\mathcal{L}(x, y ; W)=\ell\left(y_{o u t}(x ; W), y\right),$$

其中 $W = (W_k)_{1 \le k \le K}$，$W_k$ 是第 $k$ 层的可调参数。给定上述损失函数，学习过程的目标是找到最小化经验风险的参数 $W$

$$\min_{W \in \mathcal{W}^\theta} \mathcal{L}(W) \triangleq \frac{1}{n} \sum_{i=1}^n \ell(y_{\text{out}}(x_i; W), y_i).$$

大模型的部署挑战。方程(3)中陈述的最小化问题通常通过数值技术来解决，其中基于梯度的方法，如随机梯度下降（Stochastic Gradient Descent）【【14】A stochastic approximation method，1951】和 Adam【【15】Adam: A method for stochastic optimization，2014】是常见的选择。为了实现低泛化误差，根据神经网络的幂律缩放预测，需要一个在权重矩阵 W 中具有大量参数的深度网络。例如，最先进的基于 Transformer 的语言模型，如 LLaMA【【16】Llama: Open and efficient foundation language models，2023】和 PaLM【【17】Palm: Scaling language modeling with pathways，2022】，分别包含高达 650 亿和 5400 亿个参数，并由 80 和 118 个层组成。这种机器学习模型日益增大的趋势在其他任务中也很明显，例如视觉【【18】Machine learning model sizes and the parameter gap，2022】。例如，ResNet-110 是一个流行的图像分类架构，包含 170 万个参数【【3】Kaiming He, et al. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016.】。这些模型的大尺寸给在边缘设备上的部署带来了巨大障碍，特别是在可能需要多个模型的时间敏感环境中。因此，设计有效的技术来在这些环境中部署这些模型至关重要。

FlexTrain的层采样策略。解决这一挑战的一个潜在方案是独立训练不同大小的模型，然后根据边缘设备的能力进行部署。例如，LLaMA 模型提供了从 7B 到 65B 参数的多种尺寸选项，从而允许更灵活的部署选择【【16】Llama: Open and efficient foundation language models，2023】。然而，由于相关的成本和环境足迹，训练多个大型模型很快会变得不切实际。FlexTrain 是一个解决此问题的方案，它在训练过程中对同一神经网络 $N$ 的各种配置进行采样，以创建一个与完整模型相比精度损失最小的灵活模型。这些配置是通过从最深层开始选择性地停用一定数量的层来获得的。这种从深到浅的停用方法旨在允许在训练过程中学习基本特征，并使更深的层能够在此基础上构建并学习更复杂的特征。这使得我们的方法特别适用于残差神经架构，如残差网络（ResNets）和 Transformers，它们在现代深度学习应用中被广泛使用。因此，在本文的其余部分，我们将重点关注 ResNets 和 Transformers。图 1 直观地展示了这一概念，其中蓝色和白色块分别代表活动和非活动块。

FlexTrain的优化目标。考虑到这一点，我们让 $N_k$ 表示包含前 $k$ 层及其相应权重 $W_k$ 被激活的神经网络，而其余的块则作为恒等函数。我们使用 $\tilde{W}_k$ 来表示前 $k$ 层的权重，包括预处理层和决策层（例如，分类层）。接下来，我们让 $\pi = [\pi_1, \dots, \pi_K]$ 表示网络层在集合 $\{1, \dots, K\}$ 上的激活分布。具体来说，$\pi_K$ 表示在任何训练周期中所有层都被激活的概率，而 $\pi_1$ 表示只有第一层被激活的概率。这种特定分布的选择与边缘设备的约束密切相关。这是因为当某个特定的模型配置被更频繁地选择时，它可以在只能处理该配置的设备上获得更好的性能。通过仔细选择最适合边缘设备限制的分布，我们可以确保整个系统的最佳性能和效率。例如，$\pi_K$ 的一个可能选项是将其设置为能够承担部署完整模型的设备比例。鉴于此，可以得出结论，FlexTrain 旨在最小化以下损失函数：

$$\overline{\mathcal{L}}(W) = \mathbb{E}_{k \sim \boldsymbol{\pi}}[\mathcal{L}(\tilde{W}_k)] \triangleq \sum_{k=1}^K \pi_k \mathcal{L}(\tilde{W}_k).$$

优势1：减少训练时间。这种方法有两个主要好处：首先是减少训练时间。训练时间的减少是由于前向和后向传播的路径更短，这意味着训练时间不再与网络的完整深度成比例。相反，它与网络的较短预期深度成比例。这可以通过将任何训练周期的预期调整参数定义为 $r = \sum_{k=1}^K \pi_k A_k$ 来看出，其中 $A_k$ 表示构成 $\tilde{W}_k$ 的参数数量，而 $A$ 是神经网络 $N$ 的可调参数总数。鉴于只要至少有一个设备无法运行完整模型，就有 $\pi_K < 1$，可以得出 $r < 1$ 的结论，这表明训练时间确实减少了。

优势2：灵活部署。在训练过程中，会遇到不同的模型配置，这使得部署变得简单。如果一个设备的存储和处理能力有限，只能容纳大小为 $A^*$ 的模型，那么就部署已训练模型中满足此约束的前 $k^*$ 层。如果设备的能力随时间提高，可以向设备提供额外的层，而无需对先前交付的 $k^*$ 层进行任何更新。从另一个角度看这种灵活性，设备可以在等待更大型号交付的同时，从使用小型号中受益，这在时间敏感的环境中非常有利。所有这些都展示了训练模型的灵活性，并展示了该方法对不同部署场景的适应性。这些好处至关重要，但同样重要的是要确保任何模型配置的准确性与从头开始训练相应模型大小时获得的准确性相匹配。这引出了下一节。

适用范围。我们选择专注于神经网络，是因为像 ResNets 和 Transformers 这样的模型在涉及视觉和文本任务的现代应用中被广泛使用。尽管如此，值得注意的是，FlexTrain 也可以应用于其他使用与残差连接有相似之处的 boosting 方法的模型，使它们能够获得类似的灵活性优势。

2.2 自动蒸馏

Auto-distillation的目标。只有在与独立训练不同大小模型的情况相比，没有显著的准确性损失时，FlexTrain 提供的灵活性才具有吸引力。为了进一步实现这一目标，并考虑到训练过程中主动层采样的动态性，我们使用了一种知识蒸馏技术的变体，该技术通常用于将知识从教师模型转移到学生模型，如文献【【10】Geoffrey Hinton, et al. Distilling the knowledge in a neural network. 2015.】、【【19】Linfeng Zhang, et al. Be your own teacher: Improve the performance of convolutional neural networks via self distillation. CoRR, 2019.】所述。具体来说，在每个训练周期和每个训练样本 $(x, y) \in D$ 中，我们在公式(2)报告的损失函数中增加一项，该项表示当前配置 $f_{\tilde{W}}(·)$ 的最终激活层与更大配置 $f_{\tilde{W}'}(·)$ 的最终激活层之间的平方 L2 距离，如下所示：

实现细节与算法总结。其中 $\beta$ 是一个固定常数。注意，在通过 $N_k$ 传播损失梯度时，项 $f_{\tilde{W}}(x)$ 不进行微分。理想情况下，人们会从包含最多知识的最大配置 $N$ 中蒸馏知识到较小的配置中。然而，对于大型神经网络来说，这在计算上可能非常昂贵。我们在第 4 节的实验表明，将 $k'$ 固定为 $k + 1$ 已经提供了一个良好的性能基准。综合以上所有内容，我们在算法 1 中总结了整个 FlexTrain 框架。

算法 1: FlexTrain

输入: 模型 $N$，分布 $\pi$，学习率 $\alpha$，数据集 $D$
1. 初始化 $W$;
2. while 未收敛 do
3. 采样配置 $k \sim \pi$;
4. 从 $D$ 中采样大小为 $B$ 的批次 $(x, y)$;
5. $\tilde{W}_k \leftarrow \tilde{W}_k - \alpha \frac{1}{B} \sum_{j=1}^B \nabla_{\tilde{W}_k} \ell_{dist}(x, y; \tilde{W}_k)$;
6. end while

3 联邦 FlexTrain

联邦学习中的异构性挑战。在联邦学习环境中，设备的异构性对大规模深度学习应用构成了重大挑战【【6】Peter Kairouz, et al. Advances and open problems in federated learning. Foundations and Trends® in Machine Learning, 2021.】。由于数据不能共享，这一问题变得更加复杂，导致同时训练多种尺寸的模型变得不太理想，因为能力较差的设备的数据只能用于训练较小的模型，从而导致数据倾斜，并限制了较大模型的丰富性，尤其是在非独立同分布（non-i.i.d.）数据的情况下。为了应对这一挑战，一个能够适应不同训练和推理场景的灵活训练算法对于利用所有设备的数据至关重要。认识到联邦学习环境在广泛应用中的重要性，例如移动键盘预测【【20】Andrew Hard, et al. Federated learning for mobile keyboard prediction. CoRR, 2018.】，我们将 FlexTrain 方法扩展到联邦学习环境，以开发能够提高整体性能的自适应模型。

Federated FlexTrain的核心思想。在联邦学习领域，首要目标是在维护数据隐私的同时，最小化跨多个设备的损失函数。通常，这涉及将损失函数定义为所有设备各自损失函数的加权和，权重为每个设备上的数据样本数量。具体而言，给定 $J$ 个设备中设备 $j$ 上的数据集 $D_j$，神经网络模型 $N \in N^\theta$，以及如公式(2)所述的损失函数 $L(·, ·; W)$，联邦学习环境下的总损失函数由下式给出：

$$\mathcal{L}^{\text{Fed}}(W) = \sum_{j=1}^{J} \frac{|\mathcal{D}_j|}{|\mathcal{D}|} \mathcal{L}_j^{\text{Fed}}(W),$$

其中 $L_{Fed}^j(W) \triangleq \frac{1}{|D_j|} \sum_{(x,y) \in D_j} L(x, y; W)$。由于设备的异构性，Federated FlexTrain 的核心思想是为每个设备提供一个与其资源能力相匹配的模型配置（即完整模型的一部分），使其能够在自己的数据集上执行本地训练周期。然后，模型配置的更新权重被传输到中央服务器，在那里它们与完整模型的其余权重拼接，并与其他设备的更新进行聚合，类似于 FedAvg 算法【【21】H. Brendan McMahan, et al. Federated learning of deep networks using model averaging. CoRR, 2016.】。这种方法使得集中式的 FlexTrain 算法能够以分布式的方式实现，同时在训练和推理过程中考虑到所有设备的资源限制。沿用第 2 节中描述的符号，并让 $\pi = [\pi_1, \dots, \pi_K]$ 表示一个概率分布，其中 $\pi_k$ 代表能够承担部署神经网络 $N$ 的前 $k$ 层的设备比例，我们可以得出结论，Federated FlexTrain 最小化了由设备提供的不同模型配置分布下的预期损失，详述如下：

$$\overline{\mathcal{L}}^{\text{Fed}}(W) = \sum_{j=1}^{J} \frac{|\mathcal{D}_j|}{|\mathcal{D}|} \mathbb{E}_{k_j \sim \boldsymbol{\pi}} [\mathcal{L}_j^{\text{Fed}}(\tilde{W}_{k_j})].$$

联邦场景下的Auto-distillation。将第 2.2 节中介绍的蒸馏概念整合到联邦学习案例中是可取的，但也具有挑战性。运行不同配置的设备无法访问共享数据集，这使得应用相同的蒸馏技术变得困难。为了解决这个问题，我们提出了一个变通方法，即每个设备（除了最弱的设备）使用其本地数据集对前一个配置进行蒸馏。我们通过在公式(2)报告的损失函数中增加一个项来实现这一点，该项表示设备可承担的当前配置 $f_{\tilde{W}}(\cdot)$ 的最终激活层与小一步的配置 $f_{\tilde{W}}(\cdot)$ 的最终激活层之间的平方 L2 距离，如下所示：

$$\ell_{\text{dist}}^{\text{Fed}}(x, y; \tilde{W}_k) = \ell(x, y; \tilde{W}_k) + \beta \| f_{\tilde{W}_k}(x) - f_{\tilde{W}_{k-1}}(x) \|_2^2,$$

其中 $\beta$ 是一个固定常数。请注意，在通过 $N_k$ 传播损失梯度时，项 $f_{\tilde{W}}(x)$ 不进行微分。通过增加这个项，知识从较大的配置蒸馏到较小的配置，使得较小的模型能够利用更强大设备的高级能力。综合以上所有内容，我们在算法 2 中总结了 FedFlexTrain 框架的全部内容。

算法 2: Federated FlexTrain

输入: 神经网络 $N$，学习率 $\alpha$，本地数据集 $D_j$，本地周期 $l$

1. 服务器端:
2. 初始化 $W$;
3. while 未收敛 do
4. 采样设备 $C \subseteq \{1, \dots, J\}$;
5. 向设备 $j \in C$ 询问其模型配置 $k_j$ 并发送 $\tilde{W}_{k_j}$;
6. 接收更新后的 $\tilde{W}'_{k_j}$;
7. 聚合模型 $W \leftarrow \frac{1}{|C|} \sum_{j \in C} W'_{k_j}$
8. end while
9. 设备 j 端:
10. 向主服务器请求模型配置 $k_j$;
11. 从主服务器接收 $\tilde{W}_k$;
12. counter $\leftarrow$ 0;
13. while counter < $l$ do
14. foreach 批次 $(x, y) \sim D_j$ (大小为 B) do
15. $\tilde{W}_{k_j} \leftarrow \tilde{W}_{k_j} - \alpha \frac{1}{B} \sum_{b=1}^B \nabla_{\tilde{W}_k} \ell_{dist}^{Fed}(x^{(b)}, y^{(b)}; \tilde{W}_k)$
16. end foreach
17. counter $\leftarrow$ counter + 1;
18. end while
19. 将更新后的 $\tilde{W}_k$ 发送给主服务器

A4 实验环境

• 数据集:

  • CIFAR-100: 用于集中式和联邦学习设置的实验。
  • CIFAR-10: 用于联邦学习设置的实验。

• 模型架构:

  • ResNet: ResNet-56, ResNet-20, ResNet-8。
  • Visual Transformers (ViT): ViT-3, ViT-2, ViT-1 (ViT-k 表示包含 k 个 Transformer 块的模型)。

• 硬件配置:

  • GPU: 所有实验均在 Tesla-V100 GPU 上进行。

• 软件与实验配置:

  • 代码实现: 实验结果是三次不同随机种子运行的平均值。

  • 集中式设置:

    • ResNet: 训练 160 个 epoch，批大小 64，使用带动量（0.9）的 SGD。正则化采用随机深度（dropout 概率 0.5）和权重衰减（0.001），学习率为 0.05。
    • ViT: 训练 300 个 epoch，使用 dropout（概率 0.1），学习率为 0.003。
    • 蒸馏参数: $\beta$ 设为 0.2。

  • 联邦设置:

    • 通信: 160 轮全局通信。
    • 本地训练: 每轮通信包含 10 个本地 epoch。本地 epoch 的设置与集中式设置相同。
    • 设备: 共 20 个设备，数据集在设备间平均分配。设备能力根据其能容纳的模型大小划分（见原文表 3）。

A4 实验结果

1. 集中式设置实验

• 实验内容:

  • 在 CIFAR-100 数据集上，将 FlexTrain 与两种基准方法进行比较：

    1. Independent: 独立训练不同大小的模型（如 ResNet-56, ResNet-20, ResNet-8）。
    2. Single: 只训练一个最大的模型（如 ResNet-56），然后直接截取其前面的层来部署到能力较弱的设备上。

  • 分别在 ResNet 和 ViT 架构上进行了评估。

• 实验结果:

  • ResNet (表1):

    • Single 方法在小模型上的性能非常差（ResNet-20 等效准确率 6.5%，ResNet-8 等效准确率 2.6%）。
    • FlexTrain 的性能略低于 Independent（例如，ResNet-56 上为 68.6% vs 69.9%），但远超 Single。
    • 效率: FlexTrain 只需一次训练，其计算量（FLOPs）仅为 Independent 方法（需三次训练）的 41%，以及 Single 方法的 62%。

  • ViT (表2):

    • 结果与 ResNet 类似。Single 方法性能很差。FlexTrain 性能略低于 Independent，但计算量仅为后者的 37%。

• 分析结论:

  • 在异构环境中部署单一训练的最大模型是不可行的。
  • FlexTrain 在准确率和训练资源（时间、能耗）之间取得了很好的平衡。虽然准确率略有下降，但它只需一次训练过程，显著节省了计算资源。

2. 联邦设置实验

• 实验内容:

  • 在 CIFAR-10 和 CIFAR-100 数据集上，将 Federated FlexTrain 与两种联邦学习基准进行比较：
    1. FedSmall: 所有设备共同训练一个最弱设备能承受的最小模型。
    2. FedClass: 将设备按能力分组，每组使用 FedAvg 独立训练一个大小合适的模型。

• 实验结果 (表4 和 图2):

  • Federated FlexTrain 在 CIFAR-10 和 CIFAR-100 上的平均设备准确率均显著优于 FedSmall 和 FedClass。

    • 在 CIFAR-10 上，Federated FlexTrain (83%) vs FedSmall (81%) vs FedClass (78.9%)。
    • 在 CIFAR-100 上，Federated FlexTrain (48.5%) vs FedSmall (42.5%) vs FedClass (37.6%)。

  • 图2 显示，FedSmall 限制了强设备性能，而 FedClass 由于数据隔离导致性能瓶颈。Federated FlexTrain 使得所有设备都能发挥其最大潜力。

• 分析结论:

  • FedSmall 方案对能力强的设备不公平，限制了它们的性能。
  • FedClass 方案存在严重的数据共享问题，能力强的设备无法从能力弱的设备的数据中学习，在非独立同分布场景下问题更严重。

  • Federated FlexTrain 解决了这些问题：

    • 它通过训练一个通用的全局模型，使得弱设备的数据可以用于训练模型的初始层，从而提升大模型的性能。
    • 提出的蒸馏技术让小模型能够利用强设备学到的高级知识，突破自身局限。

  • 这证明了 Federated FlexTrain 是一种多功能方法，能够有效地将弱设备纳入全局模型训练，从而提升模型整体准确性，并为设备提供满足其约束的模型。

A5 结论

总结：
本文提出了 FlexTrain，一个为应对大型机器学习模型在异构设备环境中所带来挑战的动态训练框架。实验表明，在集中式设置中，FlexTrain 在准确性和能耗之间取得了良好的平衡，避免了训练多种尺寸模型的需要，同时保持了较高的准确性。同样，在联邦学习场景下，Federated FlexTrain 在准确性方面优于让不同能力等级的设备各自训练一个通用模型的方法。

局限性：
FlexTrain 采用从深到浅的层失活策略，这有利于在训练中学习基础特征，并让深层网络在此基础上构建更复杂的特征。这种方法特别适用于像 ResNets 和 Transformers 这样的残差神经架构。虽然这些架构是当前最先进且广泛应用的技术，但在某些应用场景中可能更适合其他架构，这可能会限制 FlexTrain 在这些情况下的适用性。

展望：
随着机器学习模型规模的持续增长，相关挑战将变得更加突出。鉴于机器学习在各种应用中的广泛采用，解决这些挑战至关重要，因为计算和存储能力有限的弱设备很可能会参与其中。一个有前景的解决方案是结合文献中提出的各种技术，包括压缩、量化、稀疏化等。FlexTrain 提供了一个额外的灵活性层次，可以与这些技术相结合，以推动模型训练和部署效率的边界。我们相信，这种集成方法对于应对大型机器学习模型的挑战并支持其在异构设备环境中的部署至关重要。