Gemma 4 restructure

Here is the fully rewritten document.

---

# Retrofit Plan: Gemma 4 31B IT → Progressive Multi-Scale SWA with Block Attention Residuals

## 1. Executive Summary

This plan outlines a retrofit of **Gemma 4 31B Instruction-Tuned (IT)** into a progressive multi-scale sliding-window attention architecture with cross-block Attention Residuals (AttnRes). Starting from the aligned IT checkpoint preserves the expensive RL-carved reasoning structures, thinking-token pathways, and safety boundaries that define the model's performance.

We **empirically identify and permanently remove** the least important SWA layer in each of the 10 blocks. The resulting model is stabilized through a **Top-K distillation warm-start** that teaches the pruned architecture to replicate the exact output distribution of the original 31B checkpoint. AttnRes modules are inserted at global layer boundaries using **Zero-Init Uniform Softmax Initialization**, validated by Moonshot AI to prevent training volatility.

We do **not** introduce cross-layer KV cache sharing. The baseline Gemma 4 31B IT ships with `num_kv_shared_layers = 0`, and we respect this design choice. Each SWA layer maintains its own independent KV projection and cache.

---

## 2. Theoretical Foundation

### 2.1 Layer Redundancy & Hard Pruning
* **Evidence:** Fan et al. (2019) introduced LayerDrop [1], proving transformer networks have structural redundancy and can survive layer pruning. LLM-Drop (2024) [2] extended this to attention layers specifically.
* **Implication for this plan:** We identify redundancy empirically. For each of the 10 blocks, we ablate each of the 5 SWA layers individually and measure joint perplexity + instruction coherence degradation. The layer with the smallest impact is permanently removed. This yields 10 forward-pass evaluations per block—trivial compared to a gradual annealing phase. The surrounding layers and the new AttnRes mechanism learn to compensate during training, not through a structural fade.

### 2.2 Progressive Multi-Scale Window Attention
* **Evidence:** Multi-scale attention mechanisms that allocate varying window sizes to different layers have been shown to outperform uniform sliding-window attention in both perplexity and downstream reasoning tasks [3].
* **Implication for this plan:** Within each Gemma 4 block, we replace the uniform 1024-token SWA with a per-block local-to-global funnel (1024 $\to$ 2048 $\to$ 4096 $\to$ 8192). This is an independent design choice adapted to the block structure, not a direct application of Xu et al.'s depth-wise monotonic schedule. Each layer maintains its own independent KV cache capped at its assigned window size.

### 2.3 Attention Residuals (AttnRes) with Zero-Init Uniform Softmax
* **Evidence:** The Kimi Team (2026) introduced AttnRes [4], using learned attention weights to aggregate outputs from prior layers. Their ablation study explicitly validated the initialization strategy:
  > *"All pseudo-query vectors must be initialized to zero. This ensures that the initial attention weights are uniform across source layers, which reduces AttnRes to an equal-weight average at the start of training and prevents training volatility, as we validated empirically."*

  Their ablation results demonstrate that alternative formulations underperform the zero-init softmax baseline:

| Variant | Validation Loss |
| :--- | :--- |
| Baseline (PreNorm) | 1.766 |
| **AttnRes Full (zero-init softmax)** | **1.737** |
| w/ input-independent mixing | 1.749 |
| w/ sigmoid (gating) | 1.741 |
| w/ multihead (H=16) | 1.752 |
| w/o RMSNorm | 1.743 |
| SWA over depth (W=1+8) | 1.764 |

* **Implication for this plan:** We insert Block AttnRes modules at global layer boundaries using Moonshot's validated formulation. The pseudo-query vector is initialized to zero, producing a uniform attention distribution ($\alpha_i = 1/N$) at step one. This reduces the module to an equal-weight average of prior block outputs, providing a stable, well-behaved initialization regime from which training can smoothly depart. It does **not** reproduce the original sequential residual stream exactly; rather, it provides a known, stable starting point.

### 2.4 RoPE Relative Position Properties
* **Evidence:** RoPE (Su et al., 2021) encodes relative distance $(m - n)$ rotationally [5]. A query at position $T$ attending to keys in a sliding window experiences identical relative angles regardless of context scale.
* **Implication for this plan:** Each SWA layer applies its own causal mask over its independent KV cache. No position resets or rotation rollbacks are required.

---

## 3. Architectural Design

### 3.1 Baseline: Exact Gemma 4 31B IT Structure
The following are the exact values from the `google/gemma-4-31B-it` configuration:

* **vocab_size**: `262,144`
* **hidden_size** ($d_{\text{model}}$): `5,376`
* **intermediate_size** ($d_{\text{ffn}}$): `21,504`
* **num_hidden_layers**: `60`
* **tie_word_embeddings**: `true`
* **num_attention_heads** ($h_Q$): `32`
* **num_key_value_heads** ($h_{KV}$ for SWA): `16`
* **num_global_key_value_heads** ($h_{KV}$ for global): `4`
* **head_dim** (SWA): `256`
* **global_head_dim** (global): `512`
* **layer_types**: 60 layers in a fixed 5:1 block layout (50 `sliding_attention` layers alternating with 10 `full_attention` layers)

### 3.2 Hard Pruning via Empirical Ablation
For each block, we evaluate the 5 SWA layers individually:

1. Remove candidate layer $l$ (bypass its residual).
2. Measure perplexity on a mixed validation corpus and instruction coherence on a held-out complex prompt set.
3. Select the layer with the smallest joint degradation.

The 10 selected layers are **permanently deleted** before training begins. The pruned architecture is initialized immediately: 40 SWA layers (with progressive windows assigned) + 10 global layers + 10 AttnRes modules (zero-init uniform).

**Note:** While global pruning (selecting the 10 least important layers across all 50 SWA layers) could theoretically yield better redundancy exploitation, we restrict to 1 layer per block to maintain uniform block structure and simplify kernel scheduling. Global pruning is reserved as a fallback if per-block ablation reveals insufficient redundancy in shallow blocks.

### 3.3 Progressive Window Assignment (Independent KV)
After pruning, each block retains 4 SWA layers. We assign windows exponentially:

| Layer Position in Block | Window Size | Functional Role |
| :--- | :--- | :--- |
| 1st SWA | 1024 | Local syntax, entity tracking, recent coreference |
| 2nd SWA | 2048 | Sentence-level coherence, clause resolution |
| 3rd SWA | 4096 | Paragraph structure, local argument flow |
| 4th SWA | 8192 | Section-level dependencies, code block resolution |
| Global Layer | Full (256K) | Document integration + cross-block AttnRes routing |

Each layer maintains its **own independent KV cache** capped at its window size. The 1024-layer allocates a 1K cache; the 2048-layer a 2K cache; and so on. Because the layers do not share KV projections, each computes its own keys and values from the block input. This preserves per-layer query/key/value specialization and avoids the gradient-interference risks of cross-layer sharing.

**Kernel Strategy:** Standard causal SWA kernels (e.g., FlashAttention-2/3 with custom causal masks) are sufficient. Each layer invokes the kernel with its window-bound mask. No fused multi-scale kernel is required.

**Note on the 8192-window layer:** At 8K–16K training context, this layer initially behaves as "global lite," attending to nearly the full sequence. This is expected. As the curriculum expands to 32K and beyond, the 8192 boundary becomes semantically meaningful, and the layer specializes into section-level dependency tracking.

### 3.4 Block AttnRes with Zero-Init Uniform Softmax
We insert Block AttnRes modules at the 10 global layer boundaries. The input to Block $N+1$ is computed using Moonshot's validated formulation:

$$K_i = \text{RMSNorm}(V_i)$$
$$\text{logits}_{i} = \frac{w_{N+1}^T \cdot K_i}{\sqrt{d_{\text{model}}}}$$
$$\alpha_{i \to N+1} = \text{Softmax}(\text{logits})_i$$
$$H_{\text{input}, N+1} = \sum_{i=1}^{N} \alpha_{i \to N+1} \cdot V_i$$

Where:
* $V_i$ is the output of the global layer of Block $i$.
* $w_{N+1}$ is a pseudo-query vector of shape $[d_{\text{model}}]$.
* **Initialization:** $w_{N+1} = \mathbf{0}$. This guarantees uniform attention weights ($\alpha_i = 1/N$) at step one, reducing AttnRes to an equal-weight average.
* **RMSNorm:** Applied to each source vector $V_i$ before the dot product. Moonshot's ablation shows that removing RMSNorm degrades validation loss from 1.737 to 1.743.

### 3.5 Frozen Embeddings & Output Projection
Because `tie_word_embeddings = true`, the token embedding table is shared with the output projection layer (`lm_head`). Google invested significant RL-tuning resources to stabilize these weights, including a massive multilingual vocabulary. **Both the token embedding and the LM head are frozen throughout CPT.** This prevents erosion of the output distribution and preserves token-level reasoning pathways carved during post-training.

### 3.6 Final Block Structure

