anonym5035/neurips_trail_1_8experts

XpertGPT: 8-Expert MoE Architecture

XpertGPT is a specialized Mixture-of-Experts (MoE) architecture designed for efficient scaling. This specific configuration features 8 experts and is trained on the Wikitext-103-v1 dataset.

Model Configuration

• Vocab Size: 16384
• Training Block Size (Context): 1024 tokens
• Training Batch Size: 8
• Hidden Size (d_model): 256
• Thin Size (d_thin): 384
• Number of Layers: 6
• Number of Experts: 8
• Capacity Factor: 1.0
• Attention Type: Sliding Window Attention
• Window Sizes: 2 experts per window size: [64, 64, 16, 16, 8, 8, 4, 4]

Architecture Deep Dive

The architecture replaces standard Transformer blocks with MoEP-MSIT Blocks. Each block consists of:

1. Global Block: A full-dimension (d_model) branch running full sliding window attention.
2. Batch-Independent Expert Choice Router: A major enhancement in this iteration is strict batch-safety. The router computes affinities between all tokens and the 8 experts per sequence (B, T, num_blocks). The top-$k$ tokens for each expert are selected independently for each item in the batch based on k = capacity_factor * T / E. This ensures there is absolutely no cross-sequence leakage during routing, fully supporting training with BATCH_SIZE > 1.
3. Shrink Projection: The gated token representations are projected down to d_thin to save compute.
4. Thin Parallel Blocks (Experts): 8 separate MSIT Branch blocks (Attention + SwiGLU FFN), each operating on a specific sliding window size (e.g. 64, 16, 8, 4). Using torch.gather on the batched token sequences, tokens are processed cleanly.
5. Grow Projection: Output representations are projected back to d_model and accumulated back into the residual stream via torch.scatter_add_ using the original batch-aware routing gates.

Pre-Norm

The model employs strict Pre-Norm configuration: Layer normalization is applied to the residual stream before it branches into the projections, avoiding post-norms that disrupt training stability.

Flowchart

graph TD
    Input[Input Tokens<br/>B=8, T=1024] --> Embed[Token + Positional Embedding]
    
    subgraph MoEP-MSIT Layer
        Embed --> GlobalBlock[Global Branch Block<br/>Full Attention]
        GlobalBlock --> PreNorm[LayerNorm Pre-MoE]
        PreNorm --> Router[Batch-Safe Expert Choice Router<br/>top-k selection per sequence]
        
        PreNorm --> ShrinkProj[Down Projection<br/>d_model to d_thin]
        
        Router -- Selects Tokens for E1 --> E1[Expert 1: SW=64]
        Router -- Selects Tokens for E2 --> E2[Expert 2: SW=64]
        Router -- Selects Tokens for E3 --> E3[Expert 3: SW=16]
        Router -- Selects Tokens for E4 --> E4[Expert 4: SW=16]
        Router -- Selects Tokens for E5 --> E5[Expert 5: SW=8]
        Router -- Selects Tokens for E6 --> E6[Expert 6: SW=8]
        Router -- Selects Tokens for E7 --> E7[Expert 7: SW=4]
        Router -- Selects Tokens for E8 --> E8[Expert 8: SW=4]
        
        ShrinkProj -- torch.gather --> E1 & E2 & E3 & E4 & E5 & E6 & E7 & E8
        
        E1 & E2 & E3 & E4 & E5 & E6 & E7 & E8 --> GrowProj[Up Projection<br/>d_thin to d_model]
        GrowProj --> Gate[Multiply by Gate Scores]
        Gate -- torch.scatter_add_ --> Add[Residual Add]
        GlobalBlock --> Add
    end
    
    Add --> NextLayer[Next Layer or Final Norm]