Domain: Machine Learning & Artificial Intelligence

# Hybrid Efficiency Block — Forward Pass Computation

A hybrid efficiency block processes $T = 4$ tokens in a $d = 8$-dimensional embedding space. The input $X \in \mathbb{R}^{4 \times 8}$ is sampled from `numpy.random.RandomState(37).randn(4, 8)`.

The block integrates four components. The input $X$ is first normalized with RMSNorm producing $x_n$. The normalized $x_n$ feeds **only** the retention branch. The raw input $X$ feeds both the selective state-space branch and the mixture-of-experts branch.

**Numerical conventions:** All computations must be performed in float64 using NumPy semantics. Matrix multiplications follow standard row-by-column rules with explicit broadcasting as defined. Softmax is applied along the last axis unless otherwise noted. The sigmoid function is $\sigma(x) = 1/(1 + e^{-x})$. All intermediate values must be computed without rounding until the final answers, which are rounded to 4 decimal places.

---

## Component 1: RMSNorm

$$x_n[i, :] = \frac{X[i, :]}{\sqrt{\frac{1}{d}\sum_{j=1}^{d} X[i,j]^2 + \varepsilon}} \quad \text{where } d = 8,\; \varepsilon = 10^{-5}$$

No mean subtraction. Gain vector is $\mathbf{1}_8$.

---

## Component 2: Selective State-Space Model (on raw $X$)

A one-dimensional selective scan with continuous parameter $A = -1$. For each token $t = 0, \ldots, 3$:

1. **Projections:**
$$B_t = X_t \cdot W_B \in \mathbb{R}^{8}, \quad C_t = X_t \cdot W_C \in \mathbb{R}^{8}, \quad \Delta_t = \mathrm{softplus}(X_t \cdot W_\Delta) \in \mathbb{R}$$
where $\mathrm{softplus}(x) = \ln(1 + e^x)$ and $W_\Delta \in \mathbb{R}^{8 \times 1}$ so $\Delta_t$ is a scalar.

2. **Zero-order hold discretization:**
$$\bar{A}_t = e^{\Delta_t \cdot A} = e^{-\Delta_t} \in \mathbb{R}, \quad \bar{B}_t = (1 - e^{-\Delta_t}) \cdot B_t \in \mathbb{R}^{8}$$

3. **Recurrence** ($h_0 = \mathbf{0} \in \mathbb{R}^{8}$):
$$h_t = \bar{A}_t \cdot h_{t-1} + \bar{B}_t \in \mathbb{R}^{8}$$
where $\bar{A}_t$ is a scalar broadcasting across all 8 dimensions.

4. **Output** (element-wise product):
$$y_{\text{ssm}}[t] = C_t \odot h_t \in \mathbb{R}^{8}$$

---

## Component 3: Retention Mechanism (on $x_n$)

With $n_{\text{heads}} = 2$, $d_{\text{head}} = 4$, decay $\gamma = 0.9$.

1. **Projections:**
$$Q = x_n \cdot W_{Q_r} \in \mathbb{R}^{4 \times 8}, \quad K = x_n \cdot W_{K_r} \in \mathbb{R}^{4 \times 8}, \quad V = x_n \cdot W_{V_r} \in \mathbb{R}^{4 \times 8}$$

2. **Reshape** to $Q^{(h)}, K^{(h)}, V^{(h)} \in \mathbb{R}^{4 \times 4}$ for each head $h \in \{0, 1\}$, where head $h$ takes columns $[4h : 4h+4]$.

3. **Causal decay matrix** $D \in \mathbb{R}^{4 \times 4}$:
$$D_{ij} = \begin{cases} \gamma^{i-j} & \text{if } i \geq j \\ 0 & \text{otherwise} \end{cases}$$

4. **Retention scores and output per head:**
$$S^{(h)} = Q^{(h)} (K^{(h)})^\top \in \mathbb{R}^{4 \times 4}, \quad R^{(h)} = (S^{(h)} \odot D) \cdot V^{(h)} \in \mathbb{R}^{4 \times 4}$$

5. **Concatenate** heads and project:
$$y_{\text{ret}} = [R^{(0)} \| R^{(1)}] \cdot W_{O_r} \in \mathbb{R}^{4 \times 8}$$

6. **Retention entropy** $H_{\text{ret}}$: For each head $h$ and each query position $i$:
   - Compute decay-weighted scores: $z^{(h)}_{ij} = S^{(h)}_{ij} \cdot D_{ij}$ for $j \leq i$. Set $z^{(h)}_{ij} = -\infty$ for $j > i$ (causal mask).
   - Row-normalize via softmax: $W^{(h)}_{ij} = \frac{\exp(z^{(h)}_{ij})}{\sum_{k \leq i} \exp(z^{(h)}_{ik})}$
   - Entropy per position: $H^{(h)}_i = -\sum_{j \leq i} W^{(h)}_{ij} \ln W^{(h)}_{ij}$
   - Average: $H_{\text{ret}} = \frac{1}{n_{\text{heads}} \cdot T} \sum_{h,i} H^{(h)}_i$ where $T = 4$.

