Metadata-Version: 2.4
Name: ai-engram
Version: 0.6.0
Summary: Minimal, efficient covariance-based engram extraction for editing HuggingFace LLMs
Project-URL: Homepage, https://github.com/jeakwon/ai-engram
Project-URL: Repository, https://github.com/jeakwon/ai-engram
Author: Jeakwon Kim
License: MIT License
        
        Copyright (c) 2026 Jea Kwon
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: covariance,engram,llm,model-editing,pytorch,transformers,unlearning
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Requires-Dist: torch>=2.0
Requires-Dist: tqdm
Requires-Dist: transformers
Provides-Extra: dev
Requires-Dist: pytest; extra == 'dev'
Provides-Extra: docs
Requires-Dist: black; extra == 'docs'
Requires-Dist: mkdocs-material; extra == 'docs'
Requires-Dist: mkdocstrings[python]; extra == 'docs'
Description-Content-Type: text/markdown

# ai-engram

[![tests](https://github.com/jeakwon/ai-engram/actions/workflows/tests.yml/badge.svg)](https://github.com/jeakwon/ai-engram/actions/workflows/tests.yml)
[![PyPI](https://img.shields.io/pypi/v/ai-engram.svg)](https://pypi.org/project/ai-engram/)
[![Python](https://img.shields.io/pypi/pyversions/ai-engram.svg)](https://pypi.org/project/ai-engram/)
[![Docs](https://img.shields.io/badge/docs-jeakwon.github.io-7c4dff.svg)](https://jeakwon.github.io/ai-engram/)
[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](https://github.com/jeakwon/ai-engram/blob/main/LICENSE)

**Closed-form, covariance-based engram extraction for editing HuggingFace LLMs** — forward-only, no gradient descent.

An *engram* is the slice of a layer's weights attributable to a target set of inputs. `ai-engram` isolates it analytically:

```
W_engram = W · Σ_target · pinv(Σ_total)
```

`Σ_target` and `Σ_total` are input covariances over the **forget** set and the **reference** set. Subtracting it — `W ← W − α·W_engram` — removes that knowledge while keeping the rest: fast, training-free **unlearning / model editing**.

- **Closed-form** — one pseudo-inverse per layer; no optimization loop, no labels.
- **Forward-only** — covariances via forward pre-hooks; no backprop.
- **HF-native** — Llama, Mistral, Qwen, Gemma, Phi … and GPT-2 (`Conv1D`) out of the box.
- **Affine-correct** — automatic bias absorption for bias-bearing layers.
- **Tunable** — per-layer edit scaling is pluggable: the paper's `n/N` (default), relative weight-norm, effective rank, or your own.

> Statistics collection, closed-form **extraction**, and **editing** (`apply` / `edit`) are all here, and reproduce TOFU unlearning (see [Validation](#validation)).

## Install

```bash
pip install ai-engram
```

Pulls `torch`, `tqdm`, and `transformers` — HF LLMs and GPT-2 work out of the box. Distribution name `ai-engram`; **import name `engram`**.

📖 **Documentation: <https://jeakwon.github.io/ai-engram/>**

## Quickstart

Any `nn.Linear` (or GPT-2 `Conv1D`) model:

```python
import torch
from engram import EngramEditor, EditorConfig

editor = EngramEditor(model, EditorConfig())

target = editor.collect_statistics(forget_loader)   # Statistics: mean covariance + counts
total  = editor.collect_statistics(total_loader)    # over the reference set

edited = editor.edit(target, total, alpha=1.0)      # compute the engram and subtract it
# or split it: engram = editor.compute_engram_weights(target, total); editor.apply(engram, alpha=0.6)
```

### HuggingFace LLM (answer-token masked)

```python
from engram import EngramEditor, EditorConfig

editor = EngramEditor(model, EditorConfig())

batch_fn = lambda b: {"input_ids": b["input_ids"], "attention_mask": b["attention_mask"]}
mask_fn  = lambda b: b["labels"] != -100           # covariance over answer tokens only

g_forget = editor.collect_statistics(forget_loader, batch_fn=batch_fn, mask_fn=mask_fn)
g_total  = editor.collect_statistics(total_loader,  batch_fn=batch_fn, mask_fn=mask_fn)

edited = editor.edit(g_forget, g_total, alpha=0.6)   # default scaling = the paper's n/N
# selective per-layer strength, e.g. by relative weight-norm:
#   from engram import weight_norm, compose, count_ratio
#   edited = editor.edit(g_forget, g_total, alpha=1.0, scale=compose(count_ratio(1.0), weight_norm(1.0)))
```

Restrict the edit to specific modules with `target_modules` — the same convention
as LoRA/PEFT (`["down_proj"]` by name suffix, or a regex string), plus
`layers_to_transform` for decoder-layer indices. See the
[Quickstart guide](https://jeakwon.github.io/ai-engram/quickstart/) for details.

**Mixture-of-experts.** Answer-token masking reaches the experts automatically on
transformers&nbsp;<5; on transformers&nbsp;≥5 (fused experts) opt in to the
detachable `engram.moe` adapter — `EngramEditor(model, adapters=[FusedExpertAdapter()])` —
covering ~35 fused MoE architectures (Mixtral, Qwen2/3/3.5-MoE, DeepSeek-V3,
GLM4-MoE, MiniMax, Mistral4, OLMoE, Phi-MoE, …).

## How it works

| step | what | cost |
|---|---|---|
| 1. collect | forward pre-hooks accumulate the **mean** `xᵀx` + sample count per layer | one forward pass, no backward |
| 2. compute | projection `P = W · C_target · pinv(C_total)` | one pseudo-inverse per layer |
| 3. apply | `W ← W − α · f_l · P` with a pluggable per-layer scaling `f_l` | a single subtraction |

Efficient by construction — forward-only hooks, magnitude-bounded running-mean accumulation, CPU/GPU split (covariances on `storage_device`), a `float32` solve cast back to the model dtype, and answer-token masking. The per-layer edit weighting `f_l` is pluggable (`count_ratio` default = the paper's `n/N`, `weight_norm`, `effective_rank`, …). Handles `nn.Linear`, GPT-2 `Conv1D` (a transposed linear), and masked variants; full details in the [Guide](https://jeakwon.github.io/ai-engram/guide/).

### Configuration (`EditorConfig`)

| field | default | purpose |
|---|---|---|
| `storage_device` | model's device | where covariances are held; set `"cpu"` if the `D×D` matrices don't fit in VRAM (large models) |
| `absorb_bias` | `True` | absorb bias into the edit for bias-bearing layers |

## Validation

On **TOFU forget10** with `tofu_Llama-3.2-1B-Instruct`, the engram extraction reproduces the paper's 14-metric **Overall** within **~0.01**:

| condition | ai-engram | paper |
|---|---|---|
| gold (retain90) | 0.998 | 0.998 |
| plain (α=0.6) | 0.706 | 0.698 |
| adaptive-norm (α=1.0, p=1) | 0.817 | 0.818 |

Answer-token NLL confirms strong, *selective* forgetting — the forget set's NLL jumps ~16× while retain is preserved, and adaptive-norm beats plain on both axes. Runnable end-to-end in
[`tests/`](https://github.com/jeakwon/ai-engram/tree/main/tests) and
[`examples/`](https://github.com/jeakwon/ai-engram/tree/main/examples); see the [TOFU page](https://jeakwon.github.io/ai-engram/tofu/).

## API

- `collect_statistics(loader, target_modules=None, batch_fn=None, mask_fn=None, layers_to_transform=None) -> Statistics`
- `compute_engram_weights(target, total) -> EngramResult` · `apply(engram, *, alpha=1.0, scale=count_ratio(1.0)) -> Module` · `edit(target, total, *, alpha, scale)`
- scaling functions: `count_ratio` · `weight_norm` · `effective_rank` · `uniform` · `compose`
- `merge_statistics(*stats)` · `save_statistics(stats, path)` · `load_statistics(path)`

Full reference (auto-generated from docstrings): **[API docs](https://jeakwon.github.io/ai-engram/api/)**.

## License

[MIT](https://github.com/jeakwon/ai-engram/blob/main/LICENSE) © Jeakwon Kim
