moncefem/memory-lora-gemma4

Memory-LoRA — a hypernetwork that writes repo-specific LoRA adapters for Gemma-4-E2B

Give it a code repository; it returns a LoRA adapter for that repository in one forward pass. No fine-tuning, no retrieval, and zero repository tokens at inference time — the knowledge lives in the weights.

repo ──► 6-view embedding (12288-d) ──► hypernetwork ──► LoRA adapter ──► frozen Gemma-4-E2B
          frozen Qwen3-Embedding-0.6B      750M params       rank 16, α 32

On repositories absent from the training corpus, the generated adapter makes correct answers roughly 200× more likely than the frozen model (−5.31 nats) and wins 9 of 9 benchmark family/repo combinations.

question (pallets/click, never trained on)	frozen Gemma-4-E2B	+ generated adapter
What testing framework does this use?	"Jest" ✗	pytest ✓
How is this project built and packaged?	"a Dockerfile" ✗	setuptools, setup.py ✓
What documentation tool?	"a Gantt chart" ✗	Sphinx ✓

Code, demo app, and full engineering write-up: https://github.com/moncifem/memory-lora-gemma4 · deep dive in docs/DEEP_DIVE.md.

---

What's in this repository

path	what it is
runs/all_lora/all_lora_best_cpt.pth	the model. Merged corpus (2,146 repos / 55,700 QA). cr_val 2.6811, cr_test 2.6266
runs/h200_run/head.best.pt	prose-QA-only variant, kept as a fallback
data/embeddings/*.parquet	precomputed 12288-d repo embeddings
data/qna/*.jsonl	training QA corpora
data/real_code2lora/	RepoPeftBench (Code2LoRA's benchmark)
memory_lora/ scripts/ app/ deploy/	code, mirrored from GitHub

Earlier sixview_v1 / sixview_v2 weights were removed: they predate the input-standardisation fix, and sixview_v2 measurably degrades the base model (see below). Their metrics.jsonl remain for provenance.

---

Using the hypernetwork

1. Get the pieces

git lfs install
git clone https://huggingface.co/moncefem/memory-lora-gemma4
cd memory-lora-gemma4
git lfs pull --include="runs/all_lora/all_lora_best_cpt.pth"

pip install torch transformers peft safetensors pyarrow numpy
python app/engine/fetch_base_model.py     # google/gemma-4-E2B, 10.25 GB

2. Embed a repository (6 views → 12288-d)

The head is conditioned on a specific representation: six views of the repo — call graph, architecture, git history, contracts/tests, conventions, ops — each embedded by a frozen Qwen3-Embedding-0.6B and concatenated. A different embedding will not work.

git clone --depth 80 https://github.com/pallets/click /tmp/click
python app/engine/build_embedding.py --repo /tmp/click --out /tmp/click.npy

3. Generate the adapter

python app/engine/generate_and_merge.py \
    --embedding /tmp/click.npy \
    --checkpoint runs/all_lora/all_lora_best_cpt.pth \
    --adapter-out /tmp/click-adapter \
    --no-merge

That writes a standard PEFT adapter. Use --merged-out DIR instead of --no-merge for a self-contained merged model (for vLLM).

4. Use it

from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoTokenizer

tok = AutoTokenizer.from_pretrained("models/gemma-4-E2B")
model = AutoModelForImageTextToText.from_pretrained("models/gemma-4-E2B")
model = PeftModel.from_pretrained(model, "/tmp/click-adapter")

prompt = "Q: What testing framework does this repository use?\nA:"
out = model.generate(**tok(prompt, return_tensors="pt"), max_new_tokens=32)
print(tok.decode(out[0], skip_special_tokens=True))

with model.disable_adapter(): gives the frozen baseline for an A/B — the comparison that actually matters.

Calling the head directly

import numpy as np, torch, sys
sys.path.insert(0, "app/engine")
from generate_and_merge import load_head

head, cfg, alpha = load_head("runs/all_lora/all_lora_best_cpt.pth")
emb = np.load("/tmp/click.npy").astype("float32")      # 12288-d, RAW
out = head(torch.from_numpy(emb).unsqueeze(0))
# out["A"][type] -> [1, r, in_features]
# out["B"][type] -> [1, out_features, r]

The head standardises its input internally using statistics stored in the checkpoint — pass raw embeddings, do not normalise them yourself.

Its update is Δ = (α/r)·(x·Aᵀ)·Bᵀ with A:[r,in], B:[out,r] — identically PEFT's LoRA convention, so the output drops straight into a standard adapter. One (A, B) pair per shape-qualified module type, shared across the transformer layers of that shape (205 target modules, 14 types).

---

Results

Cross-repo held-out loss versus the frozen base on identical data:

checkpoint	cr_val	cr_test	vs base	verdict
sixview_v2 (removed)	2.606	—	+0.198	worse than no adapter
h200_run/head.best.pt	2.7168	2.7155	−5.191	good
all_lora_best_cpt.pth	2.6811	2.6266	−5.246 / −5.310	best

Absolute losses are not comparable across rows — the eval sets differ. The delta against the frozen model is.

Benchmark on three unseen repositories, three task families each (FACT = the trained Q&A format, CODE = real source-line completion, TEXT = repo prose):

repo	FACT	CODE	TEXT	keyword accuracy
psf/requests	−10.31	−5.38	−2.83	0% → 83%
pallets/click	−11.74	−7.66	−3.26	25% → 75%
OpenLLM-France/AudioBench	−5.08	−3.81	−1.73	0% → 0%

100% win rate on all nine.

Limits, stated plainly

• It learns a repo's stack and conventions, not what the project does. On requests it answers "XML library" instead of HTTP. That is exact factual recall, which a rank-16 LoRA structurally cannot hold — retrieval covers it.
• AudioBench keyword accuracy stayed at 0%: projects whose identity is not inferable from structure transfer poorly.
• Base losses of 12–16 on short gold targets inflate the deltas. The keyword accuracy and the generations are the trustworthy evidence.

---

The failure worth knowing about

An earlier checkpoint reached a healthy-looking eval loss of 2.606 while being worse than applying no adapter at all — and worse than random noise of matched magnitude.

Cause: 64% of every repo embedding is a constant vector shared by all repositories (the frozen encoder's mean response to "source code"). It dominated the trunk, which collapsed to emitting essentially the same adapter for every repo.

stage	mean pairwise cosine across repos
input embedding	0.73
after trunk	0.978 ← discriminability destroyed
emitted adapter	0.96
input, centered	0.00 ← the signal was there all along

Fix: MemoryLoRAHead.fit_input_stats() standardises the conditioning input with training-set statistics, stored in the checkpoint so training and inference apply the same transform. Emitted-adapter cosine went 0.96 → 0.21 in 40 steps; the shipped model sits at 0.32.

Why it went unnoticed: training logged only the adapted loss. A number like 2.606 says nothing without the frozen-model baseline beside it. Every eval now reports delta_vs_baseline and diag/adapter_cosine.

Before trusting any checkpoint:

python app/scripts/diagnose_head.py --job <id> --checkpoints runs/<run>/head.best.pt

It scores against none, random noise of matched scale, and a zero-B control that must reproduce the baseline exactly. A head that cannot beat random has not learned the mapping.

---

Training

Reproduce or continue on a single GPU:

python3 deploy/h200/preflight.py          # validates before spending GPU time
GATE=1 bash deploy/h200/train_h200.sh     # ~20-min go/no-go
bash deploy/h200/train_h200.sh            # full run

preflight.py fails fast on git-LFS pointers masquerading as data, pre-standardised embeddings, a head without fit_input_stats, and missing eval splits — then auto-tunes the largest micro-batch that fits.

---

Credits

Reimplements and extends Code2LoRA (arXiv 2606.06492) — a static hypernetwork mapping a repository embedding to a LoRA adapter — retargeted to google/gemma-4-E2B, extended from single-view code completion to a six-view representation, and wrapped in a serving stack that speaks the OpenAI and Anthropic APIs. Training data includes RepoPeftBench from that work.

Base model: google/gemma-4-E2B · Encoder: Qwen/Qwen3-Embedding-0.6B