```
INPUT
  │
  ▼
┌─────────────────────────────────────────────┐
│ Block N                                     │
│  SWA-1024 (Independent KV)                  │
│  SWA-2048 (Independent KV)                  │
│  SWA-4096 (Independent KV)                  │
│  SWA-8192 (Independent KV)                  │
│  GLOBAL   (Full 256K Attention)             │
│  AttnRes  ◄── [Zero-Init Uniform Softmax] ──┼── From prior global outputs
└─────────────────────────────────────────────┘
  │
  ▼
OUTPUT to Block N+1
```

---

## 4. Data Strategy & Synthetic Pipeline

### 4.1 Top-K Distillation Warm-Start
Standard cross-entropy with label smoothing destroys the teacher's semantic geometry. For a well-trained model, the top 12 tokens usually capture 85–95% of the probability mass, and the critical signal is not just *what* the teacher chooses, but *what it rejects*—the near-miss tokens that are semantically adjacent but incorrect.

**Generation Protocol:**
- Generate responses from the original 31B IT on a dataset of **hard multi-turn problems** (complex technical Q&A, multi-step reasoning, code generation).
- Use **temperature = 1.0**. You want the true model distribution, not a sharpened or softened one.
- For **every token position**, save the indices of the top 12 tokens and their exact log-probabilities.

**Loss Formulation (Top-K + Tail Bucket):**
Let $p_{\text{teacher}} \in \mathbb{R}^V$ be the teacher's softmax distribution and $\mathcal{I} = \text{topK}(p_{\text{teacher}}, 12)$.

$$\mathcal{L} = -\sum_{i \in \mathcal{I}} p_{\text{teacher}}(i) \cdot \log q_{\text{student}}(i) \ - \ \left(1 - \sum_{i \in \mathcal{I}} p_{\text{teacher}}(i)\right) \cdot \log\left(1 - \sum_{i \in \mathcal{I}} q_{\text{student}}(i)\right)$$

The tail bucket aggregates all non-top-12 tokens into a single "OTHER" class. This preserves total probability mass and prevents the student from dumping arbitrary mass on tokens outside the teacher's top 12 with zero penalty.

### 4.2 Training Data Blend

| Phase | Tokens | Data Mix | Loss |
|-------|--------|----------|------|
| **Warm-start** | 0–20B | 100% Top-K distillation from hard multi-turn Q&A | Top-K + Tail |
| **Blend-in** | 20–50B | 50% Top-K distillation / 50% general long-context corpus | Top-K on distillation; standard CE on corpus |
| **General** | 50B+ | 75% long text / 15% reasoning traces / 10% formatting | Standard CE |

The warm-start phase teaches the pruned model to reproduce the original 31B's exact output geometry before exposing it to raw next-token prediction.

---

## 5. Training Protocol

### 5.1 Warm-Start Validation Gate
Before transitioning from Top-K distillation to general corpus training, a validation step is enforced:
1. After 1–2B tokens of warm-start training, evaluate the pruned model on a held-out set.
2. Measure the **KL divergence** between the student's output distribution and the teacher's top-K distribution on the same prompts.
3. Verify that the KL divergence is below a threshold (e.g., $< 0.1$ bits). This confirms the pruned model has successfully absorbed the teacher's output geometry before general training begins.

### 5.2 Optimization Hyperparameters
* **Starting Checkpoint:** Gemma 4 31B IT (with 10 SWA layers permanently removed)
* **Peak Learning Rate (Pretrained Weights):** $5 \times 10^{-7}$ (Conservative SFT-scale rate to protect RL-carved structure)
* **Peak Learning Rate (AttnRes):** $2.5 \times 10^{-6}$ (5× higher; AttnRes is randomly initialized and must learn routing from scratch)
* **Warmup:** 2% of total steps
* **Frozen Parameters:** Token embeddings and LM head receive zero learning rate.
* **Layer-wise LR Decay (Pretrained Weights Only):**
  * *Blocks 1–2 (Embedding / Shallow):* 0.4$\times$ peak LR (protect basic token syntax)
  * *Blocks 3–7 (Middle SWA):* 0.7$\times$ peak LR
  * *Blocks 8–10:* 1.0$\times$ peak LR

### 5.3 Phase Schedule & Sequence-Length Curriculum

```
CPT Tokens: 0 ────────────── 50B ────────────── 500B
AttnRes:    [ Locked (zero-init uniform) ] ──► [ Unlocked / Learned ]
Context:    [   8K   ] ──► [ 16K ] ──► [ 32K ] ──► [ 128K/256K ]
Data:       [ Top-K Warm Start ] ──► [ General Corpus Mix ]
```