---

## Component 4: Sparse Mixture of Experts (on raw $X$)

With 2 experts.

1. **Router:** $\text{logits} = X \cdot W_{\text{router}} \in \mathbb{R}^{4 \times 2}$
2. **Routing probabilities:** $\text{probs} = \mathrm{softmax}(\text{logits}, \text{axis}=-1) \in \mathbb{R}^{4 \times 2}$
3. **Top-2 selection:** Since there are exactly 2 experts, both are always selected and weights equal the original softmax probabilities (no renormalization needed).
4. **Expert outputs and combination:**
$$y_{\text{moe}}[t] = \text{probs}[t, 0] \cdot (X_t \cdot W_{e_1}) + \text{probs}[t, 1] \cdot (X_t \cdot W_{e_2}) \in \mathbb{R}^{8}$$

---

## Component 5: Linear Attention Fusion

1. **Feature map** $\phi$ (ELU+1), applied element-wise:
$$\phi(x) = \begin{cases} x + 1 & \text{if } x \geq 0 \\ e^x & \text{if } x < 0 \end{cases}$$

2. **Sources:** Query from $y_{\text{ret}}$, Key and Value from $y_{\text{ssm}}$.

3. **Causal linear attention** for each position $i = 0, \ldots, 3$:
$$h_{\text{lin}}[i] = \frac{\sum_{j=0}^{i} \bigl(\phi(Q_i) \cdot \phi(K_j)\bigr) \cdot V_j}{\sum_{j=0}^{i} \phi(Q_i) \cdot \phi(K_j)}$$
where $Q_i = y_{\text{ret}}[i] \in \mathbb{R}^8$, $K_j = y_{\text{ssm}}[j] \in \mathbb{R}^8$, $V_j = y_{\text{ssm}}[j] \in \mathbb{R}^8$, and $\cdot$ denotes the dot product (scalar). The numerator is a weighted sum of $V_j$ vectors; the denominator is a scalar sum of dot products.

4. **Gated residual:**
$$h_{\text{base}} = \frac{y_{\text{ssm}} + y_{\text{ret}} + y_{\text{moe}}}{3}$$
$$\text{gate} = \sigma(h_{\text{base}} \cdot W_r) \quad \text{where } \sigma(x) = \frac{1}{1 + e^{-x}}$$
$$h = h_{\text{base}} + \text{gate} \odot h_{\text{lin}}$$

---

## Component 6: SwiGLU Feed-Forward Network

With $d_{\text{ff}} = 12$ and $\mathrm{Swish}(z) = z \cdot \sigma(z) = \frac{z}{1 + e^{-z}}$.

1. $h_n = \mathrm{RMSNorm}(h)$ (same formula as Component 1, applied to $h$).
2. $\text{gate\_val} = \mathrm{Swish}(h_n \cdot W) \in \mathbb{R}^{4 \times 12}$, $\quad \text{linear\_val} = h_n \cdot V \in \mathbb{R}^{4 \times 12}$
3. $\mathrm{SwiGLU}(h_n) = (\text{gate\_val} \odot \text{linear\_val}) \cdot W_{\text{down}} \in \mathbb{R}^{4 \times 8}$
4. Residual: $h_{\text{final}} = h + \mathrm{SwiGLU}(h_n) \in \mathbb{R}^{4 \times 8}$

---

## Task Heads

### Classification (Focal Loss)
- Logits: $\ell = h_{\text{final}} \cdot W_{\text{cls}} \in \mathbb{R}^{4 \times 4}$
- Class probabilities: $p = \mathrm{softmax}(\ell, \text{axis}=-1)$
- Label smoothing ($\varepsilon = 0.1$, $K = 4$ classes, labels $= [2, 0, 3, 1]$):
$$y_s[t, k] = \begin{cases} 1 - \varepsilon + \varepsilon/K & \text{if } k = \text{label}[t] \\ \varepsilon / K & \text{otherwise} \end{cases}$$
- Per-token cross-entropy: $\text{CE}[t] = -\sum_{k} y_s[t,k] \cdot \ln p[t,k]$
- Per-token focal weight using **unsmoothed** true class probability: $p_t = p[t, \text{label}[t]]$
- Focal loss: $L_{\text{cls}} = \frac{1}{T} \sum_t (1 - p_t)^\gamma \cdot \text{CE}[t]$ with $\gamma = 2$

