File size: 3,255 Bytes
187c341 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | ---
license: cc-by-nc-4.0
language:
- en
tags:
- babylm
- babylm-2026
- mixture-of-experts
- msit
- xpertgpt
- swiglu
- custom_code
- safetensors
library_name: transformers
pipeline_tag: text-generation
---
# XpertGPT (SwiGLU & Sliding Window 64, 16, 8, 4)
XpertGPT is a sparse Mixture of Experts (MoE) language model designed for data-efficient pretraining under the **BabyLM 2026 challenge (Strict-Small 10M track)**. It leverages **Parallelized Multi-Scale Information Transmission (MSIT)** and **Expert Choice Routing** to maximize representational capacity within restricted token budgets.
This version implements **SwiGLU Feed-Forward Networks** across all global and parallel expert blocks, along with corrected LayerNorms, redundant residual removal, and **sliding window attention sizes** of `[64, 16, 8, 4]` tokens.
---
## 1. SwiGLU Feed-Forward Networks (FFN)
Instead of the standard Feed-Forward sequential structure (Linear -> GELU -> Linear), this model replaces FFN layers with the **SwiGLU (Swish Gated Linear Unit)** variant to improve model capacity and training stability:
$$\text{FFN}_{\text{SwiGLU}}(x) = \left(\text{Swish}(x W) \otimes x V\right) W_2$$
Where the Swish function is implemented using SiLU:
* **Gate Linear projections ($W$, $V$)**: Projects input dimension `dim` to `hidden_dim`.
* **Out Linear projection ($W_2$)**: Projects back to `dim`.
* To maintain parameter counts equivalent to standard `dim * 4` sequential FFNs, the hidden dimension is scaled to:
$$\text{hidden\_dim} = \text{round\_to\_multiple\_of\_8}\left(\frac{8}{3} \times \text{dim}\right)$$
---
## 2. Expert Sliding Window Layout
The four parallel MoE experts are configured with distinct sliding window attention constraints:
* **Expert 1**: Window size `64` tokens
* **Expert 2**: Window size `16` tokens
* **Expert 3**: Window size `8` tokens
* **Expert 4**: Window size `4` tokens
---
## 3. Architectural Layout & Changes
This model implements:
1. **Removal of Redundant Residual (`res3`)**:
* Removed redundant residual connection around the global dense block. Gated input is now simply $X_2 = X_1$.
2. **Introduction of Post-Block LayerNorm (`ln3`)**:
* LayerNorm `ln3` is added after the SwiGLU addition inside every `MSITBranchBlock`.
* Formulation: $X_{\text{out}} = \text{LayerNorm}(X^{(2)})$.
3. **Introduction of Post-MoE LayerNorm (`ln_post_moe`)**:
* LayerNorm `ln_post_moe` is added after the Residual 4 MoE aggregation.
* Formulation: $X_{\text{out}} = \text{LayerNorm}(X_2 + X_{3, \text{full}})$.
---
## 4. How to Load and Use Checkpoints (Bypass Retraining)
### A. Loading the Final Model (`main` branch)
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
"SRJ5035/sw_glu_sw_64_16_8_4_xpert_gpt",
revision="main",
trust_remote_code=True
).eval()
tokenizer = AutoTokenizer.from_pretrained(
"SRJ5035/sw_glu_sw_64_16_8_4_xpert_gpt",
revision="main"
)
```
### B. Loading an Intermediate Milestone (e.g. `chck_5M`)
```python
model_5m = AutoModelForCausalLM.from_pretrained(
"SRJ5035/sw_glu_sw_64_16_8_4_xpert_gpt",
revision="chck_5M",
trust_remote_code=True
).eval()
```
|