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.

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

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:

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):

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.

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`