### Regression (Huber Loss)
- Predictions: $\hat{y} = h_{\text{final}} \cdot W_{\text{reg}} \in \mathbb{R}^{4 \times 2}$
- Targets: $y = \text{RandomState}(71).\text{randn}(4, 2)$
- Element-wise error: $e_{ij} = y_{ij} - \hat{y}_{ij}$
$$\text{Huber}(e) = \begin{cases} \frac{1}{2} e^2 & \text{if } |e| \leq \delta \\ \delta(|e| - \frac{1}{2}\delta) & \text{otherwise} \end{cases} \quad \text{with } \delta = 1.0$$
- $L_{\text{reg}} = \frac{1}{4 \cdot 2} \sum_{i,j} \text{Huber}(e_{ij})$

### KL Divergence
- Token-averaged class distribution: $\bar{p} = \frac{1}{T} \sum_t p[t, :]$
- $D_{\text{KL}} = \sum_{k=1}^{K} \bar{p}_k \ln(K \cdot \bar{p}_k)$

---

## Total Loss

$$H = H_{\text{ret}}$$
$$E = L_{\text{cls}} + 0.5 \cdot L_{\text{reg}} + 0.01 \cdot H + 0.1 \cdot D_{\text{KL}}$$

---

## Weight Initialization

All 16 weight matrices drawn from a single `RandomState(97)` in the order below, each as `rng.randn(rows, cols) * std`:

| # | Weight | Shape | std | Role |
|---|--------|-------|-----|------|
| 1 | $W_B$ | $8 \times 8$ | $1/\sqrt{8}$ | SSM B projection |
| 2 | $W_C$ | $8 \times 8$ | $1/\sqrt{8}$ | SSM C projection |
| 3 | $W_\Delta$ | $8 \times 1$ | $1/\sqrt{8}$ | SSM step size |
| 4 | $W_{Q_r}$ | $8 \times 8$ | $1/\sqrt{8}$ | Retention Q |
| 5 | $W_{K_r}$ | $8 \times 8$ | $1/\sqrt{8}$ | Retention K |
| 6 | $W_{V_r}$ | $8 \times 8$ | $1/\sqrt{8}$ | Retention V |
| 7 | $W_{O_r}$ | $8 \times 8$ | $1/\sqrt{8}$ | Retention output |
| 8 | $W_{\text{router}}$ | $8 \times 2$ | $1/\sqrt{8}$ | MoE router |
| 9 | $W_{e_1}$ | $8 \times 8$ | $1/\sqrt{8}$ | MoE expert 1 |
| 10 | $W_{e_2}$ | $8 \times 8$ | $1/\sqrt{8}$ | MoE expert 2 |
| 11 | $W_r$ | $8 \times 8$ | $1/\sqrt{8}$ | Gated residual |
| 12 | $V$ | $8 \times 12$ | $1/\sqrt{8}$ | SwiGLU linear |
| 13 | $W$ | $8 \times 12$ | $1/\sqrt{8}$ | SwiGLU gate |
| 14 | $W_{\text{down}}$ | $12 \times 8$ | $1/\sqrt{12}$ | SwiGLU down |
| 15 | $W_{\text{cls}}$ | $8 \times 4$ | $1/\sqrt{8}$ | Classification |
| 16 | $W_{\text{reg}}$ | $8 \times 2$ | $1/\sqrt{8}$ | Regression |

All biases initialized to zero.

---

## Questions

**Q1:** What is $H$ (retention entropy)? Round to 4 decimal places.

**Q2:** What is $L_{\text{cls}}$ (focal loss)? Round to 4 decimal places.

**Q3:** What is $L_{\text{reg}}$ (Huber loss)? Round to 4 decimal places.

**Q4:** What is $D_{\text{KL}}$? Round to 4 decimal places.

**Q5:** What is $E$ (total loss)? Round to 4 decimal places.

Grading rubric:
criterion: Does the answer compute x_n[0,0] = -0.0638 and x_n[1,3] = 1.8106 from the RMSNorm step (values rounded to 4 decimal places)?
weight: 10
criterion: Does the answer compute y_ssm[0,0] = -0.3166 and y_ssm[2,3] = -0.1890 from the selective SSM recurrence (values rounded to 4 decimal places)?
weight: 10
criterion: Does the answer compute y_ret[0,0] = -0.3135 and y_ret[2,3] = 1.1332 from the retention mechanism (values rounded to 4 decimal places)?
weight: 10
criterion: Does the answer report H_ret = 0.2587 (rounded to 4 decimal places)?
weight: 15
criterion: Does the answer compute h_final[0,0] = 2.2156 and h_final[2,3] = 0.5309 after the SwiGLU residual (values rounded to 4 decimal places)?
weight: 10
criterion: Does the answer report L_cls (focal loss) = 2.1703 (rounded to 4 decimal places)?
weight: 15
criterion: Does the answer report L_reg (Huber loss) = 1.1017 (rounded to 4 decimal places)?
weight: 10
criterion: Does the answer report D_KL = 0.1544 and E (total loss) = 2.7392 (rounded to 4 decimal places)?
weight: 10
criterion: Does the answer include checksums sum(x_n) = 6.3585, sum(y_ssm) = -0.3665, sum(y_ret) = -0.3220, and sum(h_final) = 13.0751 (rounded to 4 decimal places)?
weight: 10

