Mixture of Lookup Experts

Mixture of Lookup Experts

Speakers: Shibo Jie, Peking University; Yehui Tang; Kai Han; Yitong Li; Duyu Tang; Zhi-Hong Deng; Yunhe Wang

Conference: ICML 2025

YouTube: https://slideslive.com/39043992

Overview

This presentation introduces Mixture of Lookup Experts (MoLA), a novel architectural design aimed at making large language models (LLMs) more friendly for deployment on edge devices such as mobile phones and personal computers. The talk, delivered on behalf of primary author Shibo Jie from Peking University, addresses the significant challenges faced by conventional Mixture of Experts (MoE) models in VRAM- and computation-constrained edge scenarios. MoE models, while efficient in reducing computation by activating only a subset of experts per token, suffer from large parameter sizes that often exceed the memory capacity of single-GPU edge devices.

MoLA proposes a radical rethinking of the expert paradigm, moving away from computation-intensive Feed-Forward Networks (FFNs) to lookup experts that directly retrieve results from a precomputed key-value mapping. This innovative approach significantly reduces the VRAM footprint and loading overhead during inference, which are critical bottlenecks for MoE deployment on the edge. By strategically converting computation experts into lookup tables based on embedding tokens, MoLA achieves inference speeds comparable to dense models while maintaining MoE-level performance, albeit with a trade-off in storage footprint. The work offers a compelling solution for expanding the accessibility and deployment of advanced LLM architectures to a broader range of resource-limited environments.

Background

▶ Watch: Challenge of MoE deployment on edge devices (0:00)

The evolution of large language models has seen Mixture of Experts (MoE) architectures emerge as a powerful paradigm for scaling model capacity while controlling computational cost. MoE models achieve this by sparsely activating a small subset of specialized "experts" (typically FFNs) for each input token, rather than processing the token through all experts. This conditional computation allows for models with billions or even trillions of parameters to be trained and inferred more efficiently than their dense counterparts, by effectively increasing model capacity without a proportional increase in FLOPs.

However, the deployment of these sophisticated MoE models on edge devices presents a formidable challenge. Edge scenarios, characterized by mobile phones, personal computers, or embedded systems, typically feature a single GPU with severely limited VRAM (Video Random Access Memory) and constrained storage-to-VRAM bandwidth. The sheer number of parameters in MoE models, even with sparse activation, often exceeds the available VRAM, making it impossible to load the entire model into memory.

A common strategy to address VRAM limitations in such scenarios is expert offloading. This technique involves storing the majority of expert parameters in slower secondary storage (e.g., SSD) and loading them into VRAM only when they are required during inference. While expert offloading alleviates the VRAM constraint, it introduces a new bottleneck: the frequent and large data transfers between secondary storage and VRAM. On edge devices, this storage-to-VRAM bandwidth is often limited, leading to significant latency and hindering inference speed, especially as different experts need to be loaded for different tokens or batch items. The core problem, therefore, is to reduce the VRAM usage of MoE models without incurring prohibitively high loading overheads that negate the benefits of offloading. The fundamental question guiding this research was: "Why must experts be placed in or loaded into VRAM?" The answer lies in their traditional implementation as FFNs, which inherently demand intensive matrix computations best performed on GPUs. This realization sparked the idea that if experts could be designed to require no intensive computation, they would no longer necessitate GPU involvement or VRAM residency.

Key Findings

▶ Watch: MoLA: Converting experts to lookup tables via embedding tokens (2:00)

The Mixture of Lookup Experts (MoLA) architecture presents several pivotal findings that challenge conventional MoE design and offer a new path for edge deployment:

1. Redefinition of Experts: The central innovation is the replacement of traditional computation-intensive FFN experts with lookup experts. These lookup experts directly retrieve precomputed results from a key-value mapping based on input, eliminating the need for GPU-intensive matrix computations at inference time.
2. Leveraging Embedding Tokens: MoLA identifies LLM embedding tokens (the output of the embedding layer) as ideal inputs for lookup experts. Since embedding tokens are finite (equal to the vocabulary size), their corresponding expert outputs can be precomputed for every possible token ID, enabling the construction of an efficient lookup table.
3. Significant Reduction in Loading Overhead: During inference, MoLA only needs to load compact output vectors from the lookup table into VRAM, in contrast to MoE which must load large expert parameter matrices. This results in a reduction of loaded parameters by "several orders of magnitude," drastically mitigating the storage-to-VRAM bandwidth bottleneck on edge devices.
4. Dense Model Inference Efficiency: Despite its MoE heritage, MoLA achieves inference speeds comparable to that of a dense model. The overhead introduced by its lookup tables and lightweight routers is negligible, and its performance does not degrade with increasing batch sizes in the way offloaded MoE models do due to varied expert loading.
5. Comparable Performance to MoE: Under controlled conditions—specifically, matching inference-time VRAM usage, computation, and training-time parameter count—MoLA demonstrates performance levels comparable to conventional MoE models. This indicates that the architectural shift does not compromise model quality.
6. Higher Storage, Lower Latency Trade-off: MoLA's design entails a larger overall storage footprint for the lookup table, which scales with vocabulary size. However, this increased storage is a deliberate trade-off for significantly reduced loading overhead and superior inference latency, a crucial advantage in bandwidth-constrained environments.
7. Simplified Training with Dense Activation: During training, MoLA activates all routed experts densely. This approach, while increasing training cost (a focus on edge inference efficiency), eliminates the need for auxiliary loss functions typically used to balance expert utilization in sparse MoE training, simplifying the training process.
8. Scaling Flexibility: The research explored two scaling strategies. While increasing the size of routed experts yielded limited returns, increasing the number of experts proved more effective. This latter strategy allows for continued performance improvements under fixed compute/VRAM constraints by judiciously trading off additional storage and loading costs.
9. Reparameterization Enables Performance Recovery: An ablation study revealed that directly changing expert inputs to embedding tokens initially causes a performance drop due to lost contextual information. However, the subsequent reparameterization into lookup form, which allows for dense activation of all experts at inference-time without additional cost, effectively recovers and surpasses this initial performance deficit.

Technical Deep Dive

▶ Watch: Efficiency analysis: reduced loading overhead, VRAM savings (4:40)

The core innovation of Mixture of Lookup Experts (MoLA) lies in its fundamental re-engineering of the expert mechanism within an MoE framework to circumvent the VRAM and bandwidth limitations of edge devices. This involves distinct modifications during training, a critical precomputation step, and a streamlined inference process.

At its heart, MoLA challenges the assumption that experts must be computation-intensive FFNs requiring GPU acceleration. Instead, it proposes lookup experts, which are essentially precomputed key-value mappings. The feasibility of this transformation hinges on the nature of the expert's input. If an expert's inputs are drawn from a finite, enumerable set, then all possible input-output pairs can be precomputed. MoLA identifies LLM embedding tokens—the outputs of the embedding layer—as naturally fitting this criterion, as they correspond to a finite vocabulary size and are uniquely determined by their token IDs.

MoLA Architecture: Training Phase Modifications

During the training phase, MoLA introduces two critical deviations from standard MoE:

1. Input Transformation: For all routed experts, MoLA replaces their traditional inputs (which are intermediate model features rich in contextual information) with embedding tokens. This is a crucial step as it ensures that the inputs to these experts are drawn from a finite set, making them convertible into lookup experts.
2. Dense Expert Activation: Unlike conventional MoE, which employs sparse activation during training to reduce computational cost, MoLA activates all routed experts densely. The rationale here is that since lookup experts require no computation at inference time, the need for sparse activation to reduce compute is eliminated. While this increases training cost, the primary focus of MoLA is on inference efficiency for edge-scale models. This dense activation also removes the need for auxiliary loss functions often used in sparse MoE training to ensure balanced expert utilization.

Precomputation and Lookup Table Construction

Following training and prior to inference, a critical precomputation step occurs:

The outputs of all routed experts are precomputed for every possible embedding token. Since each embedding token is uniquely determined by its token ID (a consecutive integer), the corresponding expert outputs are also uniquely determined by this ID. This allows for the construction of a massive lookup table. This table effectively maps each token ID to the concatenated outputs of all routed experts that would be activated for that token.

Inference Phase: A Paradigm Shift

The inference phase in MoLA is where the architectural shift delivers its most significant benefits:

1. Expert Parameter Removal: All parameters associated with the routed computation experts (the FFN weights and biases) are removed. They are no longer needed because their functionality has been encapsulated in the lookup table.
2. Lookup Table Offloading: The newly constructed lookup table is stored in secondary storage (e.g., an SSD or flash memory on an edge device).
3. Efficient Lookup and Transfer: During inference, when a token is processed, its ID is used for direct addressing into the lookup table. This direct addressing is highly efficient, as it bypasses complex computations. Once the relevant expert outputs are retrieved from the lookup table in secondary storage, they are transferred into VRAM for subsequent combination and use by other model components.

Other Components and Interpretation

Other essential components of the MoLA architecture, such as the shared FFNs, attention layers, and routers, retain their original design and inputs (intermediate features or outputs from previous layers). The router's role remains to determine the weights for combining the outputs from the lookup experts.

From an alternative perspective, MoLA can be interpreted as a dense model augmented with an offloaded embedding layer. If one conceptually removes the routers and the lookup table at inference time, what remains is essentially a dense model. The lookup table, combined with the router, acts as an additional input of information, much like an expanded embedding space. This "offloaded embedding layer" contributes to improved model performance by providing richer, expert-specific information with almost no extra VRAM or compute cost during runtime. This perspective also opens avenues for future research into novel ways of combining embedding and intermediate features.

