专栏名称: 计算机视觉深度学习和自动驾驶

讨论计算机视觉、深度学习和自动驾驶的技术发展和挑战

ICML 2024 || FLEXTRON: 嵌套弹性结构的MOE

计算机视觉深度学习和自动驾驶 · 公众号 · · 2024-07-22 00:51

正文

1. 基本信息和摘要

论文题目

FLEXTRON: Many-in-One Flexible Large Language Model

Arxiv: https://arxiv.org/pdf/2406.10260

作者

Ruisi Cai, Saurav Muralidharan, Greg Heinrich, Hongxu Yin, Zhangyang Wang, Jan Kautz, Pavlo Molchanov

研究单位

NVIDIA
The University of Texas at Austin

解决问题

Mixture-of-Expert (MoE) 网络是一种由多个专家（experts）组成的系统，每个专家负责处理输入数据的不同部分或不同种类的任务。与传统的密集（dense）模型相比，MoE网络在某些情况下更加高效，因为它们可以并行地利用多个专家来处理数据，从而提高计算效率和模型的扩展性。然而，MoE网络 通常受限于仅在前馈层（feedforward layers）中使用 ，并且 专家的数量和大小通常是固定的 ，这限制了模型的灵活性和对不同任务的适应性。

作者提出自适应弹性网络（ 拓展到attention 和feedforward层），优点有：

不同大小的子网络 ：根据不同的输入数据和任务需求，动态选择不同大小的子网络，以实现性能和效率的最大化。
输入驱动的适应性 ：模型能够根据输入数据的特性自动调整其结构和行为，从而更好地适应不同的任务。

通过这种方式，FLEXTRON模型能够灵活地适应不同的延迟和准确性目标，同时保持输入数据的自适应性，这使得它 在资源受限的环境中具有很大的潜力。

摘要

本文提出了FLEXTRON，一种新型的网络架构和模型优化框架，支持大型语言模型（LLMs）的灵活部署。FLEXTRON利用嵌套弹性结构， 在推理过程中无需额外微调即可快速适应特定的用户定义的延迟和准确性目标 。它还具有 输入自适应性 ，能够自动将令牌路由（route token）通过其子网络以提高性能和效率。作者提出了一种高效的样本训练方法和相关的路由算法，系统地将现有的训练有素的LLM转换为FLEXTRON模型。在GPT-3和LLama-2系列LLMs上评估FLEXTRON，展示了其在多个端到端训练变体和其他最先进的弹性网络中的优越性能，所有这些都仅通过一次预训练运行实现，仅消耗原始预训练7.63%的令牌。

2. 介绍

主要贡献

Multiple Models in One , 主要贡献包括：

提出了一种新颖的架构FLEXTRON，它在推理期间能够灵活适应不同的延迟和准确性目标， 无需额外微调 。
提出了一种后训练优化框架，用于系统地将现有的训练有素的LLMs转换为动态（输入自适应）弹性网络。
提出了新的静态和动态路由算法，自动选择给定延迟目标和/或输入令牌的最优子网络。引入了一种新颖的替代模型，用于有效训练路由。
提出了一种高效的基于采样的训练方法，用于弹性网络，其计算需求远低于现有方法。

3. 方法

Step 1 & Step2

Elastic Network Continued-Training(弹性网络训练)

弹性网络训练是FLEXTRON模型的关键步骤之一，旨在将一个已经训练好的大型语言模型（LLM）转换为具有弹性的网络，使其能够在不同的延迟和准确性目标下灵活地进行推理。

重要性排序 ：

使用一小部分数据样本，基于激活的累积幅度计算每个神经元/头的重要性。
对于多头注意力（MHA）层，每个头的重要性计算公式为：
对于多层感知器（MLP）层：其中，表示权重矩阵的第行。

权重矩阵重排列 ：

根据计算出的重要性，对MLP和MHA层中的权重矩阵进行排序，使得每个层的神经元/头按照重要性顺序存储。

构建弹性层 ：

通过简单地索引每层的前几个神经元/头，构建具有参数共享的嵌套弹性层，从而形成多个具有不同运行时和准确性特征的子网络。

同时训练所有候选网络 ：

使用组合损失项同时训练所有候选的弹性网络：
其中，是随机模型的索引，是候选操作的数量，是损失函数。

随机采样（Random Sampling） ：

由于候选网络的数量可能非常大，通过随机采样一小部分网络进行训练，以保持总预训练时间的可行性。

弹性网络继续训练的目标是保持模型核心知识的同时，通过结构化地调整网络的容量，使其能够适应不同的运行时和准确性要求。这一过程不仅提高了模型的灵活性，还为后续的自动网络选择和路由算法的训练奠定了基础。通过这种方法，FLEXTRON模型能够在不同的部署环境中实现最优的性能和效率。

Step 3

Automatic Network Selection（自动网络选择）

这一过程涉及 两个主要的路由场景 ：静态路由 (Statstic) 和动态路由(Dynamic Model Selection)。

静态模型选择

在静态模型选择中，目标是自动选择给定目标延迟的子网络，而不考虑输入数据的特定特征。这是通过在每层中插入可学习的路由器来实现的，路由器根据延迟要求选择最合适的专家组。

路由器输出选择为：
其中，是一个小的多层感知器（MLP），它将延迟作为输入并输出路由器的 logits。

动态模型选择

动态模型选择进一步扩展了静态选择的概念，通过考虑当前的隐藏状态来实现输入自适应性。这意味着模型可以根据输入数据的难度或特性选择不同的子网络。

动态路由器的输出选择为：
其中，计算为：
是用于投影延迟的权重矩阵，是用于投影隐藏状态的权重矩阵，是路由器的输出权重。

替代模型（Surrogate Model, SM）

为了有效地训练路由器，尤其是在存在梯度消失问题的情况下，作者提出了使用一个替代模型来预测基于路由器输出的LLM的语言损失值。

替代模型