File size: 8,932 Bytes
481fbb6
ef8744f
 
 
 
 
 
 
 
 
 
481fbb6
 
ef8744f
481fbb6
ef8744f
 
 
481fbb6
 
ef8744f
 
481fbb6
 
ef8744f
 
 
481fbb6
ef8744f
 
 
 
 
481fbb6
ef8744f
 
 
481fbb6
 
 
ef8744f
481fbb6
ef8744f
 
 
 
 
 
 
 
481fbb6
ef8744f
 
 
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
ef8744f
 
 
 
 
481fbb6
ef8744f
 
481fbb6
 
ef8744f
481fbb6
ef8744f
 
 
 
481fbb6
ef8744f
 
 
481fbb6
 
ef8744f
481fbb6
ef8744f
 
 
 
 
 
481fbb6
 
ef8744f
 
481fbb6
ef8744f
481fbb6
ef8744f
 
 
481fbb6
ef8744f
 
 
481fbb6
ef8744f
 
 
 
481fbb6
ef8744f
 
481fbb6
ef8744f
481fbb6
ef8744f
 
 
 
481fbb6
ef8744f
 
 
 
 
 
481fbb6
ef8744f
 
481fbb6
ef8744f
 
 
 
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
ef8744f
 
 
 
 
481fbb6
ef8744f
 
481fbb6
ef8744f
 
481fbb6
ef8744f
 
 
 
 
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
ef8744f
 
 
 
 
 
 
481fbb6
 
 
ef8744f
481fbb6
ef8744f
 
 
481fbb6
ef8744f
 
 
 
481fbb6
ef8744f
 
 
 
 
 
481fbb6
ef8744f
 
 
 
481fbb6
ef8744f
 
 
481fbb6
ef8744f
481fbb6
 
ef8744f
481fbb6
 
ef8744f
 
 
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
ef8744f
481fbb6
 
ef8744f
 
 
481fbb6
 
ef8744f
 
 
481fbb6
 
 
ef8744f
481fbb6
ef8744f
 
 
 
 
481fbb6
ef8744f
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
---
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)