Answer:
## Step 1: RMSNorm

Compute $\text{RMS}[i] = \sqrt{\frac{1}{8}\sum_j X[i,j]^2 + 10^{-5}}$, then $x_n[i,:] = X[i,:] / \text{RMS}[i]$.

- Token 0: $\text{RMS} = 0.85387136$, $x_n[0] = [-0.06378433,\; 0.78970686,\; 0.40597102,\; -1.52288300,\; 1.77838484,\; 1.15921877,\; 0.32520226,\; -0.52535941]$
- Token 1: $\text{RMS} = 0.67845533$, $x_n[1] = [1.41787705,\; -1.21979827,\; 0.78805050,\; 1.81056310,\; 0.76584605,\; -0.09338097,\; -0.05128320,\; 0.06716072]$
- Token 2: $\text{RMS} = 1.17611472$, $x_n[2] = [1.23119379,\; 1.60996670,\; 0.34268112,\; 0.16361166,\; -0.25821720,\; -0.73102681,\; -0.22731360,\; -1.75930791]$
- Token 3: $\text{RMS} = 0.94388047$, $x_n[3] = [-0.75235293,\; 0.65850531,\; 1.29411117,\; -1.81850700,\; -0.41938624,\; 1.34950962,\; -0.01024295,\; -0.14624853]$

Therefore, $x_n[0,0] = \mathbf{-0.0638}$ and $x_n[1,3] = \mathbf{1.8106}$ (rounded to 4 d.p.).

Checksum: $\text{sum}(x_n) = \mathbf{6.3585}$ (rounded to 4 d.p.).

---

## Step 2: Selective SSM (on raw $X$)

For each token: project $B_t = X_t \cdot W_B$, $C_t = X_t \cdot W_C$, compute scalar $\Delta_t = \mathrm{softplus}(X_t \cdot W_\Delta)$, discretize $\bar{A}_t = e^{-\Delta_t}$, $\bar{B}_t = (1 - e^{-\Delta_t}) \cdot B_t$, update recurrence $h_t = \bar{A}_t \cdot h_{t-1} + \bar{B}_t$, output $y_{\text{ssm}}[t] = C_t \odot h_t$.

**Token 0:**
- $X_0 \cdot W_\Delta = -0.93833015$, $\Delta_0 = \mathrm{softplus}(-0.93833015) = \ln(1 + e^{-0.93833015}) = 0.33022467$
- $\bar{A}_0 = e^{-0.33022467} = 0.71876223$, $(1 - \bar{A}_0) = 0.28123777$
- $h_0 = \mathbf{0} \cdot 0.71876223 + \bar{B}_0 = \bar{B}_0$
- $y_{\text{ssm}}[0] = [-0.31657334,\; 0.06090617,\; -0.23396682,\; -0.14475338,\; 0.21397316,\; 0.05904054,\; -0.01836393,\; 0.00552339]$

**Token 1:**
- $X_1 \cdot W_\Delta = 0.41990715$, $\Delta_1 = 0.92498096$
- $\bar{A}_1 = e^{-0.92498096} = 0.39653897$
- $h_1 = 0.39653897 \cdot h_0 + \bar{B}_1$
- $y_{\text{ssm}}[1] = [0.01080373,\; 0.04961334,\; 0.07608341,\; -0.03520193,\; 0.14811167,\; 0.03234749,\; -0.51456410,\; 0.03720833]$

**Token 2:**
- $X_2 \cdot W_\Delta = 0.06514678$, $\Delta_2 = 0.72625099$
- $\bar{A}_2 = e^{-0.72625099} = 0.48371906$
- $h_2 = 0.48371906 \cdot h_1 + \bar{B}_2$
- $y_{\text{ssm}}[2] = [-0.03860931,\; -0.25908878,\; -0.36513994,\; -0.18896820,\; -0.30955984,\; -0.51076764,\; 0.35481431,\; -0.01951336]$

