PaperlayerView exact revision on arXiv

Source-linked AI summary

MoEfication: Transformer Feed-forward Layers are Mixtures of Experts

Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, Jie Zhou

arXiv:2110.01786v3cs.CL

TL;DR

Transformer FFN computational patterns remain unclear despite FFNs storing substantial knowledge, and most inputs activate only a small fraction of their neurons. The paper proposes MoEfication, which partitions FFNs into experts and routes inputs conditionally. Across models and downstream tasks, it preserves over 95% original performance using 10% to 30% of FFN parameters while reducing inference computation.

  • Problem

    FFN computational patterns remain unclear, despite FFNs comprising nearly two-thirds of Transformer parameters and storing linguistic and factual knowledge.

  • Method

    MoEfication splits FFN parameters into functional experts and builds routers to select experts conditionally for each input.

  • Results

    10% to 30% of FFN parameters maintain over 95% original performance across models and downstream tasks, with 2x speedup using 20% of FFN parameters.

  • Takeaways & Limitations

    MoEfication reduces inference FLOPS and provides a fine-grained perspective for studying FFN routing patterns and inner mechanisms.

  • Takeaways & Limitations

    The paper notes that some MoE operations remain unresolved and may limit practical efficiency improvements.

Abstract

from arXiv · show

Recent work has shown that feed-forward networks (FFNs) in pre-trained Transformers are a key component, storing various linguistic and factual knowledge. However, the computational patterns of FFNs are still unclear. In this work, we study the computational patterns of FFNs and observe that most inputs only activate a tiny ratio of neurons of FFNs. This phenomenon is similar to the sparsity of the human brain, which drives research on functional partitions of the human brain. To verify whether functional partitions also emerge in FFNs, we propose to convert a model into its MoE version with the same parameters, namely MoEfication. Specifically, MoEfication consists of two phases: (1) splitting the parameters of FFNs into multiple functional partitions as experts, and (2) building expert routers to decide which experts will be used for each input. Experimental results show that MoEfication can conditionally use 10% to 30% of FFN parameters while maintaining over 95% original performance for different models on various downstream tasks. Besides, MoEfication brings two advantages: (1) it significantly reduces the FLOPS of inference, i.e., 2x speedup with 25% of FFN parameters, and (2) it provides a fine-grained perspective to study the inner mechanism of FFNs. The source code of this paper can be obtained from https://github.com/thunlp/MoEfication.

1 Introduction

Transformer FFN computation was underexplored despite FFNs comprising nearly two-thirds of model parameters. The paper finds sparse neuron activation and introduces MoEfication to partition FFNs into experts and route inputs conditionally.

  • Motivation: FFNs constitute nearly two-thirds of Transformer parameters, yet their computational patterns remain unclear.Prior interpretation work largely focused on attention, while FFNs were recognized as knowledge-storing memory networks.
  • Sparse Activation: 90% of inputs to fine-tuned 700-million-parameter T5-Large activate fewer than 5% of neurons.The paper treats this sparse activation as analogous to sparsity in the human brain and asks whether FFNs contain functional partitions.
  • MoEfication: MoEfication splits FFN parameters into functional experts and builds routers that select experts for each input.Expert construction groups frequently co-activated neurons, while expert selection chooses experts covering as many activated neurons as possible.
  • Results: 10% to 30% of FFN parameters preserve over 95% original performance across T5 and BERT on GLUE and QA benchmarks.The results support the existence of functional partitions in pretrained Transformer FFNs.
  • Benefits: Using 25% of FFN parameters yields 2x CPU and 1.2x GPU inference speedups while enabling fine-grained FFN interpretation.The paper studies routing patterns as an interpretability perspective on FFN inner mechanisms.

2 Related Work

Related work interprets Transformer components, develops large-scale MoE models, and accelerates pretrained language models. MoEfication differs by exploiting naturally sparse FFN activation for inference-time conditional computation.

  • Interpretation: Transformer interpretation studies have examined layer functionality, attention mechanisms, and FFNs as knowledge-storing memory networks.This paper focuses on how knowledge stored in FFNs is conditionally activated.
  • Large-scale MoE: Prior MoE research builds separate networks or experts that independently handle subsets of training examples.Large-scale language models such as GShard and Switch-Transformer study MoE architectures and training strategies.
  • Distinction: MoEfication converts an existing model into an MoE version by exploiting naturally sparse FFN activation for inference efficiency.Unlike prior large-scale MoE work, it does not primarily introduce MoE through new model training.
  • Model Acceleration: PLM acceleration includes distillation, pruning, attention approximation, quantization, and dynamic inference.These approaches target reduced time or space complexity for pretrained language models.

3 MoEfication