Efficiency Analysis

The efficiency gains of MoLA are rooted in a fundamental trade-off:

• Storage Footprint: The lookup table's size scales directly with the vocabulary size and the output dimension of the experts. Consequently, its overall storage footprint is generally larger than that of the offloaded expert parameters in a traditional MoE.
• Loading Overhead: This is where MoLA shines. During inference, MoLA loads only the precomputed output vectors corresponding to the token ID. In contrast, offloaded MoE must load entire expert parameter matrices (weights and biases) for the activated experts. Since parameter matrices are typically much larger than output vectors, MoLA achieves a "significant reduction in transfer overhead," potentially by "several orders of magnitude." This reduction is critical for mitigating the limited storage-to-VRAM bandwidth on edge devices. The direct addressing mechanism further ensures that this loading is fast and predictable, unlike the potentially varied and fragmented loading patterns of MoE.

Experimental Setup & Results

▶ Watch: Scaling insights: embedding tokens enable efficient reparameterization (7:20)

To validate the efficacy of Mixture of Lookup Experts (MoLA), the researchers conducted a series of experiments, primarily focusing on performance, VRAM usage, loading overhead, and decoding speed, particularly in comparison to traditional Mixture of Experts (MoE) models.

Dataset and Fair Comparison

The models were trained and evaluated on a subset of the Pi dataset. To ensure a fair comparison between MoE and MoLA, the experimental setup meticulously controlled key parameters:

• Inference-time VRAM usage: Both architectures were constrained to use the same amount of VRAM during inference.
• Computation: The computational budget for both models was kept equivalent.
• Training-time parameter count: The total number of parameters during training was controlled, allowing for a direct comparison of their respective performance under similar resource envelopes.

Headline Results and Observations

1. Performance Parity: Under the aforementioned controlled conditions, MoE and MoLA achieved comparable performance. This is a crucial result, indicating that MoLA's architectural modifications do not lead to a degradation in model quality despite its significant changes to expert implementation.
2. Storage vs. Loading Trade-off:

• Storage Footprint: As anticipated, MoLA generally uses more storage space for its lookup table compared to the expert parameters of MoE. This is a direct consequence of storing precomputed outputs for every vocabulary token.
• Loading Overhead: Despite the larger storage, the amount of parameters MoLA needs to load into VRAM during inference is dramatically reduced – by several orders of magnitude. This confirms the core hypothesis that loading output vectors is far more efficient than loading entire expert parameter matrices.

1. Training Process: The researchers verified that because MoLA activates all experts densely during training, its training process does not rely on auxiliary loss for guidance. This simplifies the training pipeline compared to many sparse MoE implementations.
2. Decoding Speed on V100 GPU:

• Experiments on a V100 GPU revealed a significant advantage for MoLA in decoding speed, especially with increasing batch sizes.
• For MoE, as batch size increases, multiple tokens might require different experts, leading to a greater number of unique expert parameters needing to be loaded into VRAM. This amplifies the loading cost and becomes a bottleneck.
• In stark contrast, MoLA's inference overhead is primarily due to its lookup tables and lightweight routers, which is negligible. As a result, MoLA achieves inference speeds comparable to that of a dense model, demonstrating its efficiency in real-world decoding scenarios.

Scaling Strategies

The study explored two distinct scaling strategies for MoLA to understand its performance characteristics:

1. Scaling the Size of Routed Experts: This approach involves increasing the internal dimension or capacity of the individual experts that are eventually converted into lookup tables. While this increases training cost, it does not add inference-time overhead because the experts are replaced by static lookup tables. The observation was that performance initially improves with larger expert sizes but eventually becomes bounded by the fixed inference-time architecture. This suggests that simply making individual experts larger has diminishing returns beyond a certain point.
2. Increasing the Number of Experts: This strategy involves adding more independent lookup experts to the model. This increases both training cost and, crucially, inference-time storage and loading overhead (as the lookup table becomes larger and more outputs need to be retrieved). However, this approach proved more effective for improving MoLA's performance. This finding highlights that under fixed compute and VRAM constraints, MoLA's performance can still be enhanced by trading off additional storage and loading costs, offering a clear path for scaling model capacity.

Reparameterization Study (Ablation)

An ablation study was conducted to understand the impact of the key architectural changes, particularly the switch to embedding tokens as expert inputs and the reparameterization into lookup form:

1. Full MoE Activation: Activating all MoE experts densely (without sparse routing) significantly improved performance, confirming the benefit of increased capacity.
2. Expert Granularity: Using fewer, larger experts resulted in a slight drop in performance, suggesting a benefit to having more granular, specialized experts.
3. Embedding Token Input: Changing the expert input from the rich intermediate features to simpler embedding tokens led to a notable performance drop. This was attributed to the loss of contextual information that intermediate features provide.
4. Lookup Reparameterization: Finally, applying the reparameterization to convert these experts into lookup form did not affect the model output when compared to the embedding-token-input-only baseline.

• Key Insight: The study concluded that while switching to embedding tokens as inputs initially weakens performance, this step is essential because it enables the reparameterization into lookup experts. This reparameterization, in turn, allows all experts to be activated densely during inference without incurring additional computational cost. Ultimately, this dense activation of lookup experts effectively recovers the initial performance loss and delivers the desired efficiency benefits.

Practical Implications

▶ Watch: Summary: Lookup experts for efficient edge LLMs (8:00)

The Mixture of Lookup Experts (MoLA) architecture introduces significant practical implications for the deployment and design of large language models, particularly in resource-constrained environments.

For practitioners, infrastructure teams, and model deployers working on edge devices (such as mobile phones, personal computers, or embedded systems with single GPUs), MoLA offers a compelling solution to overcome the long-standing challenge of deploying large MoE models. Traditional MoE models, with their vast parameter counts, are often too large to fit into the limited VRAM of edge hardware, and even with expert offloading, the slow storage-to-VRAM bandwidth creates a critical bottleneck. MoLA directly addresses this by drastically reducing the amount of data that needs to be loaded into VRAM during inference, moving from loading large parameter matrices to small output vectors. This translates directly into faster inference speeds that are comparable to dense models, even as batch sizes increase. This capability unlocks the potential to deliver more powerful and complex LLM experiences directly on user devices, enhancing privacy, reducing latency, and enabling offline functionality.

However, adopting MoLA involves certain trade-offs and limitations:

1. Increased Storage Footprint: The most prominent trade-off is the larger storage footprint required for the lookup table. This table scales with the vocabulary size and the output dimension of the experts. While storage on edge devices (e.g., flash memory) is often more abundant than VRAM, the absolute size of this lookup table could still be substantial for very large vocabularies or highly dimensional experts. Infrastructure teams would need to consider this increased disk space requirement.
2. Precomputation Overhead: The need for a precomputation step before inference means that model updates or fine-tuning would necessitate regenerating the lookup tables, adding a step to the deployment pipeline.
3. Initial Performance Dip: The ablation study revealed an initial performance drop when experts are forced to take only embedding tokens as input, due to a loss of contextual information. While MoLA's full reparameterization and dense activation strategy ultimately recovers and even improves performance, model builders need to understand that this architectural constraint requires a holistic design approach to achieve optimal results.
4. Fixed Inference Architecture: The scaling strategy experiments indicated that while increasing the number of experts is effective, simply scaling the size of individual lookup experts has diminishing returns because the inference-time architecture becomes fixed after precomputation. This implies that MoLA's design might favor breadth (more experts) over depth (larger individual experts) for performance scaling.
5. Applicability: The core idea relies on the expert inputs being convertible into a finite, enumerable set. While embedding tokens are a natural fit for LLMs, applying lookup experts to other domains or types of model layers might require careful consideration of input characteristics.

Despite these considerations, MoLA opens up new avenues for efficient model design, particularly for VRAM-constrained environments. The concept of "intermediate embedding layers within the model" or "offloaded embedding layers" can inspire future research into hybrid architectures that judiciously combine computation-intensive components with precomputed, low-latency lookup structures. This could lead to a new generation of LLMs that are not only powerful but also inherently designed for ubiquitous deployment across a diverse ecosystem of hardware.

Key Takeaways

• MoLA addresses MoE deployment on edge devices by replacing traditional computation-intensive FFN experts with efficient lookup experts.
• Lookup experts are derived from precomputing outputs for embedding tokens, stored in a lookup table that maps token IDs to expert outputs.
• MoLA significantly reduces VRAM loading overhead during inference, loading orders of magnitude less data than offloaded MoE, achieving dense model-like speeds on a V100 GPU.
• It trades a higher storage footprint for the lookup table for drastically lower loading latency and faster inference, making it suitable for bandwidth-constrained edge environments.
• Performance is comparable to MoE under similar inference-time VRAM, computation, and training parameter constraints, without requiring auxiliary loss functions during training.
• The architecture enables dense activation of experts during inference without additional computational cost, a key factor in recovering performance lost by switching to embedding token inputs.

About the Speaker(s)

The primary author of the work, Shibo Jie, is affiliated with Peking University. The presentation at ICML 2025 was given by another individual on Shibo Jie's behalf. The collaborative nature of this research is further highlighted by the contributions of co-authors Yehui Tang, Kai Han, Yitong Li, Duyu Tang, Zhi-Hong Deng, and Yunhe Wang.