**Token 3:**
- $X_3 \cdot W_\Delta = -1.69920205$, $\Delta_3 = 0.16790933$
- $\bar{A}_3 = e^{-0.16790933} = 0.84543049$
- $h_3 = 0.84543049 \cdot h_2 + \bar{B}_3$
- $y_{\text{ssm}}[3] = [-0.12913567,\; 0.27411209,\; 0.44832762,\; 0.16763803,\; 0.05048704,\; -0.01371588,\; 0.82524138,\; -0.08276301]$

Therefore, $y_{\text{ssm}}[0,0] = \mathbf{-0.3166}$ and $y_{\text{ssm}}[2,3] = \mathbf{-0.1890}$ (rounded to 4 d.p.).

Checksum: $\text{sum}(y_{\text{ssm}}) = \mathbf{-0.3665}$ (rounded to 4 d.p.).

---

## Step 3: Retention

$Q = x_n \cdot W_{Q_r}$, $K = x_n \cdot W_{K_r}$, $V = x_n \cdot W_{V_r}$. Reshape to 2 heads of $d_{\text{head}}=4$.

Decay matrix $D$:
$$D = \begin{pmatrix} 1.0000 & 0 & 0 & 0 \\ 0.9000 & 1.0000 & 0 & 0 \\ 0.8100 & 0.9000 & 1.0000 & 0 \\ 0.7290 & 0.8100 & 0.9000 & 1.0000 \end{pmatrix}$$

**Retention entropy per head and position:**

Head 0: $H^{(0)}_0 = 0.00000000$, $H^{(0)}_1 = 0.00186584$, $H^{(0)}_2 = 0.58683032$, $H^{(0)}_3 = 0.95104853$
Head 1: $H^{(1)}_0 = 0.00000000$, $H^{(1)}_1 = 0.00425098$, $H^{(1)}_2 = 0.15385849$, $H^{(1)}_3 = 0.37143897$

For Head 0, Position 3 (the largest contributor): $z = [3.01066747,\; -3.74591348,\; 1.66717806,\; 3.24714433]$. After softmax: $W = [0.39543275,\; 0.00045997,\; 0.10318168,\; 0.50092560]$. Entropy $= -\sum W_j \ln W_j = 0.95104853$.

$y_{\text{ret}} = [R^{(0)} \| R^{(1)}] \cdot W_{O_r}$

Therefore, $y_{\text{ret}}[0,0] = \mathbf{-0.3135}$ and $y_{\text{ret}}[2,3] = \mathbf{1.1332}$ (rounded to 4 d.p.).

Checksum: $\text{sum}(y_{\text{ret}}) = \mathbf{-0.3220}$ (rounded to 4 d.p.).

$$H_{\text{ret}} = \frac{1}{8}(0.00000000 + 0.00186584 + 0.58683032 + 0.95104853 + 0.00000000 + 0.00425098 + 0.15385849 + 0.37143897) = \frac{2.06929313}{8} = 0.25866164$$

Therefore, $H_{\text{ret}} = \mathbf{0.2587}$ (rounded to 4 d.p.).

---

## Step 4: Mixture of Experts

Router probabilities (softmax of $X \cdot W_{\text{router}}$):
- Token 0: $[0.45508597,\; 0.54491403]$
- Token 1: $[0.73548553,\; 0.26451447]$
- Token 2: $[0.85534109,\; 0.14465891]$
- Token 3: $[0.35432635,\; 0.64567365]$

$y_{\text{moe}}[t] = \text{probs}[t,0] \cdot (X_t \cdot W_{e_1}) + \text{probs}[t,1] \cdot (X_t \cdot W_{e_2})$

Therefore, $y_{\text{moe}}[0,0] = \mathbf{0.3095}$ (rounded to 4 d.p.).

Checksum: $\text{sum}(y_{\text{moe}}) = \mathbf{2.4294}$ (rounded to 4 d.p.).

---

## Step 5: Linear Attention Fusion

$h_{\text{base}} = (y_{\text{ssm}} + y_{\text{ret}} + y_{\text{moe}}) / 3$

Therefore, $h_{\text{base}}[0,0] = (-0.31657334 + (-0.31349594) + 0.30952753) / 3 = -0.10684725$.

Causal linear attention with $\phi = \text{ELU}+1$, Query from $y_{\text{ret}}$, Key/Value from $y_{\text{ssm}}$:

For position 0: only $j=0$ contributes, so $h_{\text{lin}}[0] = V_0 = y_{\text{ssm}}[0]$.

Therefore, $h_{\text{lin}}[0,0] = -0.31657334$.

Gated residual: $\text{gate} = \sigma(h_{\text{base}} \cdot W_r)$, $h = h_{\text{base}} + \text{gate} \odot h_{\text{lin}}$.

