paper/method_section.md

2. Architecture

Our model stacks $L{=}6$ identical MoEP-MSIT blocks between a shared token embedding layer and a final layer-normalised language-modelling head. Input tokens are embedded into $d_\text{model}{=}256$ dimensions; the embedding matrix is tied with the output projection to reduce parameter count. Each MoEP-MSIT block processes the token sequence through two sequential stages: a global dense block that establishes full-sequence context, followed by a sparse multi-scale expert stage that selectively refines token representations through parallel sliding-window pathways. Figure~1 illustrates the full block architecture.

2.1 Global Dense Block

The first stage of each MoEP-MSIT block operates on the full sequence at dimension $d_\text{model}$. It comprises a pre-LayerNorm multi-head causal self-attention layer with $h_g{=}4$ heads (head dimension $d_h{=}64$) followed by a gated feed-forward network. Self-attention uses Rotary Position Embeddings (RoPE) \cite{su2024roformer} without any window constraint, so every token can attend to all preceding positions. Given query and key vectors $\mathbf{q}_t, \mathbf{k}_s \in \mathbb{R}^{d_h}$ at positions $t$ and $s$, RoPE encodes relative position by applying a rotation matrix $R_\theta$:

$$\begin{array}{cccc}& \text{Attn}(t,s)=\frac{(R_{\theta ,t}{\mathbf{q}}_t{)}^{\mathrm{\top }}(R_{\theta ,s}{\mathbf{k}}_s)}{\sqrt{d_h}}& & \text{(1)}\end{array}\backslash text\{Attn\}(t,s)\; =\; \backslash frac\{(R\_\{\backslash theta,t\}\backslash ,\backslash mathbf\{q\}\_t)^\backslash top\; (R\_\{\backslash theta,s\}\backslash ,\backslash mathbf\{k\}\_s)\}\{\backslash sqrt\{d\_h\}\}\; \backslash tag\{1\}$$

where $R_{\theta,p}$ rotates consecutive dimension pairs by $p\theta_i$ with frequencies $\theta_i = 10000^{-2i/d_h}$. The output is projected through a linear layer and added to the residual stream. All feed-forward sub-layers use the SwiGLU activation \cite{shazeer2020glu}, defined as $\text{SwiGLU}(\mathbf{x}) = (\text{SiLU}(W_1 \mathbf{x}) \odot W_2 \mathbf{x}) W_3$ with hidden dimension $\lfloor 8d/3 \rceil_8$ (rounded to the nearest multiple of 8), which keeps the parameter count comparable to a standard $4d$-wide FFN while improving gradient flow.

2.2 Sparse Multi-Scale Expert Stage

After the global block produces a contextually enriched representation $\mathbf{X} \in \mathbb{R}^{B \times T \times d_\text{model}}$, the expert stage routes tokens through $E{=}4$ parallel expert branches.

Routing. A lightweight router computes per-token expert affinities. The router first normalises $\mathbf{X}$ via LayerNorm, then projects each token to $E$ logits through a bias-free linear layer $W_R \in \mathbb{R}^{d_\text{model} \times E}$:

$$\begin{array}{cccc}& {\mathbf{r}}_t=\text{softmax}\text{\hspace{0.17em}⁣}(W_R\text{LN}({\mathbf{x}}_t))\in {\mathbb{R}}^E& & \text{(2)}\end{array}\backslash mathbf\{r\}\_\{t\}\; =\; \backslash text\{softmax\}\backslash !\backslash left(W\_R\backslash ,\backslash text\{LN\}(\backslash mathbf\{x\}\_t)\backslash right)\; \backslash in\; \backslash mathbb\{R\}^\{E\}\; \backslash tag\{2\}$$

where $r_{t,i}$ is the affinity of token $t$ for expert $i$.

Expert-Choice Selection. We adopt Expert-Choice routing \cite{zhou2022mixture}, in which each expert independently selects its most relevant tokens rather than each token choosing an expert. Expert $i$ selects the top-$k$ tokens by their affinity $r_{t,i}$ along the token axis:

$$\begin{array}{cccc}& {\mathcal{S}}_i=\text{Top-}k\text{\hspace{0.17em}⁣}(\{r_{t,i}{\}}_{t=1}^N,k),k=\lfloor \frac{c\cdot N}{E}\rceil & & \text{(3)}\end{array}\backslash mathcal\{S\}\_i\; =\; \backslash text\{Top\}\backslash text\{-\}k\backslash !\backslash left(\backslash \{r\_\{t,i\}\backslash \}\_\{t=1\}^\{N\},\backslash ;\; k\backslash right),\; \backslash quad\; k\; =\; \backslash left\backslash lfloor\; \backslash frac\{c\; \backslash cdot\; N\}\{E\}\; \backslash right\backslash rceil\; \backslash tag\{3\}$$

where $N{=}B \cdot T$ is the total number of tokens in the batch and $c{=}2.0$ is the capacity factor. With $E{=}4$ experts and $c{=}2.0$, each expert processes $k{=}N/2$ tokens, and the union of all selections covers $2N$ token–expert slots in total. Because every expert selects exactly $k$ tokens, load balancing is guaranteed by construction, and no auxiliary balancing loss is required (see Appendix~A for an empirical comparison with token-choice top-1 routing and its associated load-balancing loss).

