stamsam/Gemma-4-12B-a4B-MoE

Gemma 4 12B-a4B-MoE

⚠️ Works, but is not as efficient as the stock model with MTP.
The stock gemma-4-12B-it-oQ4 is faster (~23-30 t/s vs ~17.7 t/s), smaller (6.9 GB vs 11 GB), and doesn't need a custom bridge. This MoE variant demonstrates the viability of sparse upcycling for Gemma 4 at the 12B scale, but the dense oQ4 is the better daily driver on this hardware.

A sparse Mixture-of-Experts upcycle of Google's Gemma 4 12B-it, healed via QLoRA distillation.

This model takes the dense Gemma 4 12B instruction-tuned checkpoint and converts it into a sparse MoE architecture (4 experts, top-2 routing, low-rank experts) using the upcycling approach described in works like He et al. (2024) and ToMoE (Gao et al., 2025). The expert weights were then healed via sequence-level QLoRA distillation from a Gemma 4 31B teacher.

Architecture

Property	Value
Base model	Google Gemma 4 12B-it
Architecture	Sparse MoE (upcycled from dense)
Experts	4, top-2 routing per token
Expert type	Low-rank (rank 16, intermediate 1536)
Mode	Residual MoE (dense_keep=0.75)
Active params/token	~4B
Total stored	~12B
Quantization	4-bit, group size 64 (~7.57 bpw)
Size on disk	11 GB
MLX bridge	gemma4_moe_unified (custom, included in repo)

Performance

Metric	MoE (a4B)	Stock oQ4 (with MTP)
Size	11 GB	6.9 GB
Speed	~17.7 t/s	~23-30 t/s
Memory	11.3 GB	~7 GB

Quality is strong across coding, reasoning, agentic, and conceptual tasks (A-grade on all hard eval questions). The stock dense oQ4 is faster and smaller for daily use; this MoE variant demonstrates the viability of sparse upcycling for Gemma 4 at the 12B scale.

Quick Start

Prerequisites

pip install mlx-lm

Load and generate

from mlx_lm import load, generate

# The custom bridge (gemma4_moe_unified.py) is included in the repo
# and will be loaded automatically.
model, tokenizer = load("stamsam/Gemma-4-12B-a4B-MoE")

# Set residual mode (required for quality):
for i in range(48):
    model.language_model.model.layers[i].mlp.dense_keep = 0.75

# Generate
prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "Your prompt here"}],
    tokenize=False,
    add_generation_prompt=True
)
result = generate(model, tokenizer, prompt=prompt, max_tokens=200)
print(result)

That's it. The bridge is included in the repo and mlx-lm will load it automatically with trust_remote_code.

oMLX

omlx start
# Select "Gemma-4-12B-a4B-MoE" from the UI

Chat template

Uses the standard Gemma 4 instruction-tuned chat template with an empty thought channel by default. Set enable_thinking: true in chat_template_kwargs to enable the reasoning channel.

Status: Work in Progress

This is an experimental build. The pure sparse path (dense_keep=0.0) does not yet match the residual version in quality. Future work includes:

• Larger SFT dataset from the stock oQ4 teacher for a second healing pass
• Training directly on the pure sparse 4-bit base
• MTP/speculative decoding integration
• Multi-modal evaluation

Training Details

• Upcycle method: Dense → MoE conversion via weight factorization (router + low-rank experts initialized from dense MLP projections)
• Healing: QLoRA on the residual 4-bit base, 2000 iterations, rank 32, alpha 64, 12 layers
• Teacher: Gemma 4 31B (Ollama Cloud), sequence-level distillation
• SFT dataset: 77 rows covering coding, agentic, reasoning, general chat, and concise categories
• Quantization: Post-training 4-bit gs64 via mlx_lm convert
• Hardware: Apple Silicon Mac, 48 GB unified memory

Files

• model-*.safetensors — Quantized 4-bit weights
• config.json — Model configuration (model_type: gemma4_moe_unified)
• gemma4_moe_unified.py — Custom MLX bridge (loaded automatically)
• chat_template.jinja — Gemma 4 instruction chat template
• tokenizer.json / tokenizer_config.json — Gemma 4 tokenizer
• moe_manifest.json — MoE build metadata
• REFERENCES.md — Papers and resources that informed this build

Bridge

The custom MLX bridge (gemma4_moe_unified) is included in this repository. It implements the PartialGatedResidualMoE layer with:

• 4 experts, top-2 routing via mx.gather_mm
• Low-rank expert factorization (rank 16, intermediate 1536)
• Residual dense path (dense_keep=0.75) for quality
• mx.stop_gradient on routing indices for training compatibility

mlx-lm will load it automatically when you call load() on this repo.

License

This model is derived from Google's Gemma 4 12B-it, which is governed by the Gemma Terms of Use. The upcycled weights and bridge code are released under the same terms.