Therefore, $h[0,0] = -0.19915215$.

---

## Step 6: SwiGLU FFN

$h_n = \text{RMSNorm}(h)$.

Therefore, $h_n[0,0] = -0.41589836$.

$\text{gate\_val} = \text{Swish}(h_n \cdot W)$, $\text{linear\_val} = h_n \cdot V$, $\text{SwiGLU} = (\text{gate\_val} \odot \text{linear\_val}) \cdot W_{\text{down}}$.

$h_{\text{final}} = h + \text{SwiGLU}(h_n)$.

Therefore, $h_{\text{final}}[0,0] = \mathbf{2.2156}$ and $h_{\text{final}}[2,3] = \mathbf{0.5309}$ (rounded to 4 d.p.).

Checksum: $\text{sum}(h_{\text{final}}) = \mathbf{13.0751}$ (rounded to 4 d.p.).

---

## Step 7: Focal Loss

Logits $= h_{\text{final}} \cdot W_{\text{cls}}$, $p = \text{softmax}(\text{logits})$.

Label smoothing: $y_s[t, \text{label}[t]] = 0.925$, $y_s[t, k \neq \text{label}[t]] = 0.025$.

**Token 0** (label=2):
- $p[0] = [0.16931738,\; 0.55573335,\; 0.19369955,\; 0.08124972]$
- $p_t = p[0,2] = 0.19369955$
- $\text{CE}[0] = -(0.025 \ln 0.16931738 + 0.025 \ln 0.55573335 + 0.925 \ln 0.19369955 + 0.025 \ln 0.08124972) = 1.64018037$
- $(1 - 0.19369955)^2 = 0.65012041$
- Focal contribution: $0.65012041 \times 1.64018037 = 1.06631474$

**Token 1** (label=0):
- $p[1] = [0.02615431,\; 0.06495379,\; 0.87720469,\; 0.03168721]$
- $p_t = 0.02615431$, $\text{CE} = 3.52838405$, $(1-p_t)^2 = 0.94837543$
- Focal contribution: $3.34623272$

**Token 2** (label=3):
- $p[2] = [0.00156533,\; 0.02789972,\; 0.86946387,\; 0.10107108]$
- $p_t = 0.10107108$, $\text{CE} = 2.37450335$, $(1-p_t)^2 = 0.80807321$
- Focal contribution: $1.91877254$

**Token 3** (label=1):
- $p[3] = [0.73411755,\; 0.06762737,\; 0.00147339,\; 0.19678169]$
- $p_t = 0.06762737$, $\text{CE} = 2.70308512$, $(1-p_t)^2 = 0.86931872$
- Focal contribution: $2.34984249$

$$L_{\text{cls}} = \frac{1}{4}(1.06631474 + 3.34623272 + 1.91877254 + 2.34984249) = \frac{8.68116249}{4} = 2.17029062$$

Therefore, $L_{\text{cls}} = \mathbf{2.1703}$ (rounded to 4 d.p.).

---

## Step 8: Huber Loss

$\hat{y} = h_{\text{final}} \cdot W_{\text{reg}}$, targets $y = \text{RandomState}(71).\text{randn}(4,2)$.

| Token | Dim | Prediction | Target | Error | Huber |
|-------|-----|-----------|--------|-------|-------|
| 0 | 0 | 0.18738558 | -0.43060315 | -0.61798873 | 0.19095504 |
| 0 | 1 | -0.67642805 | -1.19392809 | -0.51750004 | 0.13390314 |
| 1 | 0 | 1.34143401 | -0.44429899 | -1.78573301 | 1.28573301 |
| 1 | 1 | -0.51755654 | 0.48941221 | 1.00696875 | 0.50696875 |
| 2 | 0 | 3.25990536 | -0.45155715 | -3.71146251 | 3.21146251 |
| 2 | 1 | 0.51102252 | 0.58569596 | 0.07467344 | 0.00278806 |
| 3 | 0 | -2.62122138 | 1.17731982 | 3.79854120 | 3.29854120 |
| 3 | 1 | -0.35902921 | -0.96500906 | -0.60597984 | 0.18360579 |

$$L_{\text{reg}} = \frac{1}{8}(0.19095504 + 0.13390314 + 1.28573301 + 0.50696875 + 3.21146251 + 0.00278806 + 3.29854120 + 0.18360579) = \frac{8.81395749}{8} = 1.10174469$$

Therefore, $L_{\text{reg}} = \mathbf{1.1017}$ (rounded to 4 d.p.).

---

## Step 9: KL Divergence

$\bar{p} = \frac{1}{4}\sum_t p[t,:] = [0.23278864,\; 0.17905356,\; 0.48546038,\; 0.10269742]$