MoEfication uses sparse FFN activations to preserve the original computation with fewer active parameters. It constructs balanced experts from neuron relationships, then routes each input to a small subset of experts.

  • Overall Framework: MoEfication divides an FFN into independent experts and designs a router to select suitable experts for each input.The framework has two phases: expert construction and expert selection.
  • Sparse Activation: Inactive zero-valued intermediate activations can be omitted, while input-dependent neuron activation motivates selecting neurons before FFN computation.This reduces computation without pruning the model parameters.
  • Illustration: For a 2-dimensional input and 4-neuron intermediate layer, two positive activations require only corresponding 2 × 2 submatrices and one selected expert.The example illustrates how the MoEfied FFN reproduces the original output for that input.
  • Expert Construction: Expert construction groups frequently co-activated neurons into equally sized parts, retaining dmodel input and output dimensions.Equal expert sizes are chosen to improve parallel computation performance.
  • Parameter Rearrangement: A permutation mapping groups neurons by expert and rearranges FFN parameters without changing the output representation.The mapping uses each neuron's expert index and within-expert position to construct the permutation.
  • Expert Construction: Parameter Clustering Split uses balanced K-Means on W1 columns, whereas Co-Activation Graph Split partitions neurons with strong internal co-activation connections.The two methods use parameter similarity and activation statistics, respectively.
  • Expert Selection: Expert selection assigns scores to experts and chooses the n highest-scoring experts to approximate the original FFN output.Groundtruth Selection minimizes intermediate-output error approximately, while Similarity and MLP Selection use parameter or predicted-activation information.
  • Expert Selection: Similarity Selection scores experts by cosine similarity between input and expert representation, while MLP Selection predicts each expert's positive activation sum.These routers approximate the ideal groundtruth selection using input or learned prediction signals.

4 Experiment

Experiments across T5 and BERT models show sparse FFN activation and effective MoEfication, which preserves most performance while selecting a fraction of neurons. Results also compare model scales, activation functions, expert strategies, and evaluation settings.

  • Experimental setup: The experiments evaluate T5 variants and distilled models on SST-2, MNLI, and RACE, with additional GLUE and SQuAD results reported in the appendix.The datasets represent single-sentence classification, sentence-pair classification, and reading comprehension.
  • Activation sparsity: Most inputs activate fewer than 10% of neurons, while larger T5 models exhibit greater activation sparsity than smaller models.In T5-XLarge, 80% of inputs activate fewer than 3% of neurons, compared with more than 3% for 40% of T5-Small inputs.
  • MoEfication performance: 10% to 30% of FFN parameters preserve over 95% of original performance across models and downstream tasks.T5-XLarge reaches nearly 98% relative performance on SST-2 and MNLI with 10% neurons, whereas T5-Small requires 30% to 40%.
  • MoEfication performance: MoEfication performs similarly on T5-Small and distilled T5-Small, indicating compatibility with knowledge distillation for efficient inference.The reported sparsity curves for T5-Small-Distill and T5-Small are also close.
  • GeLU-based models: BERT retains sparse activation, with over 80% of inputs activating fewer than 10% of neurons, while selecting 30% to 40% recovers 97% performance.BERT requires more neurons than T5 because its GeLU-derived activation is slightly denser; the model is first adapted to ReLU before MoEfication.
  • MoEfication strategies: Co-activation Graph Split is the best expert-construction method, while MLP Selection is the best expert-selection method across the tested combinations.Co-activation Graph Split groups neurons that activate simultaneously, and MLP Selection works with both tested clustering-based split methods.

5 Analysis

MoEfication improves inference efficiency while revealing uneven expert-routing patterns in MoEfied models. Its benefits are constrained by unchanged attention costs and difficult GPU parallelization.

  • Routing Patterns: Expert selection is highly unbalanced across the 64 experts in each encoder layer, rather than matching the ideal frequency of 0.2.Some experts are selected by more than 80% of inputs, while others are long-tail experts.
  • Efficiency Improvement: 2x speedup is achieved with a 25% selected-parameter ratio, substantially reducing total FLOPS for T5-Large inference.CPU speedup is close to FLOPS speedup, while GPU speedup is smaller because some MoE operations are difficult to parallelize.
  • Efficiency Improvement: Reducing the selected ratio produces diminishing gains because attention costs remain unchanged and become the bottleneck.Halving 25% to 12.5% yields 1.2x gain, whereas halving 50% to 25% yields 1.3x.
  • Efficiency Improvement: GPU acceleration remains below FLOPS-based expectations because some MoE operations cannot be easily parallelized.Distributed MoE packages are being developed to address this parallel-computing bottleneck.
  • Routing Patterns: Most-selected experts have similar self- and inter-input similarities, whereas least-selected experts show higher self-similarity than inter-similarity.The latter pattern suggests that inputs routed to least-selected experts form obvious clusters.
  • Routing Patterns: Routing patterns comprise general experts that serve most inputs and input-specific experts that are seldom used for particular domains or tasks.The authors suggest this observation may inform future MoE training.

