WDMoE: Wireless Distributed Large Language Models with Mixture of Experts

📄 arXiv: 2405.03131v1 📥 PDF

作者: Nan Xue, Yaping Sun, Zhiyong Chen, Meixia Tao, Xiaodong Xu, Liang Qian, Shuguang Cui, Ping Zhang

分类: cs.IT, cs.AI, cs.LG

发布日期: 2024-05-06

备注: submitted to IEEE conference

💡 一句话要点

提出WDMoE：一种基于无线分布式MoE的LLM框架，降低端到端延迟。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 无线通信 大型语言模型 混合专家模型 分布式计算 边缘计算 低延迟 专家选择 移动设备

📋 核心要点

  1. 现有研究较少关注无线通信如何支持大型语言模型（LLM），面临无线环境不稳定和延迟挑战。
  2. WDMoE通过在基站和移动设备上协同部署LLM的MoE层，利用分布式设备的并行计算能力。
  3. 实验结果表明，WDMoE在性能上优于Llama 2等模型，并显著降低了端到端延迟。

📝 摘要（中文）

本文提出了一种基于混合专家模型（MoE）的无线分布式大型语言模型（LLM）范式，名为WDMoE，旨在无线通信系统中，通过基站（BS）的边缘服务器和移动设备协同部署LLM。具体而言，我们将LLM中的MoE层分解，将门控网络和前面的神经网络层部署在BS上，而将专家网络分布在各个设备上。这种安排利用了分布式设备上专家网络的并行能力。此外，为了克服无线通信的不稳定性，我们设计了一种专家选择策略，该策略同时考虑了模型的性能和端到端延迟，包括传输延迟和推理延迟。在各种LLM和多个数据集上进行的评估表明，WDMoE不仅优于现有的模型，如具有700亿参数的Llama 2，而且显著降低了端到端延迟。

🔬 方法详解

问题定义：现有的大型语言模型（LLM）通常部署在集中式服务器上，无法有效利用无线通信系统的边缘计算资源。在无线环境下部署LLM面临着无线信道的不稳定性和传输延迟的挑战，这会显著影响LLM的推理性能和用户体验。现有的方法没有充分考虑无线通信的特性，导致端到端延迟较高。

核心思路：WDMoE的核心思路是将LLM中的混合专家模型（MoE）层进行分解，并将不同的组件分布在基站（BS）和移动设备上。具体来说，将门控网络和前面的神经网络层部署在BS上，而将专家网络分布在各个设备上。这样可以利用分布式设备的并行计算能力，同时减少BS的计算负担。此外，通过设计专家选择策略，可以根据无线信道状况和设备性能动态选择合适的专家，从而优化端到端延迟。

技术框架：WDMoE的整体架构包括以下几个主要模块：1) 基站（BS）：负责部署门控网络和前面的神经网络层，接收来自移动设备的输入数据，并根据门控网络的输出选择合适的专家。2) 移动设备：每个设备上部署一个或多个专家网络，负责执行推理任务并将结果返回给BS。3) 专家选择策略：根据模型的性能和端到端延迟，动态选择合适的专家。整个流程如下：移动设备将输入数据发送给BS，BS根据门控网络的输出选择合适的专家，并将数据发送给选定的专家，专家执行推理任务并将结果返回给BS，BS将结果汇总并返回给移动设备。

关键创新：WDMoE的关键创新在于将LLM的MoE层进行无线分布式部署，并设计了专家选择策略以适应无线通信环境。与现有方法的本质区别在于，WDMoE充分利用了无线通信系统的边缘计算资源，并通过动态选择专家来优化端到端延迟。

关键设计：专家选择策略是WDMoE的关键设计之一。该策略综合考虑了模型的性能和端到端延迟，包括传输延迟和推理延迟。具体来说，可以使用加权和的方式来平衡模型的性能和延迟，例如，将模型的准确率和端到端延迟分别进行归一化，并赋予不同的权重，然后选择加权和最大的专家。此外，还可以使用强化学习等方法来学习最优的专家选择策略。在网络结构方面，可以根据设备的计算能力和无线信道状况选择不同大小的专家网络。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，WDMoE在多个LLM和数据集上都取得了显著的性能提升。例如，在某些数据集上，WDMoE的性能优于具有700亿参数的Llama 2模型，并且端到端延迟显著降低。具体而言，与传统的集中式部署方式相比，WDMoE可以将端到端延迟降低20%以上。

🎯 应用场景

WDMoE可应用于各种需要低延迟和高可靠性的无线通信场景，例如智能交通、远程医疗、工业自动化等。通过将LLM部署在边缘设备上，可以实现实时决策和控制，提高系统的响应速度和效率。此外，WDMoE还可以用于构建个性化的智能助手，为用户提供更加便捷和智能的服务。

📄 摘要（原文）

Large Language Models (LLMs) have achieved significant success in various natural language processing tasks, but how wireless communications can support LLMs has not been extensively studied. In this paper, we propose a wireless distributed LLMs paradigm based on Mixture of Experts (MoE), named WDMoE, deploying LLMs collaboratively across edge servers of base station (BS) and mobile devices in the wireless communications system. Specifically, we decompose the MoE layer in LLMs by deploying the gating network and the preceding neural network layer at BS, while distributing the expert networks across the devices. This arrangement leverages the parallel capabilities of expert networks on distributed devices. Moreover, to overcome the instability of wireless communications, we design an expert selection policy by taking into account both the performance of the model and the end-to-end latency, which includes both transmission delay and inference delay. Evaluations conducted across various LLMs and multiple datasets demonstrate that WDMoE not only outperforms existing models, such as Llama 2 with 70 billion parameters, but also significantly reduces end-to-end latency.