$$D_{\text{KL}} = \sum_k \bar{p}_k \ln(4 \cdot \bar{p}_k) = 0.23278864 \ln(0.93115457) + 0.17905356 \ln(0.71621423) + 0.48546038 \ln(1.94184151) + 0.10269742 \ln(0.41078970)$$
$$= -0.01660481 + (-0.05976377) + 0.32216935 + (-0.09136722) = 0.15443355$$

Therefore, $D_{\text{KL}} = \mathbf{0.1544}$ (rounded to 4 d.p.).

---

## Step 10: Total Loss

$$E = L_{\text{cls}} + 0.5 \cdot L_{\text{reg}} + 0.01 \cdot H + 0.1 \cdot D_{\text{KL}}$$
$$= 2.17029062 + 0.55087234 + 0.00258666 + 0.01544335 = 2.73919298$$

Therefore, $E = \mathbf{2.7392}$ (rounded to 4 d.p.).

---

## Final Answers

- **Q1:** $H = 0.2587$
- **Q2:** $L_{\text{cls}} = 2.1703$
- **Q3:** $L_{\text{reg}} = 1.1017$
- **Q4:** $D_{\text{KL}} = 0.1544$
- **Q5:** $E = 2.7392$

---

## Verification Code

The following Python script reproduces all values in this solution. To run it:

```
python3 -m venv venv
source venv/bin/activate        # Linux/macOS
venv\Scripts\activate            # Windows
pip install numpy
python verify.py
```

