temp / README.md
Soham Jain
Update strict-small architecture files (sliding window [64, 16, 8, 4], ln3, ln_post_moe, no res3)
3a7a00c verified
|
Raw
History Blame Contribute Delete
2.92 kB
---
license: cc-by-nc-4.0
language:
- en
tags:
- babylm
- babylm-2026
- mixture-of-experts
- msit
- xpertgpt
- custom_code
- safetensors
library_name: transformers
pipeline_tag: text-generation
---
# XpertGPT (Sliding Window 16, 16, 4, 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 corrected LayerNorms, redundant residual removal, and **expanded sliding window sizes** across parallel expert channels.
---
## 1. Expert Sliding Window Layout
The four parallel MoE experts are configured with distinct sliding window attention constraints to capture varying context ranges (multi-scale sequence features):
* **Expert 1**: Window size `16` tokens
* **Expert 2**: Window size `16` tokens
* **Expert 3**: Window size `4` tokens
* **Expert 4**: Window size `4` tokens
---
## 2. 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 Residual 2 Feed-Forward Network addition inside every `MSITBranchBlock` (global block and all expert blocks).
* 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}})$.
---
## 3. Checkpoint Branch Layout
Checkpoints are saved as separate git branches (revisions) on this repository:
* **1M to 10M words**: Saved every 1M words (`chck_1M` through `chck_10M`).
* **10M to 100M words**: Saved every 10M words (`chck_10M` through `chck_100M`).
* **Final Model**: Saved under the `main` branch.
---
## 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/correct_small_sw_16_16_4_4_norm_residuls_xpert_strcit_small",
revision="main",
trust_remote_code=True
).eval()
tokenizer = AutoTokenizer.from_pretrained(
"SRJ5035/correct_small_sw_16_16_4_4_norm_residuls_xpert_strcit_small",
revision="main"
)
```
### B. Loading an Intermediate Milestone (e.g. `chck_5M`)
```python
model_5m = AutoModelForCausalLM.from_pretrained(
"SRJ5035/correct_small_sw_16_16_4_4_norm_residuls_xpert_strcit_small",
revision="chck_5M",
trust_remote_code=True
).eval()
```