Gemma-4-E4B-it — Weight-Space Programming Patch (0.002% Weights Edited, Zero Training)

This repository provides a direct weight-space edit for Google's gemma-4-E4B-it (8B total parameters / 4B active text parameters), inspired by Sparks of Weight-Space Programming: LLMs can natively read & write neural network weights (Takehi et al., MIT CSAIL).

Without a single step of gradient descent, backpropagation, or fine-tuning, modifying 64 MLP down_proj columns across Layers 38–41 (163,840 parameters = 0.00205% of the model) resolves two critical pathologies in Gemma 4 E4B's <|think|> mode:

1. Thinking + Tool-Calling Paralysis (0.0% → 75.0%): In baseline gemma-4-E4B-it, enabling <|think|> (enable_thinking=True) causes the model to either get trapped in 256+ token verbose markdown outlines (1. **Analyze the Request:** ...) or hallucinate temporal refusals ("Since 2026 is in the future...") instead of invoking declared tools like web_search.
2. Multi-Step Math Token Runaway (33.3% → 66.7%): Suppresses the meta-commentary outline loop ("Here's a thinking process to solve this problem: 1. **Identify...") so the model jumps directly into algebraic derivation ("Let N be...").

---

Before vs. After Benchmark Results

Full per-prompt transcripts, parsed tool calls, latencies, and column indices are included in benchmark_before_after.json.

---

Quickstart: Two Ways to Use

Option A: Load the Full Standalone Model Directly (transformers)

This repository contains the complete 16 GB bfloat16 model.safetensors, config.json, tokenizer.json, and chat_template.jinja, so you can load it directly via transformers or our 12 GB VRAM-optimized loader:

from transformers import AutoProcessor, AutoModelForMultimodalLM
import torch

model_id = "jstarks/gemma-4-E4B-it-weight-space-programming"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForMultimodalLM.from_pretrained(
    model_id, device_map="auto", torch_dtype=torch.bfloat16
)

Or on a 12 GB GPU (like an RTX 4070 Ti) using the included model_loader.py (which keeps the 42 active layers on GPU in 9.29 GB VRAM and offloads the 5.64 GB static Per-Layer Embedding table to CPU):

from model_loader import load_gemma4_e4b_text

model, tokenizer = load_gemma4_e4b_text("jstarks/gemma-4-E4B-it-weight-space-programming")

Option B: Apply the 641 KB Instant Patch to an Existing Checkpoint in <1 Second

If you already have google/gemma-4-E4B-it or unsloth/gemma-4-E4B-it downloaded locally, you can apply gemma4_e4b_weight_space_patch.safetensors (641 KB) in-memory without downloading 16 GB:

from apply_patch import apply_gemma4_weight_space_patch
from model_loader import load_gemma4_e4b_text

# 1. Load your existing Gemma-4-E4B-it checkpoint
model, tokenizer = load_gemma4_e4b_text("unsloth/gemma-4-E4B-it")

# 2. Apply the 64-column weight-space patch (163,840 parameters across Layers 38-41)
apply_gemma4_weight_space_patch(model, "gemma4_e4b_weight_space_patch.safetensors")

---

Key Mechanistic Findings on Gemma 4 E4B

1. The Layer 23 Boundary Hub (post_ffn_norm = 26.16, layer_scalar = 0.0654):
   • In gemma-4-E4B-it, Layers 0–23 have independent KV heads while Layers 24–41 (num_kv_shared_layers = 18) share KV states. At the boundary (Layer 23, a global full_attention layer), post_feedforward_layernorm has a mean weight of 26.16 (compared to 0.5–2.7 in other layers), but its residual output is immediately multiplied by a learned scalar buffer layer_scalar = 0.0654.
2. RMSNorm-Preconditioned Surgery in Layers 38–41:
   • Because each Gemma 4 decoder layer passes mlp(x) through post_feedforward_layernorm before adding to the residual stream, modifying down_proj columns in Layers 38–41 requires preconditioning the target residual direction d_target by max(|w_post|, 0.25)^(-1) so high-norm outlier coordinates in w_post do not distort the vocabulary projection.