Single-author paper, Google proposes millions of expert Mixture, surpassing dense feedforward and sparse MoE

Unlock the potential to further scale Transformer while maintaining computational efficiency.

Feedforward (FFW) layers in the standard Transformer architecture result in a linear increase in computational cost and activation memory as the hidden layer width increases. As the size of large language models (LLM) continues to increase, the sparse mixed expert (MoE) architecture has become a feasible method to solve this problem, which separates the model size from the computational cost. Many emerging MoE models can achieve better performance and more powerful performance at the same size.

The recently discovered fine-grained MoE expansion law shows that higher granularity leads to better performance. However, existing MoE models are limited to a low number of experts due to computational and optimization challenges.

This Tuesday, New research from Google DeepMind introduces a parameter-efficient expert retrieval mechanism that leverages product key technology to perform sparse retrieval from one million micro-experts.

Link: https://arxiv.org/abs/2407.04153

This approach attempts to decouple the computational cost from parameter count by efficiently concatenating to a large number of tiny experts through a learned index structure for routing . Demonstrates superior efficiency compared to dense FFW, coarse-grained MoE, and Product Key Memory (PKM) layers.

This work introduces the Parameter Efficient Expert Retrieval (PEER) architecture (parameter efficient expert retrieval), which uses product key retrieval to efficiently route to a large number of experts, separating the computational cost from the amount of parameters. This design demonstrated superior computational performance levels in experiments, positioning it as a competitive alternative to dense FFW layers for extending base models. The main contributions of this work are:

Exploration of extreme MoE settings: Unlike the focus on a few large experts in previous MoE studies, this work investigates the underexplored situation of numerous small experts.

Learned Index Structure for Routing: First demonstration that a learned index structure can efficiently route to over a million experts.

New layer design: Combining product key routing with single-neuron experts, we introduce the PEER layer, which scales layer capacity without significant computational overhead. Empirical results demonstrate higher efficiency compared to dense FFW, coarse-grained MoE, and Product Key Memory (PKM) layers.

Comprehensive ablation study: We study the impact of different design choices for PEER (such as number of experts, activity parameters, number of heads, and query batch normalization) on language modeling tasks.

Method Introduction

In this section, the researcher explains in detail the Parametric Efficient Expert Retrieval (PEER) layer, a hybrid expert architecture that uses product keys in routing and single-neuron MLP As an expert. Figure 2 below shows the calculation process within the PEER layer.

PEER Layer Overview. Formally, the PEER layer is a function f : R^n → R^m, which consists of three parts: a pool of N experts E := {e_i}^N_i=1, where each expert e_i : R^n → R^m shares the same signature as f; a corresponding set of N product keys K := {k_i}^N_i=1 ⊂ R^d; and a query network q : R^n → R ^d, which maps the input vector x ∈ R^n to the query vector q (x).

Let T_k represent the top-k operator. Given an input x, first retrieve the subset of k experts whose corresponding product keys have the highest inner product with the query q (x).

Then apply a nonlinear activation (such as softmax or sigmoid) to the query key inner product of the top k experts to obtain the routing score.

Finally the output is calculated by linearly combining the expert outputs weighted by routing scores.

Product Key Retrieval. Since the researchers intend to use a large number of experts (N ≥ 10^6), simply computing the top k indexes in Equation 1 can be very costly, so the product key retrieval technique is applied. Instead of using N independent d-dimensional vectors as keys k_i, they create them by concatenating vectors from two independent d/2-dimensional subkey sets (i.e., C, C ′ ⊂ R d/2):

Parametric efficient expert and multi-head search. Unlike other MoE architectures, these architectures typically set each expert's hidden layer to the same size as other FFW layers. In PEER, each expert e_i is a singleton MLP, in other words, it has only one hidden layer with a single neuron:

The researchers did not change the size of the individual expert, but used Multi-head retrieval is used to adjust the expressive ability of the PEER layer, which is similar to the multi-head attention mechanism in the transformer and the multi-head memory in PKM.

Specifically, they use h independent query networks, each network computes its own query and retrieves a separate set of k experts. However, different heads share the same expert pool, with the same set of product keys. The output of these h heads is simply summarized as follows:

Why do we need a large number of small experts ? A given MoE layer can be characterized by three hyperparameters: the total number of parameters P, the number of active parameters per token P_active, and the size of a single expert P_expert. Krajewski et al. (2024) showed that the scaling law of the MoE model has the following form:

For PEER, the researcher uses the smallest possible expert size by setting d_expert = 1, and the number of activated neurons is the search head The number is multiplied by the number of experts retrieved per head: d_active = hk. Therefore, the granularity of PEER is always G = P_active/P_expert = d_active/d_expert = hk.

Experimental results

Let’s first look at the evaluation results on the language modeling data set.

After determining the computationally optimal model for each method based on the isoFLOP curve, the researchers evaluated the performance of these pre-trained models on the following popular language modeling datasets:

• Curation Corpus
• Lambada
• Pile
• Wikitext
• Pre-training dataset C4

Table 1 below shows the evaluation results. We grouped the models based on the FLOP budget used during training. As can be seen, PEER has the lowest perplexity on these language modeling datasets.

In the ablation experiment, the researchers changed the total number of experts. The models shown in the isoFLOP curve in Figure 1 below all have over one million (1024^2) experts.

The researcher chose the model with the optimal position of isoFLOP and changed the number of experts in the PEER layer (N = 128^2, 256^2, 512^2, 1024^2), while keeping the number of active experts unchanged (h = 8, k = 16). The results are shown in Figure 3(a) below.

It can be seen that the isoFLOP curve interpolates between the PEER model with 1024^2 experts and the corresponding dense backbone without replacing the FFW layer in the middle block with a PEER layer. This shows that model performance can be improved simply by increasing the number of experts.

At the same time, the researchers changed the number of active experts. They systematically varied the number of active experts (hk = 32, 64, 128, 256, 512) while keeping the total number of experts constant (N = 1024^2). For a given hk, the researcher then jointly changes h and k to determine the best combination. Figure 3(b) below plots the isoFLOP curve with respect to the number of heads (h).

Table 2 below lists the expert usage and unevenness for different numbers of experts with and without BN. It can be seen that even for 1M experts, the expert utilization rate is close to 100%, and using BN can make the expert utilization rate more balanced and the confusion level lower. These findings demonstrate the effectiveness of the PEER model in leveraging a large number of experts.

The researchers also compared the isoFLOP curves with and without BN. Figure 4 below shows that PEER models with BN can generally achieve lower perplexity. Although the difference is not significant, it is most noticeable near the isoFLOP optimal region.

Introduction to the author

PEER study has only one author, Xu He (Owen), who is a research scientist at Google DeepMind and graduated with a PhD from the University of Groningen in the Netherlands in 2017.

The above is the detailed content of Single-author paper, Google proposes millions of expert Mixture, surpassing dense feedforward and sparse MoE. For more information, please follow other related articles on the PHP Chinese website!