* **Phase 1: Stabilization (0–50B tokens)**
  * *Context Length:* Max 8K tokens for the first 20B; scaled to 16K for the remainder. Starting at 8K ensures the 8192-window layer receives meaningful boundary signal immediately.
  * *AttnRes:* Pseudo-query weights frozen at zero (uniform average).
  * *Data:* Top-K distillation warm-start transitioning to general corpus.
  * *Goal:* Teach the pruned model to reproduce the teacher's distribution. The dense gradients from Top-K distillation train the remaining SWA layers to compensate for the missing ones.
* **Phase 2: Integration (50B–500B tokens)**
  * *Context Length:* Ramp 32K → 128K/256K.
  * *AttnRes:* Pseudo-query weights $w_{N+1}$ and RMSNorm parameters unlocked and allowed to train.
  * *Data:* Full general corpus blend.
  * *Goal:* Let AttnRes learn cross-block routing as the context length forces the progressive windowing to specialize.

### 5.4 AttnRes Monitoring
Whether AttnRes actually learns meaningful cross-block routing is **unknown a priori** and must be determined at train time. The following diagnostics should be computed per-batch once AttnRes is unlocked:

| Metric | Purpose | Healthy Signal |
| :--- | :--- | :--- |
| **Pseudo-query L2 norm** | Whether the module is learning at all | Gradual increase from ~0 to stable value |
| **Attention entropy** $H(\alpha)$ | Whether specialization is occurring | Decrease from $\log(N)$ toward moderate values |
| **Source block distribution** | Which prior blocks are being queried | Structured patterns (e.g., Block 8 weights Block 3 and 7) |
| **Gradient norm ratio** $\|\nabla w\| / \|\nabla W_{\text{SWA}}\|$ | Whether AttnRes is receiving training signal | Comparable magnitude to SWA gradients |

If attention entropy remains near maximum after 10B tokens of Phase 2, increase the AttnRes LR multiplier or add a sparsity-inducing auxiliary loss. Exact thresholds and mitigation strategies are unknown and should be tuned empirically.

### 5.5 Benchmark Monitoring
Monitor a composite score (GSM8K, MATH, HumanEval average). If the composite drops by $>2\%$ relative to the original IT baseline, pause CPT to reduce the learning rate and increase the distillation ratio.

---

## 6. Implementation Roadmap

### Phase 1: Redundancy Mapping & Hard Pruning (Week 1)
1. Evaluate Gemma 4 31B IT to locate one redundant SWA layer per block via ablation.
2. Permanently delete the 10 selected layers.
3. Pre-test the pruned model on instruction-following tasks to confirm baseline coherence.

### Phase 2: Architecture Surgery & AttnRes Setup (Week 2)
1. Initialize the 10 Block AttnRes modules with **zero-init pseudo-queries** and RMSNorm scaling.
2. Configure per-layer causal window masking ([1024, 2048, 4096, 8192]) using standard SWA kernels.
3. Verify forward-pass activation parity.

### Phase 3: Continual Pre-Training (CPT) (Months 2–4)
1. Execute the two-phase training curriculum over 200B–500B tokens.
2. Monitor benchmark composite and AttnRes diagnostics every 5B tokens.

### Phase 4: Light Alignment Refresh (Month 5)
1. **SFT Polish:** 2B tokens of combined Gemini 3 formatted responses and Gemma 4 reasoning chains.
2. **DPO Step:** 100K preference pairs to realign safety and conversational tone.

---

## 7. Technical Specifications

### 7.1 Exact Parameter Decomposition

Using the exact configuration values from the 31B checkpoint:

**1. Non-Decoder Parameters (Frozen Embeddings)**
Because `tie_word_embeddings = true`:
$$P_{\text{embed}} = 262{,}144 \times 5{,}376 = 1{,}409{,}286{,}144 \text{ parameters } (\approx 1.41\text{B})$$

**2. Inside a Single SWA Layer**
*   **`q_proj`**: $5{,}376 \times (32 \times 256) = 44{,}040{,}192$
*   **`k_proj`**: $5{,}376 \times (16 \times 256) = 22{,}020{,}096$
*   **`v_proj`**: $5{,}376 \times (16 \times 256) = 22{,}020{,}096$
*   **`o_proj`**: $(32 \times 256) \times 5{,}376 = 44{,}040{,}192$
*   **Attention Subtotal:** $132{,}120{,}576$

*   **`gate_proj`**: $5{,}376 \times 21{,}504 = 115{,}605{,}504$
*   **`up_proj`**: $5{,}376 \times 21{,}504 = 115{,}605{,}504$
*   **`down_proj`**: $21{,}504 \times 5{,}376 = 115{,}605{,}504$
*   **FFN Subtotal:** $346{,}816{,}512$

