Table of Contents
Fetching ...

L-MoE: End-to-End Training of a Lightweight Mixture of Low-Rank Adaptation Experts

Shihao Ji, Zihui Song

TL;DR

L-MoE introduces a fully differentiable framework that unifies Mixture of Experts with Low-Rank Adaptation by defining each MoE expert as a LoRA adapter and learning a lightweight gating network to softly compose these adapters per token. The core mechanism computes a soft mixture of LoRA updates $\Delta W^{(l)}(z) = \sum_i p_i(z) B_i^{(l)} A_i^{(l)}$ based on token context, and applies this composite update atop the frozen backbone with $W_0^{(l)} + \alpha \Delta W^{(l)}(z)$. The model is trained end-to-end using a joint objective combining an autoregressive loss $\mathcal{L}_{\text{AR}}$ and a load-balancing loss $\mathcal{L}_{\text{LB}}$, $\\mathcal{L}_{\text{Total}} = \mathcal{L}_{\text{AR}} + \lambda \mathcal{L}_{\text{LB}}$, to promote balanced expert utilization. This architecture offers substantial parameter efficiency, modular interpretability, and dynamic skill composition, enabling scalable specialization of LLMs with end-to-end trainability and potential for continual extension through additional LoRA experts or routing refinements.

Abstract

The Mixture of Experts (MoE) architecture enables the scaling of Large Language Models (LLMs) to trillions of parameters by activating a sparse subset of weights for each input, maintaining constant computational cost during inference. Concurrently, Low-Rank Adaptation (LoRA) has emerged as a dominant technique for parameter-efficiently fine-tuning LLMs on specialized tasks. In this work, we unify these two paradigms into a novel, end-to-end trainable framework named L-MoE: a Lightweight Mixture of LoRA Experts. L-MoE redefines MoE experts not as dense feed-forward networks, but as a collection of task-specialized, low-rank adapters. A lightweight gating network, trained jointly with the experts, learns to dynamically compose these LoRA adapters by computing a weighted average of their parameters for each input token. This composition is fully differentiable, allowing gradients from a standard auto-regressive language modeling objective to flow back through the entire architecture, simultaneously refining both the expert adapters and the routing strategy. This approach creates a highly parameter-efficient MoE model that is modular by design, allows for dynamic skill composition, and is trainable from end-to-end. We present the formal mathematical framework for L-MoE, detailing the differentiable routing mechanism and the joint optimization objective, thereby providing a new path toward building more efficient, scalable, and specialized language models.

L-MoE: End-to-End Training of a Lightweight Mixture of Low-Rank Adaptation Experts

TL;DR

L-MoE introduces a fully differentiable framework that unifies Mixture of Experts with Low-Rank Adaptation by defining each MoE expert as a LoRA adapter and learning a lightweight gating network to softly compose these adapters per token. The core mechanism computes a soft mixture of LoRA updates based on token context, and applies this composite update atop the frozen backbone with . The model is trained end-to-end using a joint objective combining an autoregressive loss and a load-balancing loss , , to promote balanced expert utilization. This architecture offers substantial parameter efficiency, modular interpretability, and dynamic skill composition, enabling scalable specialization of LLMs with end-to-end trainability and potential for continual extension through additional LoRA experts or routing refinements.

Abstract

The Mixture of Experts (MoE) architecture enables the scaling of Large Language Models (LLMs) to trillions of parameters by activating a sparse subset of weights for each input, maintaining constant computational cost during inference. Concurrently, Low-Rank Adaptation (LoRA) has emerged as a dominant technique for parameter-efficiently fine-tuning LLMs on specialized tasks. In this work, we unify these two paradigms into a novel, end-to-end trainable framework named L-MoE: a Lightweight Mixture of LoRA Experts. L-MoE redefines MoE experts not as dense feed-forward networks, but as a collection of task-specialized, low-rank adapters. A lightweight gating network, trained jointly with the experts, learns to dynamically compose these LoRA adapters by computing a weighted average of their parameters for each input token. This composition is fully differentiable, allowing gradients from a standard auto-regressive language modeling objective to flow back through the entire architecture, simultaneously refining both the expert adapters and the routing strategy. This approach creates a highly parameter-efficient MoE model that is modular by design, allows for dynamic skill composition, and is trainable from end-to-end. We present the formal mathematical framework for L-MoE, detailing the differentiable routing mechanism and the joint optimization objective, thereby providing a new path toward building more efficient, scalable, and specialized language models.
Paper Structure (14 sections, 10 equations)