gemma-4-e2b-it-text-mlx-4bit

A text-only, PLE-safe, 4-bit MLX build of Gemma 4 E2B, derived from Google's official quantization-aware-trained (QAT) checkpoint. Purpose-built for on-device inference on iPhone/Mac via mlx-swift-lm / mlx-lm: the vision tower and audio encoder are removed at conversion, and the Per-Layer-Embedding (PLE) pathway is quantized according to Google's own mobile QAT scheme rather than uniformly.

Base	google/gemma-4-E2B-it-qat-q4_0-unquantized
Base commit	6befbaca7398925921802abd1f277b495b78b738
Base lineage	google/gemma-4-E2B-it → QAT pipeline → half-precision QAT weights
Architecture	gemma4_text (text-only; no vision_config, no audio_config)
Precision	4-bit affine, group 64, with PLE-safe exclusions (see below)
Bits per weight	4.559
Size on disk	2.49 GiB (2.67 GB), 4 files + tokenizer
Parameters	4.629 B (text stack only)

Why this exists

1. Text-only. Google publishes E2B as a single 10.2 GB multimodal shard — file-level exclusion of the towers is impossible against the upstream repo. Converting ourselves makes the download manifest text-only by construction.
2. PLE-safe. E2B's Per-Layer Embeddings feed each decoder layer through scalar-scaled projections. Quantizing that pathway uniformly at 4 bits measurably degrades the model (numbers below). Google's own mobile QAT scheme keeps those projections at 8-bit or full precision; this build mirrors that.
3. QAT-derived. Starting from Google's QAT weights rather than post-quantizing the plain bf16 checkpoint is, by a wide margin, the single largest quality factor measured.

What was removed

All vision and audio parameters are dropped at conversion. The model cannot accept image or audio input; the tokenizer and chat template retain their media special tokens, which are inert.

Dropped prefix	Tensors	Parameters
model.vision_tower.*	(part of 1,411)
model.audio_tower.*
model.embed_vision.*
model.embed_audio.*
Total dropped	1,411	0.4757 B (0.886 GiB bf16)
Total kept (text stack incl. PLE, embeddings)	540	4.6286 B (8.621 GiB bf16)

Size math: 10.209 GB source shard → 8.621 GiB text stack in bf16 → 2.49 GiB after PLE-safe 4-bit quantization. The Per-Layer-Embedding table alone (embed_tokens_per_layer, 262,144 × 8,960 = 2.349 B parameters) is 46% of the whole model and ~53% of this artifact; it is the reason a text-only 4-bit E2B lands near 2.5 GB rather than the ~1.5 GB one would estimate from "2B-class model at 4 bits".

Quantization recipe

Derived directly from google/gemma-4-E2B-it-qat-mobile-ct's quantization_config (Google's own mobile QAT scheme), read off its config_groups and ignore list:

Module	Google mobile QAT	This build
model.per_layer_model_projection	ignored — full precision	bf16, not quantized
layers.N.per_layer_input_gate	8-bit	8-bit, group 64
layers.N.per_layer_projection	8-bit	8-bit, group 64
model.embed_tokens_per_layer (PLE table)	4-bit, group 256	4-bit, group 64 (finer ⇒ lower error)
layers.N.self_attn.{q,k,v,o}_proj	4-bit	4-bit, group 64
layers.N.mlp.{gate,up,down}_proj	4-bit (15 layers) / 2-bit (20 layers)	4-bit, group 64 (we do not take the 2-bit tier)
model.embed_tokens (tied to lm_head)	2-bit channel	4-bit, group 64 (more conservative)
all RMSNorms	not quantized	bf16

Measured quality

Perplexity on 16,384 tokens of long structured agent-context text (2,048-token chunks), all measured with the same mlx-lm build on the same machine:

Model	Perplexity
bf16 text-only reference (this build's QAT source, unquantized)	21.13
this build — text-only, PLE-safe 4-bit	20.51
ablation: same QAT source, naive uniform 4-bit (no PLE exclusions)	21.36
mlx-community/gemma-4-E2B-it-qat-4bit (QAT, PLE at 4-bit, MLP at 8-bit)	22.29
mlx-community/gemma-4-e2b-it-4bit (non-QAT base, uniform 4-bit)	106.85

Two findings worth stating plainly:

• QAT lineage dominates. Uniform 4-bit applied to the QAT checkpoint is fine (21.36); uniform 4-bit applied to the plain bf16 checkpoint is catastrophic (106.85). If you take one thing from this card: quantize the QAT release, not the base release.
• PLE-safety is a real but second-order win. Protecting the scalar-scaled PLE projections moves 21.36 → 20.51 and puts the 4-bit build slightly below its own bf16 reference — expected behaviour for QAT weights, which are trained to sit on a 4-bit grid.

Throughput on an Apple M5 Max (mlx-lm, greedy): ~197 tok/s decode on short prompts, ~84 tok/s decode with a 2,300-token prompt in context; prompt prefill ~29 tok/s cold on the first call. Generation was checked on 11 prompts spanning factual QA, arithmetic reasoning, strict formatting, JSON-only output, narrative, dialogue, code, summarization, negative constraints, and one 2,300-token structured agent context — coherent throughout, no repetition collapse.

Usage

from mlx_lm import load, generate

model, tokenizer = load("<this-repo>")
prompt = tokenizer.apply_chat_template(
    [{"role": "user", "content": "Why does a village need a granary?"}],
    add_generation_prompt=True, tokenize=False,
)
print(generate(model, tokenizer, prompt, max_tokens=256))

Requires mlx-lm with gemma4_text support (≥ 0.31). Text input only.

Reproducing

source : google/gemma-4-E2B-it-qat-q4_0-unquantized @ 6befbaca7398925921802abd1f277b495b78b738
convert: mlx 0.32.0 / mlx-lm 0.31.3
         drop model.{vision_tower,audio_tower,embed_vision,embed_audio,multi_modal_projector}.*
         rename model.language_model.* -> model.*
         config = source config.json "text_config", model_type=gemma4_text
         quantize group_size=64 bits=4 mode=affine with the PLE-safe predicate above

License and attribution

This is a derivative of Google's Gemma 4. It is distributed under the Apache License 2.0, the license of the base model — see the Gemma 4 license. Gemma is a family of open models built by Google DeepMind; the original model card, capability descriptions, evaluation results, and limitations are in the base repository and apply to this derivative.

Users of this derivative are directed to Google's responsible-use guidance for Gemma, which the base model card carries forward:

• Responsible Generative AI Toolkit
• Google's AI Principles

Known limitations of the base model — including that quality and diversity of training data bound its capabilities, that it can generate false or misleading text, and that it should not be relied on for factual accuracy without verification — carry over unchanged. This build additionally removes all image and audio capability: any multimodal claim in the base model card does not apply here.

Attribution: Gemma 4 © Google DeepMind. Technical report: arXiv:2607.02770.