---
license: apache-2.0
base_model: google/gemma-4-E2B
library_name: peft
tags:
  - hypernetwork
  - lora
  - code
  - repository-adaptation
  - gemma
pipeline_tag: text-generation
---

# 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`](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

```bash
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.**

```bash
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

```bash
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

```python
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

```python
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:

```bash
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:

```bash
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`](https://huggingface.co/google/gemma-4-E2B) ·
Encoder: [`Qwen/Qwen3-Embedding-0.6B`](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B)