6 Conclusion

The paper argues that Transformer FFNs naturally contain expert-like functional partitions and introduces MoEfication to expose them without changing model parameters. MoEfied models preserve comparable performance while reducing inference computation and revealing general and input-specific routing patterns.

  • Conclusion: MoEfication converts a normal Transformer into an MoE version with the same parameters by exploiting sparse FFN activation.The method supports both efficiency improvements and analysis of FFN routing behavior.
  • Conclusion: 10% to 30% of FFN parameters suffice for MoEfied models to achieve comparable performance to the original models.This result supports the view that Transformer FFNs contain naturally reusable expert partitions.
  • Conclusion: 2x inference speedup is reported with 20% of FFN parameters, alongside substantially reduced FLOPS.The conclusion presents efficiency as a practical benefit of conditional FFN computation.
  • Conclusion: MoEfied routing reveals both general experts and input-specific experts, offering a fine-grained perspective on FFN mechanisms.The authors suggest these routing patterns may inspire future work on training MoE models.

A MoEfication on Other Datasets

On downstream NLU tasks with T5-Large, MoEfication reduces FFN computation while causing only a small average performance loss. Parameter calibration corrects precision errors and can recover performance on specific tasks.

  • Parameter calibration effectively addresses precision errors introduced by MoEfication.
  • Nearly 4% improvement on CoLA brings MoEfied+Calib to the same average performance as MoEfied+GT.

B Activation Statistics before Fine-tuning

Before fine-tuning, pre-trained language models already exhibit sparse FFN activation, with most inputs activating only a small fraction of neurons. Fine-tuning generally raises average activation ratios, though the increase varies across datasets.

  • Before fine-tuning, pre-trained language models exhibit sparse activation, and fine-tuning changes this phenomenon little overall.Activation statistics were collected on pre-training data containing about 50,000 input tokens.
  • Figure 7 compares cumulative distributions of activated-neuron ratios across models before fine-tuning.
  • Fine-tuning increases the average activation ratio for most models.

C Results of Graph Partition

The co-activation graphs have low edgecut ratios, supporting partitioning neurons into experts that preserve frequently co-activated groups. Separately, trained MLP selectors achieve moderate prediction accuracy in encoder and decoder layers.

  • Results of Graph Partition: Co-activation Graph Split groups neurons with frequent simultaneous activation into partitions with relatively few crossing edges.
  • Results of Graph Partition: Low overall edgecut ratios indicate that co-activation graphs are suitable for partitioning.All decoder-layer edgecut ratios are smaller than 0.001.
  • Results of Graph Partition: MLP selection accuracy is about 0.8 for encoder layers and about 0.7 for decoder layers.

E Relative Cost of Routing

Routing adds a model-independent relative cost compared with FFN computation. The surrounding evaluations also examine graph partitioning, pruning, and pre-training strategies for MoE models.

  • Relative Cost of Routing: The relative cost of routing compared with FFN computation is constant across different models and similar to MLP Selection.
  • Results of Graph Partition: Graph partitioning minimizes edges crossing sub-graphs through coarsening, partitioning, and refinement phases.Coarsening groups highly connected nodes, including neurons that often activate simultaneously.
  • MoE Pre-training: Standard pre-training with MoEfication performs better on validation MLM loss than pre-training a MoE model from scratch.
  • Comparison with Model Pruning: MoEfication maintains good performance by selectively activating network parts according to each input, whereas model pruning significantly degrades performance.

H MoEfication vs. MoE pre-training

This subsection compares MoE models pre-trained from scratch with standard models transformed through MoEfication under matched pre-training conditions. The reported validation loss favors standard pre-training and indicates that MoEfication performs better than direct MoE pre-training.

  • Experimental setup: The comparison uses a pre-trained-from-scratch MoE model and a standard model transformed by MoEfication.Both models are evaluated against a standard pre-trained model under the experimental setup.
  • Experimental setup: The models share the same model size and use T5-Small architecture with WikiText-103 masked language-model pre-training.The setup uses a batch size of 4096.
  • Results: The standard pre-trained model achieves lower validation MLM loss than the MoE model pre-trained from scratch.The authors attribute this difference to the greater optimization difficulty of MoE models under restricted expert selection.
  • Results: MoEfied models achieve better performance than the MoE model pre-trained from scratch.This supports pre-training a standard model before applying MoEfication rather than directly pre-training the MoE model.
Loading 2110.01786v3…