Skip to content

feat(cli): larql memit — batch fact editing (Phase C of RFC-0001)#5

Open
mikeumus wants to merge 3 commits into
mainfrom
feat/memit-command
Open

feat(cli): larql memit — batch fact editing (Phase C of RFC-0001)#5
mikeumus wants to merge 3 commits into
mainfrom
feat/memit-command

Conversation

@mikeumus

Copy link
Copy Markdown

Summary

Phase C of RFC-0001: wraps the existing covariance-MEMIT solver (`larql_inference::forward::memit::run_memit`) with a CLI, edits.json file format, automatic crown-layer discovery per edit, and a dense patch output format. Stacked on #4 (Phase B).

What this PR adds

Extended patch file format

  • Bumped version 1 → 2 with a `kind` field (defaults to `"rank_one"` — Phase B patches round-trip unchanged).
  • New `kind = "dense"` variant: flat row-major ΔW matrix. Larger on disk (~72 MB per Gemma 4 4B layer) but semantically exact — no SVD approximation.
  • `write_patch` / `read_patch` / `apply_patch` all dispatch on `kind`.
  • New `compute_dense()` helper.

New CLI: `larql memit`

  1. Reads `edits.json` (a list of `{label, src, new_token, layer?}` records).
  2. For each edit: tokenises `src`, resolves `target_token_id`, resolves crown layer (explicit or auto-scan).
  3. Calls the existing `run_memit` with a `Vec`, receives one `MemitResult` per affected layer.
  4. Serialises each layer's ΔW as a Dense patch into the output directory, plus a `manifest.json`.
  5. Prints the `apply-patch` invocation that installs the batch.

Usage

