File size: 3,020 Bytes
fab62cc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
---

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 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 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 `64` tokens
* **Expert 2**: Window size `16` tokens
* **Expert 3**: Window size `8` 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(

    "anonymous_user/correct_small_sw_64_16_8_4_norm_residuls_xpert_strcit_small",

    revision="main",

    trust_remote_code=True

).eval()



tokenizer = AutoTokenizer.from_pretrained(

    "anonymous_user/correct_small_sw_64_16_8_4_norm_residuls_xpert_strcit_small",

    revision="main"

)

```

### B. Loading an Intermediate Milestone (e.g. `chck_5M`)

```python

model_5m = AutoModelForCausalLM.from_pretrained(

    "anonymous_user/correct_small_sw_64_16_8_4_norm_residuls_xpert_strcit_small",

    revision="chck_5M",
    trust_remote_code=True

).eval()

```