```python
import numpy as np

# ── INPUT ──
X = np.random.RandomState(37).randn(4, 8)

# ── WEIGHTS (all drawn from RandomState(97) in order) ──
rng = np.random.RandomState(97)

def make_weight(rows, cols, std):
    return rng.randn(rows, cols) * std

W_B      = make_weight(8, 8,  1 / np.sqrt(8))
W_C      = make_weight(8, 8,  1 / np.sqrt(8))
W_Delta  = make_weight(8, 1,  1 / np.sqrt(8))
W_Qr     = make_weight(8, 8,  1 / np.sqrt(8))
W_Kr     = make_weight(8, 8,  1 / np.sqrt(8))
W_Vr     = make_weight(8, 8,  1 / np.sqrt(8))
W_Or     = make_weight(8, 8,  1 / np.sqrt(8))
W_router = make_weight(8, 2,  1 / np.sqrt(8))
W_e1     = make_weight(8, 8,  1 / np.sqrt(8))
W_e2     = make_weight(8, 8,  1 / np.sqrt(8))
W_r      = make_weight(8, 8,  1 / np.sqrt(8))
V_swigl  = make_weight(8, 12, 1 / np.sqrt(8))
W_swigl  = make_weight(8, 12, 1 / np.sqrt(8))
W_down   = make_weight(12, 8, 1 / np.sqrt(12))
W_cls    = make_weight(8, 4,  1 / np.sqrt(8))
W_reg    = make_weight(8, 2,  1 / np.sqrt(8))

# ── STEP 1: RMSNorm ──
def rmsnorm(X):
    rms = np.sqrt(np.mean(X**2, axis=1, keepdims=True) + 1e-5)
    return X / rms, rms.flatten()

x_n, rms_vals = rmsnorm(X)

# ── STEP 2: Selective SSM ──
y_ssm = np.zeros((4, 8))
h_state = np.zeros(8)
for t in range(4):
    Bt = X[t] @ W_B
    Ct = X[t] @ W_C
    delta_raw = (X[t] @ W_Delta).item()
    delta_t = np.log(1 + np.exp(delta_raw))
    Abar = np.exp(-delta_t)
    Bbar = (1 - np.exp(-delta_t)) * Bt
    h_state = Abar * h_state + Bbar
    y_ssm[t] = Ct * h_state

# ── STEP 3: Retention ──
Q = x_n @ W_Qr; K = x_n @ W_Kr; V_ret = x_n @ W_Vr
gamma = 0.9
D = np.zeros((4, 4))
for i in range(4):
    for j in range(4):
        if i >= j: D[i, j] = gamma ** (i - j)
n_heads = 2; d_head = 4
R_all = []; H_ret_total = 0.0
for h in range(n_heads):
    Qh = Q[:, h*d_head:(h+1)*d_head]
    Kh = K[:, h*d_head:(h+1)*d_head]
    Vh = V_ret[:, h*d_head:(h+1)*d_head]
    Sh = Qh @ Kh.T
    Rh = (Sh * D) @ Vh
    R_all.append(Rh)
    for i in range(4):
        z = Sh[i, :i+1] * D[i, :i+1]
        z_max = np.max(z)
        z_exp = np.exp(z - z_max)
        W_h = z_exp / np.sum(z_exp)
        ent = -np.sum(W_h * np.log(W_h + 1e-30))
        H_ret_total += ent
R_concat = np.concatenate(R_all, axis=1)
y_ret = R_concat @ W_Or
H_ret = H_ret_total / (n_heads * 4)

# ── STEP 4: MoE ──
def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)

logits_moe = X @ W_router; probs_moe = softmax(logits_moe)
y_moe = np.zeros((4, 8))
for t in range(4):
    e1 = X[t] @ W_e1; e2 = X[t] @ W_e2
    y_moe[t] = probs_moe[t, 0] * e1 + probs_moe[t, 1] * e2

# ── STEP 5: Linear Attention Fusion ──
def elu_plus1(x):
    return np.where(x >= 0, x + 1, np.exp(x))

h_base = (y_ssm + y_ret + y_moe) / 3.0
Q_lin = y_ret; K_lin = y_ssm; V_lin = y_ssm
h_lin = np.zeros((4, 8))
for i in range(4):
    phi_Qi = elu_plus1(Q_lin[i])
    num = np.zeros(8); den = 0.0
    for j in range(i + 1):
        phi_Kj = elu_plus1(K_lin[j])
        w = np.dot(phi_Qi, phi_Kj)
        num += w * V_lin[j]; den += w
    h_lin[i] = num / den
gate = 1.0 / (1.0 + np.exp(-(h_base @ W_r)))
h = h_base + gate * h_lin

# ── STEP 6: SwiGLU FFN ──
h_n, _ = rmsnorm(h)
def swish(x): return x / (1.0 + np.exp(-x))
gate_val = swish(h_n @ W_swigl)
linear_val = h_n @ V_swigl
swiglu_out = (gate_val * linear_val) @ W_down
h_final = h + swiglu_out

# ── STEP 7: Focal Loss ──
logits_cls = h_final @ W_cls; p = softmax(logits_cls)
labels = [2, 0, 3, 1]; eps = 0.1; K_cls = 4; gamma_focal = 2
L_cls = 0.0
for t in range(4):
    ys = np.full(K_cls, eps / K_cls)
    ys[labels[t]] = 1 - eps + eps / K_cls
    ce = -np.sum(ys * np.log(p[t]))
    pt = p[t, labels[t]]
    focal_w = (1 - pt) ** gamma_focal
    L_cls += focal_w * ce
L_cls /= 4

# ── STEP 8: Huber Loss ──
y_pred = h_final @ W_reg
y_target = np.random.RandomState(71).randn(4, 2)
e = y_target - y_pred; delta_h = 1.0
huber_total = 0.0
for i in range(4):
    for j in range(2):
        ae = abs(e[i, j])
        if ae <= delta_h: hub = 0.5 * e[i, j]**2
        else: hub = delta_h * (ae - 0.5 * delta_h)
        huber_total += hub
L_reg = huber_total / 8

# ── STEP 9: KL Divergence ──
p_bar = np.mean(p, axis=0)
D_KL = np.sum(p_bar * np.log(K_cls * p_bar))

# ── STEP 10: Total Loss ──
E = L_cls + 0.5 * L_reg + 0.01 * H_ret + 0.1 * D_KL

# ── PRINT ALL CHECKPOINTS ──
print(f"x_n[0,0]     = {round(x_n[0,0], 4)}")
print(f"x_n[1,3]     = {round(x_n[1,3], 4)}")
print(f"sum(x_n)     = {round(np.sum(x_n), 4)}")
print(f"y_ssm[0,0]   = {round(y_ssm[0,0], 4)}")
print(f"y_ssm[2,3]   = {round(y_ssm[2,3], 4)}")
print(f"sum(y_ssm)   = {round(np.sum(y_ssm), 4)}")
print(f"y_ret[0,0]   = {round(y_ret[0,0], 4)}")
print(f"y_ret[2,3]   = {round(y_ret[2,3], 4)}")
print(f"sum(y_ret)   = {round(np.sum(y_ret), 4)}")
print(f"H_ret        = {round(H_ret, 4)}")
print(f"y_moe[0,0]   = {round(y_moe[0,0], 4)}")
print(f"sum(y_moe)   = {round(np.sum(y_moe), 4)}")
print(f"h_final[0,0] = {round(h_final[0,0], 4)}")
print(f"h_final[2,3] = {round(h_final[2,3], 4)}")
print(f"sum(h_final) = {round(np.sum(h_final), 4)}")
print(f"L_cls        = {round(L_cls, 4)}")
print(f"L_reg        = {round(L_reg, 4)}")
print(f"D_KL         = {round(D_KL, 4)}")
print(f"E            = {round(E, 4)}")
```