ai-engram 0.1.0

Surgical unlearning of classes, concepts, and facts in any PyTorch model

ai-engram

Surgical unlearning of classes, concepts, and facts in any PyTorch model.

ai-engram identifies the directions in a network's weights that store specific information — a class, a concept, a fact — and edits them out in closed form, with no gradient descent. It works on Linear, Conv1d, and Conv2d layers (grouped and depthwise convolutions included) and has been verified end-to-end against 24 pretrained vision, language, and diffusion models.

Installation

pip install ai-engram

Quickstart

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from ai_engram import EditorConfig, EngramEditor

model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 4))

forget_loader = DataLoader(TensorDataset(torch.randn(32, 8)), batch_size=8)
total_loader = DataLoader(TensorDataset(torch.randn(128, 8)), batch_size=8)

editor = EngramEditor(model, EditorConfig(device="cpu", damping_factor=1e-3))

forget_cov = editor.collect_statistics(forget_loader)
total_cov = editor.collect_statistics(total_loader)
edited = editor.edit(forget_cov, total_cov)

forget_cov carries the second-moment statistics of the data you want the model to forget; total_cov carries the statistics of the full training distribution. The edit removes the projection of the weights onto the forget subspace while preserving the rest.

This exact snippet is executed in CI on every commit and on every release, so a pip install always produces a runnable library.

Documentation

Full guides, API reference, and the 24-model compatibility report live at https://jeakwon.github.io/ai-engram/.

License

MIT — see also CHANGELOG.md.