Hitchhiker's Guide — Model Architecture Mastery

"Architecture is not what you add; it is what you understand well enough to subtract."

1. The Attention Equation and Its Variants — Precise Math

1.1 Standard Multi-Head Attention

$$\text{MHA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

$$\text{head}_i = \text{Attention}(XW^Q_i, XW^K_i, XW^V_i)$$

$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$

where $M$ is the causal mask ($-\infty$ for future positions).

Memory during inference: For sequence length $n$, each layer stores KV cache of shape $(2, \text{batch}, n, h, d_k)$ = $2 \cdot b \cdot n \cdot h \cdot d_k$ floats. For LLaMA-7B: $b=1, n=4096, h=32, d_k=128$ → 32 layers × $2 × 4096 × 32 × 128 × 2$ bytes = 8 GB in FP16.

1.2 Grouped Query Attention

Reduce KV heads from $h$ to $g$ heads:

$$\text{GQA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

$$\text{head}i = \text{Attention}(XW^Q_i, XW^K{\lfloor i/r \rfloor}, XW^V_{\lfloor i/r \rfloor})$$

where $r = h/g$ is the repetition factor. When $g=h$: standard MHA. When $g=1$: MQA.

KV cache reduction: $g/h$ fraction of MHA's KV cache. LLaMA-3 uses $g=8, h=32$ → 4× reduction → 2 GB for same scenario.

Implementation trick: no need to actually repeat the KV tensors — use torch.repeat_interleave(k, h//g, dim=1) only during attention computation, not in cache.

1.3 Rotary Position Embeddings (RoPE)

RoPE encodes position $m$ by rotating query/key vectors in 2D subspaces:

$$f_q(x_m, m) = \begin{pmatrix} \cos(m\theta_1) & -\sin(m\theta_1) \ \sin(m\theta_1) & \cos(m\theta_1) \end{pmatrix} \ldots \begin{pmatrix} x_1 \ x_2 \end{pmatrix}$$

where $\theta_i = 10000^{-2(i-1)/d}$ (same base as sinusoidal embeddings).

Key property: $\langle f_q(x_m), f_k(x_n) \rangle$ depends only on $x_m$, $x_n$, and $(m-n)$ — position enters only as relative offset. This is why RoPE can generalize to unseen sequence lengths.

Implementation (the complex number trick):

def apply_rotary_emb(xq, xk, freqs_cis):
    # xq: (batch, seq, n_heads, head_dim)
    # freqs_cis: (seq, head_dim//2) complex
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

1.4 ALiBi (Attention with Linear Biases)

Instead of position embeddings, add bias $-m \cdot |i - j|$ to attention logits:

$$A_{ij} = \frac{Q_i K_j^T}{\sqrt{d}} - m \cdot |i - j|$$

where $m$ is a per-head slope: $m_h = 2^{-8h/H}$ for head $h$ of $H$ total heads.

Key advantage: completely eliminates position embedding parameters. Generalizes to longer sequences at inference time (slopes penalize large distances). Disadvantage: slightly lower peak performance than RoPE on standard benchmarks.

1.5 SwiGLU vs Standard FFN

Standard FFN: $\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$
SwiGLU: $\text{FFN}(x) = (\text{Swish}(xW_1) \odot xW_2)W_3$

where $\text{Swish}(x) = x \cdot \sigma(x)$.

Parameter count: SwiGLU adds $W_2$ (same size as $W_1$) but reduces $d_{ff}$ from $4d$ to $\frac{2}{3} \times 4d \approx 2.67d$ to maintain FLOPs parity. LLaMA uses $d_{ff} = \frac{8}{3} d_{model}$ rounded to multiple of 256.

2. Vision Transformer Internals

2.1 Patch Embedding

An image of size $(H, W, C)$ is divided into $N = \frac{H}{P} \times \frac{W}{P}$ patches of size $(P, P, C)$.

Each patch is linearly projected to $d_{model}$: nn.Conv2d(C, d_model, kernel_size=P, stride=P) — this is equivalent to dividing into patches and projecting each independently, but implemented efficiently as a single strided convolution.

Sequence length: $N + 1$ (N patches + 1 CLS token). For ViT-L with $P=16$, $H=W=224$: $N = 196$.

2.2 Why ViT Attention Heads Learn to Segment

Each patch is a $(P \times P)$ region. The class token attends to all patches to aggregate a global representation. In DINO training, the self-supervised objective forces the model to be invariant to local augmentations while sensitive to semantic content.

Mechanistically: early layers learn local edge/texture features (similar to CNN filters), while deep layers learn to cluster semantically related patches (foreground vs background). The CLS token's attention weights in the last layer correspond to object foreground because:

1. The DINO loss forces the CLS token to represent the "semantic essence" of the full image
2. Background patches carry less semantic information and receive lower attention weight
3. Multiple views (global + local crops in DINO) force the model to find consistent semantic regions

This is not hand-designed — it emerges purely from the training objective.

3. Mixture of Experts Architecture

3.1 Routing Mechanism

Each token independently selects the top-K experts out of E total:

def route(x, W_gate):
    # x: (batch, seq, d_model)
    logits = x @ W_gate.T  # (batch, seq, num_experts)
    gates, indices = logits.topk(k=2, dim=-1)  # select top-2 experts
    gates = F.softmax(gates, dim=-1)  # normalize selected expert weights
    return gates, indices

Auxiliary load balancing loss: without it, all tokens collapse to 1-2 experts ("router collapse"). The auxiliary loss penalizes unequal expert utilization:

$$L_{\text{aux}} = \alpha \cdot E \cdot \sum_{i=1}^{E} f_i \cdot P_i$$

where $f_i$ = fraction of tokens assigned to expert $i$, $P_i$ = average routing probability to expert $i$.

3.2 MoE Impact on NPU Deployment

MoE models (Mixtral 8x7B, GPT-4-style) are problematic for NPUs because:

1. Sparse activation: different tokens activate different experts → no batching across tokens for expert computation
2. Expert parameter loading: each expert's weights must be available — 8 experts × 7B params = 56B total (even though only 14B active per token)
3. Load imbalance: if multiple tokens route to the same expert simultaneously, that expert becomes a bottleneck

NPU mitigation strategies:

• Expert pruning: drop low-frequency experts for a specific deployment domain
• Expert merging: average 8 experts into 4 (reduces params at slight accuracy cost)
• MoE → dense distillation: distill the sparse model into a single dense FFN

4. Mamba / SSM Architecture

4.1 The Core SSM Equation

Mamba models are discrete-time state space models:

$$h_t = \bar{A} h_{t-1} + \bar{B} x_t$$ $$y_t = C h_t$$

where $h_t \in \mathbb{R}^N$ is the hidden state, $x_t, y_t \in \mathbb{R}$ are input/output, and $\bar{A}, \bar{B}$ are discretized from continuous-time parameters.

Selective SSM (what makes Mamba work): $B$, $C$, and $\Delta$ (discretization timescale) are input-dependent:

$$\bar{B}t = f_B(x_t), \quad C_t = f_C(x_t), \quad \Delta_t = \text{softplus}(f\Delta(x_t))$$

This gives the model the ability to selectively forget or remember based on content — equivalent to attention's dynamic weighting but in $O(n)$ time.

4.2 Mamba vs Attention for NPU Deployment

5. Key Interview Numbers

6. Resources

• Attention is All You Need (Vaswani et al., 2017) — the foundation
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
• RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al., 2021)
• Train Short, Test Long: Attention with Linear Biases (ALiBi) (Press et al., 2021)
• GLU Variants Improve Transformer (Noam Shazeer, 2020) — SwiGLU origin
• An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale (Dosovitskiy et al., 2020)
• Emerging Properties in Self-Supervised Vision Transformers (DINO) (Caron et al., 2021)
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Gu & Dao, 2023)
• Mixtral of Experts (Mistral AI, 2024)