File size: 7,654 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 86 87 88 89 90 91 92 | # 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$:
$$
\text{Attn}(t,s) = \frac{(R_{\theta,t}\,\mathbf{q}_t)^\top (R_{\theta,s}\,\mathbf{k}_s)}{\sqrt{d_h}}
\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}$:
$$
\mathbf{r}_{t} = \text{softmax}\!\left(W_R\,\text{LN}(\mathbf{x}_t)\right) \in \mathbb{R}^{E}
\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:
$$
\mathcal{S}_i = \text{Top}\text{-}k\!\left(\{r_{t,i}\}_{t=1}^{N},\; k\right), \quad k = \left\lfloor \frac{c \cdot N}{E} \right\rceil
\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$:
$$
[w_1,\, w_2,\, w_3,\, w_4] = [64,\, 16,\, 8,\, 4]
\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:
$$
\mathbf{y}_{t,i} = \text{MSITBranch}_i\!\left(W_{\downarrow}\,\mathbf{x}_t;\; w_i\right) \in \mathbb{R}^{d_\text{thin}}
\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:
$$
\mathbf{o}_t = \sum_{i:\, t \in \mathcal{S}_i} r_{t,i} \cdot W_{\uparrow}\,\mathbf{y}_{t,i}
\tag{6}
$$
The final block output combines the expert contributions with the residual from the global block, followed by a post-MoE LayerNorm:
$$
\mathbf{h}_t = \text{LN}\!\left(\mathbf{x}_t + \mathbf{o}_t\right)
\tag{7}
$$
## 2.3 Training Objective
The model is trained with the standard autoregressive cross-entropy loss over the token vocabulary:
$$
\mathcal{L} = -\frac{1}{T}\sum_{t=1}^{T} \log\, p_\theta(x_t \mid x_{<t})
\tag{8}
$$
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.
|