Research map / Architectures and efficiency / Mixture of experts
How does Architecture Influence the Base Capabilities of Pre-trained Language Models? A Case Study Based on FFN-Wider and MoE Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Explains why FFN-Wider Transformers lose base capabilities by attributing it to a reduced contribution ratio of MHA (a combination function), proposes Combination Enhanced Architecture (CEA) that moves part of the FFN inside MHA, and applies it to a 14B MoE model.
Problem
FFN-Wider and MoE Transformers achieve comparable pre-training loss to vanilla Transformers but show degraded base capabilities (OOD language modeling, transfer, few-shot); the architectural cause is unexplained.
Approach
Hypothesizes that the actual contribution ratio of MHA (a combination function) to pre-training is the key factor for base capabilities; proposes CAA/CEA that splits the wider FFN into an Outer-FFN (transformation) and an Inner-FFN placed inside MHA with a direct current-position pathway, and extends this to MoE layers.
Key results
- On a 14B MoE with 100B tokens, MoE 14B w/ CEA lowered SlimPajama OOD loss (e.g., Arxiv 2.320 -> 2.239, Book 2.761 -> 2.585) versus Vanilla MoE 14B under matched pre-training loss, steps, params and compute.
- For FFN-Wider BERT (H=768), CEA raised GLUE average from 76.66 to 78.20 at H=128 and improved Pile average loss from -2.196 to -2.149 at H=768.
- On the MoE 14B, 5-shot MMLU rose from 30.6 to 31.4 and LAMBADA 0-shot accuracy from 36.4 to 39.7 with CEA.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the standard MoE placement with CEA, placing MoE inside MHA as Inner-MoE | 14B MoE Transformer · SlimPajama-Arxiv, Book, Github, Stack and Wiki | loss: 2.320 -> 2.239; 2.761 -> 2.585; 1.989 -> 1.840; 2.588 -> 2.465; 3.155 -> 2.964 | Helped |
| Replace the standard MoE placement with CEA, placing MoE inside MHA as Inner-MoE | 14B MoE Transformer · LAMBADA (0-shot) | perplexity: 24.8 -> 19.3 | Helped |
| Replace the standard MoE placement with CEA, placing MoE inside MHA as Inner-MoE | 14B MoE Transformer · LAMBADA (0-shot) | accuracy: 36.4 -> 39.7 | Helped |
| Replace the standard MoE placement with CEA, placing MoE inside MHA as Inner-MoE | 14B MoE Transformer · BoolQ (5-shot) | accuracy: 62.2 -> 62.1 | No clear effect |
| Use CEA instead of the FFN-Wider architecture | BERT H=768 · Pile (22 parts, average) | language modeling score: -2.196 -> -2.149 | Helped |
| Use CAA with 37.5% Outer-FFN width ratio instead of FFN-Wider | GPT H=768 · 0-shot (9 tasks, average) | score: 44.60 -> 45.27 | Helped |
| Remove the direct pathway in MHA (ablation) | FFN-Wider models with CEA · Base-capability evaluations | base capability | Hurt |
| Reduce Outer-FFN width ratio from 100% toward 0% (ablation) | Combination Adjustable Architecture (CAA) · BERT OOD language modeling | OOD performance | Helped |
| Set Outer-FFN width to zero (ablation) | GPT Combination Adjustable Architecture · GPT base capabilities | base capability | Hurt |
Surprising findings
- Pushing the Outer-FFN width to 0% for GPT caused anomalous degradation, suggesting a minor transformation contribution is necessary.
- The Inner-FFN design had to be paired with a direct pathway in MHA; removing it caused a significant decline in base capabilities.
Limitations
- Analysis is restricted to language-modeling pre-training objectives, so conclusions may not generalize to other objectives.
- Experiments are limited to small/medium BERT and GPT models and a single 14B MoE, with max sequence length 128 for the dense GPT models due to compute limits.
- Contribution ratio measures rely on proxy estimates (mutual information via clustering, token prediction), and the BERT case retains residual transformation leakage.
Experimental setup
Datasets: Wikipedia BooksCorpus SlimPajama Pile
Benchmarks: Pile GLUE SuperGLUE HellaSwag PIQA OpenBookQA ARC WinoGrande LAMBADA MMLU BoolQ RACE
Metrics: loss perplexity accuracy fine-tuning score
Baselines: Vanilla BERT FFN-Wider BERT Vanilla GPT FFN-Wider GPT Vanilla MoE 14B Vanilla GPT 1.3B
Builds on: BERT GPT-3 Mixture of Experts RoPE FlashAttention-2
Closest papers by meaning
- How Does Sequence Modeling Architecture Influence Base Capabilities of Pre-trained Language Models? Exploring Key Architecture Design Principles to Avoid Base Capabilities Degradation (NeurIPS 2025 · similarity 0.92)
- Sparse Layers are Critical to Scaling Looped Language Models (NeurIPS 2026 · similarity 0.87)
- Exploring the LLM Journey from Cognition to Expression with Linear Representations (ICML 2024 · similarity 0.87)
- Joint Fine-tuning and Conversion of Pretrained Speech and Language Models towards Linear Complexity (ICLR 2025 · similarity 0.87)
- Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers (NeurIPS 2025 · similarity 0.87)
- Law of the Weakest Link: Cross Capabilities of Large Language Models (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.