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

TL;DR

This work proposes a two-level expert selection mechanism through which an optimality-preserving property of gating parameters across expert selections is uncovered, which enables the decomposition of the training and selection processes, significantly reducing complexity.

Abstract

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.

Figures (6)

• Figure 1: An overview of the MoE$^2$ framework, which leverages a gating network and a two-level expert selection mechanism for efficient task processing. Mobile users send prompts, which are routed to a subset of edge servers with LLM experts. The expert selection method is carefully designed to enable a scalable and efficient framework for diverse LLM applications.
• Figure 2: A detailed illustration of the MoE$^2$ framework, which consists of training stage and inference stage. In the training stage, the gating network parameters $\bm{\theta}$ are derived by optimizing $\mathcal{L}(\boldsymbol \theta, \mathcal{S})$ over all LLM $\mathcal{N}$. Then, the subset selection $\mathcal{S}$ are obtained by Subset Monotonic Optimization algorithm to satisfy system constraints. In the inference stage, the gating network computes the gating values for each LLM experts, and Top-$k$ mechanism selects a smaller subset from $\mathcal{S}$ of $k$ LLMs with the highest gating values for responding to prompt $\mathbf{x}$.
• Figure 3: Performance comparison on the MMLU dataset. The MoE$^2$ architecture achieves the highest accuracy, which outperforms the Single Agent, Majority Voting and A. E. A. by $14.4\%$, $7.3\%$, $16.4\%$, respectively.
• Figure 4: Implementation of an edge LLM testbed on NVIDIA Jetson AGX Orin 64GB, connected to NVIDIA GeForce RTX 4090 via a local area network.
• Figure 5: Studies of edge delays and edge energy consumptions under different prompt lengths. Edges equipped with less powerful computational resources and larger-scale LLM models exhibit higher delays, while edges equipped with more powerful computational resources and larger-scale LLM models exhibit higher energy consumption.

Theorems & Definitions (5)

• Remark 1
• Theorem 1: Optimality for subset
• Remark 2
• Theorem 2: Monotonic Improvement
• proof