Overview

MoFEbaseD2D studies sparse parameter expansion of GPT-2 Small on FineWeb-Edu 10BT. It compares three methods under the same training-token budget:

• Dense: continued pretraining of GPT-2 Small.
• MoFE: Mixture of Factorized Experts in the final three GPT-2 MLP blocks.
• Upcycling: sparse expansion using complete copies of the final three MLPs.

Unless explicitly labeled otherwise, MoFE means MoFE group LR. The older single-learning-rate MoFE is preserved only as historical context in timeline.md and is excluded from the final comparison. The latest MoFE ablation study is archived in results/mofe_ablation_80k.

Method

MoFE replaces the MLPs in Transformer blocks 9, 10, and 11. Each converted layer keeps one always-active dense shared expert and adds 16 factorized private experts. A token-choice router selects the top 3 private experts per token.

With a 4 x 4 Cartesian factor-bank construction and expert e = 4i + j:

W1_e = A1_i C1_e B1_j
W2_e = A2_i C2_e B2_j

The factorized path runs as x -> B -> C -> A without materializing complete expert matrices per token. The shared branch is copied from the original GPT-2 MLP. Private output cores and biases are zero-initialized, so the converted model preserves the Dense GPT-2 function at initialization.

Final Protocol

MoFE group LR uses 1e-5 for the backbone/shared experts, 2e-5 for private experts, and 3e-5 for routers. Dense and Upcycling use 1e-5.

Continuation checkpoints did not store the streaming dataloader cursor. A new process restored model and optimizer state but rebuilt the fixed-seed stream from its beginning. Token counts across continuation boundaries therefore mean processed tokens, not unique tokens. Interpret both the 50k result and the 120k extension below with this limitation.

Results at 50k

FineWeb-Edu held-out results:

Downstream results use the original acc metric for both tasks:

ARC and WikiText are excluded from the final benchmark. HellaSwag acc_norm is present in raw lm-eval JSON but is not used in the primary table.

Diagnostic Extension to 120k

Step 120k corresponds to 3.93216B processed tokens. The 50k-to-80k and 80k-to-120k continuations restarted the training stream and therefore overlap in examples; MoFE restarted it once more at 95k. These checkpoints support diagnostic and method comparisons, not an uninterrupted unique-token scaling claim.

MoFE Ablation Study at 80k

The final MoFE ablation study compares Full MoFE with four architecture or initialization variants. All ablations follow the same FineWeb-Edu 10BT protocol as the main experiment. After the 50k checkpoint, each ablation resumed in a new process to 80k; this reset the streaming dataloader cursor and matches the 50k-to-80k diagnostic protocol used for Dense, MoFE, and Upcycling.

CalFLOPs was run with batch size 1, sequence length 1024, GPT-2 eager attention, and forward FLOPs only.

No A/B sharing increases parameter count because the A/B factor bank grows from 4 groups to 16 independent parameter copies. Its active forward FLOPs match Full MoFE because each token still executes top-k=3 private expert paths with the same matrix shapes. This ablation is slightly better at 80k, but it uses about 1.76x as many MoFE parameters; the more compact Full MoFE remains the default main method.

Raw Experiment Data

The 50k archive is in results/final_50k:

• validation/validation_loss_50k.csv: 50k held-out loss data for all three methods.
• validation_prediction_accuracy/raw/: 30 original JSON points, 10 checkpoints per method.
• downstream/: original 50k LAMBADA/HellaSwag JSON outputs for all three methods.
• figures/: the final side-by-side loss and token-accuracy figure.

The 80k-to-120k diagnostic archive is in results/final_120k:

• validation/: complete held-out loss records through 120k.
• validation_prediction_accuracy/: 72 raw evaluations and the 5k-to-120k CSV.
• downstream/: consolidated and raw LAMBADA/HellaSwag results at 80k, 100k, 110k, and 120k.
• figures/: the combined 120k loss and token-accuracy plot.

The 80k MoFE ablation archive is in results/mofe_ablation_80k:

• validation/: 0-50k held-out validation loss for Full MoFE and four ablations.
• validation_prediction_accuracy/: 80k next-token accuracy summary, raw JSON, and logs for all five checkpoints.
• compute/: CalFLOPs parameter and forward-FLOP summaries for all five models.
• configs/: the four ablation configs.
• figures/: the 0-50k validation-loss comparison plot.

A downloadable zip package is available at:

/data/chenpengxu/MoFEbaseD2D_runtime/packages/mofe_ablation_results_20260731_133549.zip

All historical experiments remain available in archive.

Code and Usage

The active implementation lives in MoFE/. Install dependencies and run tests:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

python -m unittest \
  MoFE.tests.test_data \
  MoFE.tests.test_mofe \
  MoFE.tests.test_upcycling

The full experiment record is maintained in timeline.md.

Upstream

This project is an independent derivative of D2DMoE, based on upstream commit a7027cdc1f01c9c618c39eebe639d1664549b066. The upstream project and this derivative use the MIT License. The associated paper is:

Filip Szatkowski, Bartosz Wojcik, Mikolaj Piorczynski, Simone Scardapane. Exploiting Activation Sparsity with Dense to Dynamic-k Mixture-of-Experts Conversion. NeurIPS 2024. https://arxiv.org/abs/2310.04361