Dimension Projection. Before expert processing, the $d_\text{model}$-dimensional representations are projected to a higher expert dimension $d_\text{thin}{=}384$ via a shared bias-free linear projection $W_{\downarrow} \in \mathbb{R}^{d_\text{model} \times d_\text{thin}}$. This up-projection increases the representational capacity available to each expert without inflating the global model dimension.

MSIT Expert Branches. Each expert $i$ is a full MSIT Branch Block: a pre-LayerNorm sliding-window causal self-attention sub-layer with RoPE, followed by a pre-LayerNorm SwiGLU FFN and a post-block LayerNorm. All experts share the same dimensionality ($d_\text{thin}{=}384$, $h_e{=}6$ heads) but differ in their attention window size $w_i$:

$$\begin{array}{cccc}& [w_1,w_2,w_3,w_4]=[64,16,8,4]& & \text{(4)}\end{array}[w\_1,\backslash ,\; w\_2,\backslash ,\; w\_3,\backslash ,\; w\_4]\; =\; [64,\backslash ,\; 16,\backslash ,\; 8,\backslash ,\; 4]\; \backslash tag\{4\}$$

This multi-scale design is the core structural hypothesis of the architecture: different linguistic phenomena—discourse coherence, syntactic constituency, morphological agreement, and sub-word regularities—are best captured at different contextual granularities. By assigning each expert a dedicated window, the model can develop specialised representations at each scale without cross-scale interference (see Appendix~C for ablations over alternative window configurations).

Expert $i$ processes only its selected token subset $\mathcal{S}_i$. For each selected token $t \in \mathcal{S}_i$, the expert computes:

$$\begin{array}{cccc}& {\mathbf{y}}_{t,i}={\text{MSITBranch}}_i\text{\hspace{0.17em}⁣}(W_{\downarrow }{\mathbf{x}}_t;w_i)\in {\mathbb{R}}^{d_{\text{thin}}}& & \text{(5)}\end{array}\backslash mathbf\{y\}\_\{t,i\}\; =\; \backslash text\{MSITBranch\}\_i\backslash !\backslash left(W\_\{\backslash downarrow\}\backslash ,\backslash mathbf\{x\}\_t;\backslash ;\; w\_i\backslash right)\; \backslash in\; \backslash mathbb\{R\}^\{d\_\backslash text\{thin\}\}\; \backslash tag\{5\}$$

The expert output is projected back to $d_\text{model}$ via a shared projection $W_{\uparrow} \in \mathbb{R}^{d_\text{thin} \times d_\text{model}}$, gated by the routing weight, and scattered back into the full token sequence:

$$\begin{array}{cccc}& {\mathbf{o}}_t=\sum _{i:t\in {\mathcal{S}}_i}{r}_{t,i}\cdot W_{\uparrow }{\mathbf{y}}_{t,i}& & \text{(6)}\end{array}\backslash mathbf\{o\}\_t\; =\; \backslash sum\_\{i:\backslash ,\; t\; \backslash in\; \backslash mathcal\{S\}\_i\}\; r\_\{t,i\}\; \backslash cdot\; W\_\{\backslash uparrow\}\backslash ,\backslash mathbf\{y\}\_\{t,i\}\; \backslash tag\{6\}$$

The final block output combines the expert contributions with the residual from the global block, followed by a post-MoE LayerNorm:

$$\begin{array}{cccc}& {\mathbf{h}}_t=\text{LN}\text{\hspace{0.17em}⁣}({\mathbf{x}}_t+{\mathbf{o}}_t)& & \text{(7)}\end{array}\backslash mathbf\{h\}\_t\; =\; \backslash text\{LN\}\backslash !\backslash left(\backslash mathbf\{x\}\_t\; +\; \backslash mathbf\{o\}\_t\backslash right)\; \backslash tag\{7\}$$

2.3 Training Objective

The model is trained with the standard autoregressive cross-entropy loss over the token vocabulary:

Notably, because Expert-Choice routing guarantees balanced expert utilisation by construction (each expert selects exactly $k$ tokens), no auxiliary load-balancing loss is required. This is a direct advantage over token-choice routing, which typically requires an additional entropy or importance-weighted penalty to prevent expert collapse (see Appendix~A for an empirical comparison).

2.4 Tokenization

We train a custom Byte-Pair Encoding tokenizer with a vocabulary of $|\mathcal{V}|{=}16{,}384$ sub-word units on the training corpus. The final vocabulary entry is reserved as a mask token for optional masked-language-modelling pre-training phases.

2.5 Hyperparameters

All models are trained using AdamW ($\beta_1{=}0.9$, $\beta_2{=}0.999$) with weight decay $\lambda{=}0.1$, a peak learning rate of $3 \times 10^{-4}$ with cosine decay to $5%$ of peak, and 800 linear warm-up steps. The context length is $T{=}512$ tokens and the batch size is 16 sequences. Gradient norms are clipped at 1.0. We train for 8 epochs over the BabyLM Strict-Small corpus. All remaining hyperparameters are reported in Appendix~B.

2.6 Model Scale

The complete model has approximately 48.4M non-embedding parameters and 52.6M total parameters. The tied embedding accounts for $16{,}384 \times 256 \approx 4.2\text{M}$ parameters. Within each MoEP-MSIT block, the global dense block operates at $d_\text{model}{=}256$ and the four expert branches operate at $d_\text{thin}{=}384$, giving the experts higher per-token capacity than the global stream while keeping the total active parameter count per forward pass bounded by the Expert-Choice capacity factor.