cat > edits.json <<EOF
[
  {\"label\":\"france-to-tokyo\",\"src\":\"Capital of France? A:\",
   \"new_token\":\" Tokyo\",\"layer\":27},
  {\"label\":\"germany-to-rome\",\"src\":\"Capital of Germany? A:\",
   \"new_token\":\" Rome\",\"layer\":27}
]
EOF

larql memit /path/to/gemma4 --edits edits.json --output patches/
larql apply-patch /path/to/gemma4 \\
    -p patches/memit_L27.lqpatch \\
    --prompt \"Capital of France? A:\"

Design choice: dense patches vs low-rank

MEMIT's output ΔW is mathematically rank ≤ N (number of edits at that layer) but isn't stored as a factorisation — the covariance-projected solve naturally produces a dense delta. We could SVD-factorise post-hoc to emit N rank-1 patches, but that adds a lossy approximation step and doubles the code surface for apply-patch.

Choosing "dense" keeps Phase C tight and preserves the exact MEMIT result on disk. File size is large (~72 MB per layer for Gemma 4 4B) but a full batch edit typically touches 1–3 layers, so total patch bundle is 72–216 MB — acceptable for an enterprise MLOps artefact.

If future Phase-D Python consumers want rank-1 factors, they can SVD the dense delta themselves — lossless within numerical precision for ΔW of rank ≤ N.

Known ceiling

Chapter 22 established that single-layer MEMIT with correlated keys (~60% cosine) lands ~3/5 concurrent targets. Users can distribute across crown layers manually via `layer` overrides in edits.json — MEMIT runs independently per layer group.

Leverage of existing code

`forward::memit::run_memit` already implements the full ROME/MEMIT algorithm in Rust (~400 LOC with covariance estimation via `forward::trace::estimate_ffn_covariance`, solve with Tikhonov regularisation). This PR is a CLI veneer on top — no duplication, no math rewrite.

Testing

  • `cargo check --package larql-inference` ✓
  • `cargo check --package larql-cli` ✓
  • Round-trip live testing pending.

Base branch

Targets `feat/edit-command` (#4). After #3 and #4 merge to main, rebase onto main.

Remaining phase

Phase D: `larql-python` bindings (PyO3) exposing `crown`, `edit`, `memit` to Python — lets Colab experiments become one-liner Rust invocations.

🤖 Generated with Claude Code

mikeumus and others added 3 commits April 17, 2026 16:26
Implements Phase A of RFC-0001 (#2): per-layer MLP ablation scan to find
the layer whose last-position MLP output is load-bearing for a given
(prompt, expected-token) pair.

Changes:
- crates/larql-inference/src/ffn/ablating.rs — new LastPositionAblatingFfn
  that wraps any FfnBackend and zeroes its output at the last-token row for
  one target layer. Thin wrapper, no math changes.
- crates/larql-cli/src/commands/extraction/crown_cmd.rs — new `larql crown`
  subcommand. Tokenises the prompt, runs a baseline forward pass, then
  iterates layers in [start..=end] running predict_with_ffn against the
  ablating backend, reports per-layer Δ in expected-token probability and
  picks the layer whose ablation causes the top prediction to flip with the
  largest suppression magnitude.

Methodology matches Phase 125c of Divinci-AI/server
notebooks/CHAPTER_17_CORONATION.md — on Gemma 4 4B, ablating L27 MLP on
"Capital of France? A:" makes the top prediction flip from " Paris" to
"France" (the country token). The command outputs JSON (optional --json)
so downstream commands (edit, memit) can consume the crown_layer field.

Compile-checked with `cargo check --package larql-cli`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… RFC-0001)

Implements Phase B of RFC-0001 (#2): single-fact rank-1 editor with
portable patch file format. Builds on Phase A's LastPositionAblatingFfn
(#3) and adds the symmetric LastPositionInjectingFfn for scale search.

### New library module: `larql-inference/src/edit.rs`
- `EditPatch` struct (serializable via serde)
- `compute_rank1(k, d, scale, layer, provenance) -> EditPatch`
- `write_patch(path, &patch)` / `read_patch(path) -> EditPatch` with a
  simple binary format: LQPATCH magic + JSON meta + little-endian f32
  vectors for d and k_norm. ~55 KB for Gemma 4 4B.
- `apply_patch(&mut ModelWeights, &EditPatch)`: installs the rank-1
  outer product into `down_proj.weight` in place, handling both
  `[hidden, intermediate]` and `[intermediate, hidden]` layouts.

### New FFN wrapper: `larql-inference/src/ffn/injecting.rs`
- `LastPositionInjectingFfn` — adds a fixed delta vector to the inner
  backend's last-row output at one target layer. Symmetric to the
  ablating wrapper from PR #3. Used for auto-scale search.

### New CLI commands
- `larql edit <model> --src "..." --tgt "..." --new-token " Tokyo" --output f2t.lqpatch`
  Runs Phase A crown discovery (or accepts `--layer`), captures k at the
  crown layer for both prompts, computes d = W_down @ (k_tgt - k_src),
  linearly searches [0.5, 1, 1.5, 2, 2.5, 3, 4] for the minimum scale
  that flips the source's top-1 to --new-token, emits the patch.
- `larql apply-patch <model> --patch f2t.lqpatch --prompt "..."`
  Non-destructively installs one or more patches into the loaded
  weights, optionally runs a test prediction. Supports `--reverse`
  to subtract a patch (verifies reversibility).

### Supporting change
- Added `InferenceModel::weights_mut()` accessor so apply-patch can
  mutate the in-memory weight map without reloading.

Methodology validated in Python across Divinci-AI/server
notebooks/CHAPTER_20_HONEY.md (Phase 140c: France→Tokyo with 11/11
specificity at 0.9% weight perturbation) and CHAPTER_18_THE_EDIT.md
(Phase 130 scale search). The Rust port preserves the same math.

Compile-checked with `cargo check --package larql-cli`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wraps the existing covariance-MEMIT solver (larql_inference::forward::memit::
run_memit) with a CLI, an edits.json file format, and automatic crown-layer
discovery for each edit. Groups edits by crown layer, invokes the joint
least-squares solve, emits one dense `.lqpatch` per affected layer plus a
manifest.json. Phase C of RFC-0001 (#2), stacked on Phase B (#4).

### Extended patch file format (still backward compatible)
- Bumped patch version 1 → 2 with a `kind` field (defaults to "rank_one")
- New `kind = "dense"` variant carries a flat row-major ΔW matrix, needed
  because MEMIT's covariance-projected solve isn't natively a rank-1 outer
  product. Larger on disk (~72 MB per Gemma 4 4B layer) but semantically
  exact — no SVD approximation step.
- `write_patch`, `read_patch`, `apply_patch` all dispatch on kind. Phase B
  rank-1 patches continue to round-trip unchanged.
- New `compute_dense()` helper builds a Dense patch from an Array2<f32>.

### New CLI: `larql memit`
- Reads edits.json (list of {label, src, new_token, layer?} records).
- For each edit: tokenises src, resolves target_token_id, resolves crown
  layer (explicit or auto-scan).
- Calls `run_memit` with Vec<MemitFact>, receives one `MemitResult` per
  affected layer.
- Serialises each layer's ΔW as a Dense patch into the output directory,
  writes a manifest.json enumerating them.
- Prints the apply-patch command to install the batch.

### Usage

    cat > edits.json <<EOF
    [
      {"label":"france-to-tokyo","src":"Capital of France? A:",
       "new_token":" Tokyo","layer":27},
      {"label":"germany-to-rome","src":"Capital of Germany? A:",
       "new_token":" Rome","layer":27}
    ]
    EOF

    larql memit /path/to/gemma4 --edits edits.json --output patches/
    larql apply-patch /path/to/gemma4 \\
        -p patches/memit_L27.lqpatch \\
        --prompt "Capital of France? A:"

### Known ceiling
Chapter 22 established that single-layer MEMIT with correlated keys (~60%
cosine) lands ~3/5 concurrent targets. For 5+ correlated edits, users can
now distribute across multiple crown layers via `layer` overrides in
edits.json — MEMIT runs once per layer group.

Compile-checked with `cargo check --package larql-cli`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant