MoE$^2$: Optimizing Collaborative Inference for Edge Large Language Models

📄 arXiv: 2501.09410v1 📥 PDF

作者: Lyudong Jin, Yanning Zhang, Yanhan Li, Shurong Wang, Howard H. Yang, Jian Wu, Meng Zhang

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

发布日期: 2025-01-16

备注: Submitted to IEEE/ACM Transactions on Networking

💡 一句话要点

提出MoE$^2$框架，优化边缘大语言模型协同推理，提升能效与降低延迟。

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

关键词: 边缘计算 大语言模型 协同推理 专家混合 模型优化

📋 核心要点

  1. 现有边缘LLM推理方法难以有效利用异构边缘设备的计算能力，导致成本高、延迟大。
  2. MoE$^2$框架通过两级专家选择机制和离散单调优化算法，优化边缘LLM的协同推理过程。
  3. 实验表明，MoE$^2$在多种LLM模型上均有性能提升，并在不同资源约束下优于现有方法。

📝 摘要（中文）

本文提出了一种名为“边缘专家混合(MoE$^2$)”的边缘大语言模型协同推理框架。该框架旨在解决边缘LLM异构能力利用问题，以实现更高的成本效益和更低的延迟。MoE$^2$将联合门控和专家选择问题建模为在能量和延迟约束下优化推理性能。与传统MoE问题不同，LLM专家选择更具挑战性，因为边缘LLM具有组合性质和异构属性。为此，本文提出了一种两级专家选择机制，揭示了门控参数在专家选择中保持最优性的特性，从而将训练和选择过程解耦，显著降低了复杂度。此外，利用目标函数的单调性，设计了一种离散单调优化算法以实现最优专家选择。实验结果表明，MoE$^2$方法能够在不同的延迟和能量预算之间实现最佳权衡，并在各种系统资源约束下优于基线方法。

🔬 方法详解

问题定义：论文旨在解决如何在边缘设备上高效协同推理大型语言模型（LLM）的问题。现有的方法难以充分利用边缘设备异构的计算能力，导致推理延迟高、能耗大，无法满足新兴应用的需求。传统的MoE方法在边缘LLM场景下面临组合爆炸和异构性挑战，专家选择难度大。

核心思路：论文的核心思路是将边缘LLM的协同推理问题建模为联合门控和专家选择的优化问题，目标是在给定的能量和延迟约束下最大化推理性能。通过解耦训练和选择过程，并利用目标函数的单调性，降低了问题复杂度，从而实现高效的专家选择。

技术框架：MoE$^2$框架包含两个主要阶段：训练阶段和推理阶段。在训练阶段，各个边缘LLM作为专家独立训练。在推理阶段，首先进行第一级专家选择，根据门控网络输出的权重选择候选专家集合；然后，利用离散单调优化算法在候选集合中选择最优的专家组合进行协同推理。框架的关键在于两级专家选择机制和离散单调优化算法。

关键创新：论文的关键创新在于提出了两级专家选择机制，并证明了门控参数在专家选择中具有最优性保持特性。这一特性使得训练和选择过程可以解耦，极大地降低了计算复杂度。此外，针对专家选择问题，设计了一种离散单调优化算法，能够在保证最优性的前提下高效地找到最优的专家组合。

关键设计：两级专家选择机制中，第一级使用门控网络预测每个专家的权重，选择权重较高的专家作为候选集合。第二级使用离散单调优化算法，根据延迟和能量约束，从候选集合中选择最优的专家组合。离散单调优化算法利用了目标函数的单调性，避免了对所有可能的专家组合进行评估，从而提高了选择效率。具体而言，算法通过迭代地添加或删除专家，并评估目标函数的变化，最终找到满足约束条件的最优解。

📊 实验亮点

实验结果表明，MoE$^2$框架在多种LLM模型上均取得了显著的性能提升。例如，在NVIDIA Jetson AGX Orin和RTX 4090 GPU上，MoE$^2$能够在不同的延迟和能量预算下实现最佳的性能权衡，并且在各种系统资源约束下优于基线方法。具体而言，MoE$^2$在延迟降低的同时，能够显著降低能耗，提升推理效率。

🎯 应用场景

MoE$^2$框架可应用于各种需要低延迟、高能效的边缘LLM推理场景，例如智能助手、自动驾驶、智能制造等。通过优化边缘设备的协同推理能力，可以实现更快的响应速度、更低的功耗，并支持更复杂的应用。

📄 摘要（原文）

Large language models (LLMs) have demonstrated remarkable capabilities across a wide range of natural language processing tasks. Exploiting the heterogeneous capabilities of edge LLMs is crucial for diverse emerging applications, as it enables greater cost-effectiveness and reduced latency. In this work, we introduce \textit{Mixture-of-Edge-Experts (MoE$^2$)}, a novel collaborative inference framework for edge LLMs. We formulate the joint gating and expert selection problem to optimize inference performance under energy and latency constraints. Unlike conventional MoE problems, LLM expert selection is significantly more challenging due to the combinatorial nature and the heterogeneity of edge LLMs across various attributes. To this end, we propose a two-level expert selection mechanism through which we uncover an optimality-preserving property of gating parameters across expert selections. This property enables the decomposition of the training and selection processes, significantly reducing complexity. Furthermore, we leverage the objective's monotonicity and design a discrete monotonic optimization algorithm for optimal expert selection. We implement edge servers with NVIDIA Jetson AGX Orins and NVIDIA RTX 4090 GPUs, and perform extensive experiments. Our results validate that performance improvements of various LLM models and show that our MoE$^2$ method can achieve optimal trade-offs among different delay and energy budgets, and outperforms baselines under various system resource constraints.