$$\text{Total SWA Layer Mass} = 478{,}937{,}088 \text{ parameters } (\approx 478.9\text{M})$$

**3. Inside a Single Global Layer**
*   **`q_proj`**: $5{,}376 \times (32 \times 512) = 88{,}080{,}384$
*   **`k_proj`**: $5{,}376 \times (4 \times 512) = 11{,}010{,}048$
*   **`v_proj`**: $5{,}376 \times (4 \times 512) = 11{,}010{,}048$
*   **`o_proj`**: $(32 \times 512) \times 5{,}376 = 88{,}080{,}384$
*   **Attention Subtotal:** $198{,}180{,}864$

*   **FFN Subtotal:** $346{,}816{,}512$ (same intermediate size)

$$\text{Total Global Layer Mass} = 544{,}997{,}376 \text{ parameters } (\approx 545.0\text{M})$$

**4. AttnRes Module**
*   **Pseudo-query vector**: $5{,}376$
*   **RMSNorm scale**: $5{,}376$
*   **Total per module**: $10{,}752$

### 7.2 Parameter Budget

| Component | Original Gemma 4 31B | Retrofitted Model |
| :--- | :--- | :--- |
| Embeddings (frozen) | 1,409M | 1,409M |
| SWA Layers | 50 × 478.9M = **23,947M** | 40 × 478.9M = **19,157M** |
| Global Layers | 10 × 545.0M = **5,450M** | 10 × 545.0M = **5,450M** |
| AttnRes Modules | 0 | 10 × 0.01M = **~0.1M** |
| **Total Parameters** | **~30.81B** | **~26.02B** |
| **Active Trainable Parameters** | **~29.40B** | **~24.61B** |

### 7.3 Memory & Compute Analysis (At 256K Context)
* **Per-Layer SWA Cache Memory:** Each SWA layer's KV cache grows only with its window size, not sequence length. At 256K context, the 4 SWA layers per block consume a combined 15K tokens of KV cache (1K + 2K + 4K + 8K), versus 256K per layer for full global attention. The global layers retain the native Gemma 4 unified KV mechanics.
* **Overall VRAM Impact:** The retrofit maintains a similar memory footprint to the original. The primary savings come from the 10 deleted layers (~4.8B fewer parameters) rather than KV cache compression.

---

## 8. Risk Assessment & Mitigation

| Risk | Likelihood | Impact | Mitigation |
| :--- | :--- | :--- | :--- |
| **Catastrophic Alignment Drift** | Medium | High | Conservative peak LR ($5 \times 10^{-7}$), frozen embeddings, and Top-K distillation warm-start that preserves the teacher's output distribution. |
| **Post-Pruning Shock** | Medium | High | **Hard prune + Top-K warm-start**. Dense gradients from exact teacher distribution replication teach the remaining layers to compensate immediately. |
| **AttnRes Representation Collapse** | Low | Critical | **Zero-Init Uniform Softmax**. Moonshot's ablation validates that zero-initialized pseudo-queries prevent collapse by reducing to an equal-weight average at step one. |
| **SWA Triton Kernel Failures** | Low | Medium | Standard SWA kernels with custom causal masks are sufficient. No exotic fused kernel is required. |
| **AttnRes Fails to Specialize** | Unknown | High | Monitor entropy and gradient norms. If modules remain uniform after unlocking, increase AttnRes LR multiplier or add sparsity auxiliary loss. Exact thresholds unknown; determined at train time. |

---

## 9. References

1. **Fan, A., Grave, E., & Joulin, A. (2019).** Reducing Transformer Depth on Demand with Structured Dropout. *arXiv preprint arXiv:1909.11556.* (LayerDrop)
2. **LLM-Drop Authors (2024).** LLM-Drop: Dropping Layers in Large Language Models. *arXiv preprint.* (Attention layer redundancy)
3. **Xu, et al. (2025).** Multi-Scale Window Attention (MSWA). *arXiv preprint.* (General multi-scale attention principles)
4. **Kimi Team (2026).** Attention Residuals: Cross-Block Feature Preservation in Deep Transformers. *Technical Report.* (AttnRes concept; zero-init uniform softmax initialization; ablation data showing sigmoid and multihead variants underperform)
5. **Su, J., et al. (2021).** RoFormer: Enhanced Transformer with Rotary Position Embedding. *arXiv preprint arXiv:2104.09864.* (RoPE relative properties)
6. **Google DeepMind (2026).** Gemma 4 Technical Report. *Google AI.* (Interleaved global and SWA attention, 512 head dim global layer design, exact configuration values for 31B dense variant)