---
license: gemma
base_model: google/gemma-4-e2b-it
tags:
- gemma
- moe
- grouped-experts
- pruning
---

# Gemma-4 E2B with grouped expert routing (K=96, d=0.50)

This model is a sparsified variant of `google/gemma-4-e2b-it` where each MLP's
FFN dimension is partitioned into **96 groups** via k-means on activation
profiles. At inference, the top **48 of 96 groups** fire per token (50% density),
allowing entire groups of neurons to be skipped (and their corresponding rows
of `up_proj`/`gate_proj` and columns of `down_proj` to be skipped from memory).

Group sizes (D_FFN / 96): 64 in early layers (D_FFN=6144), 128 in late layers
(D_FFN=12288). 64-element groups in bf16 = exactly 128B = one L2 cache line on
NVIDIA GPUs — **memory-bandwidth-aligned routing**.

## Training recipe

1. Train `Sw_gate_only_50` baseline: 5000 steps gate-only, density=0.50, lr=1e-4,
   tau anneal 1.0→0.01, int4 QAT, on `gemma4_e2b_it_bulk_50k.jsonl`.
2. Cluster activations: 4 sequences × 2048 tokens through baseline; k-means at
   K=96 per layer on |gate*up| activation profile.
3. Polish: install GroupedMaskedMLP (K=96, K_active=48), add LoRA r128 on
   `up_proj`/`down_proj`, train 1000 steps CE on `gemma4_e2b_it_final_50k.jsonl`.

## Results (eval on `gemma4_e2b_it_final_50k.jsonl`)

| Metric | Value | vs Base | vs polish_50_per_neuron |
|---|---|---|---|
| PPL | 2.95 | +1.31× | +0.85× |
| MMLU | 30.6 | +0.2 | +1.0 |
| HellaSwag | 54.0 | +2.6 | -1.8 |
| ARC-Challenge | 35.2 | -2.2 | +1.6 |
| ARC-Easy | 40.6 | +1.2 | -0.6 |
| BoolQ | 74.2 | -2.4 | +3.2 |
| **5-task avg** | **46.9** | **-0.1** | **+0.7** |
| World-knowledge probe (172q) | 58.1 | -2.9 | +5.2 |

## Usage

```bash
git lfs install
git clone https://huggingface.co/Cactus-Compute/gemma4-e2b-grouped-k96
cd $(basename Cactus-Compute/gemma4-e2b-grouped-k96)
pip install torch transformers
python inference_k96.py --prompt "The capital of France is"
```

The inference script loads the base Gemma-4 weights from HuggingFace, installs
the grouped routing wrappers (loading cluster assignments from `groups/`), then
applies int4 QAT + LoRA + the trained state dict. `verify_grouped_routing()`
asserts every layer is genuinely a `GroupedMaskedMLP` with K=96 before generation.

## Citation

Internal research (Anthropic Claude + Noah Cylich, Cactus Compute), 2026-04.