SRJ5035/temp

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)

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)

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()