Deconstructing Pre-training: Knowledge Attribution Analysis in MoE and Dense Models

TL;DR

The paper investigates how MoE architectures reshape knowledge acquisition during pre-training compared with dense transformers, addressing scalability limits by decoupling capacity from per-token compute. It introduces Gated-LPI, a neuron-level attribution method that decomposes log-probability gains across expert and attention components, enabling time-resolved comparisons over extensive pre-training. Key findings show that MoE develops a low-entropy backbone with an early consolidation of both FFN and ATTN layers, and that this structure supports robust, distributed knowledge storage as evidenced by resilience under targeted ablations. These results provide a principled link between architectural sparsity, learning dynamics, and interpretability, suggesting design principles for stability-aware routing and pruning in large-scale sparse models.

Abstract

Mixture-of-Experts (MoE) architectures decouple model capacity from per-token computation, enabling scaling beyond the computational limits imposed by dense scaling laws. Yet how MoE architectures shape knowledge acquisition during pre-training, and how this process differs from dense architectures, remains unknown. To address this issue, we introduce Gated-LPI (Log-Probability Increase), a neuron-level attribution metric that decomposes log-probability increase across neurons. We present a time-resolved comparison of knowledge acquisition dynamics in MoE and dense architectures, tracking checkpoints over 1.2M training steps (~ 5.0T tokens) and 600K training steps (~ 2.5T tokens), respectively. Our experiments uncover three patterns: (1) Low-entropy backbone. The top approximately 1% of MoE neurons capture over 45% of positive updates, forming a high-utility core, which is absent in the dense baseline. (2) Early consolidation. The MoE model locks into a stable importance profile within < 100K steps, whereas the dense model remains volatile throughout training. (3) Functional robustness. Masking the ten most important MoE attention heads reduces relational HIT@10 by < 10%, compared with > 50% for the dense model, showing that sparsity fosters distributed -- rather than brittle -- knowledge storage. These patterns collectively demonstrate that sparsity fosters an intrinsically stable and distributed computational backbone from early in training, helping bridge the gap between sparse architectures and training-time interpretability.

Figures (5)

• Figure 1: Top-1% FFN and ATTN neurons Jaccard overlap between consecutive checkpoints.
• Figure 2: Mean FFN and ATTN importance scores across all layers over training steps. OLMoE shows smoother and earlier stabilization compared to OLMo.
• Figure 3: Snapshots of FFN layer importance distribution at different training steps.
• Figure 4: Snapshots of ATTN layer importance distribution at different training steps.
• Figure 5: Performance of the model under ablation interventions (Top-1 head, Top-10 heads and Top-1% FFN neurons) measured by the HIT@10 metric.