concept-erasure-lab

An implementation of LEACE (least-squares concept erasure, Belrose et al. 2023, arXiv:2306.03819) built from scratch, checked against its closed-form guarantees and against the reference concept-erasure package, then exercised on SmolLM2-135M residual streams.

What is in it

• leace/eraser.py, closed-form LEACE eraser for a binary concept. Whitens the data covariance, takes the SVD of the whitened cross-covariance, and builds the rank-1 oblique projection that removes the concept direction, plus the affine shift that aligns class-conditional means. apply works on any device and dtype.
• tests/test_eraser.py, five tests on synthetic data with a known concept subspace, and one parity test against the reference package.
• experiments/, dataset templates and the runnable scripts for the two experiments and the figures.
• results/, committed metrics, plots, test output.

Run it

uv venv .venv && source .venv/bin/activate
uv pip install torch transformers numpy matplotlib scikit-learn pytest concept-erasure
cd experiments
python extract_acts.py    # ~1 min on a 4090, downloads SmolLM2-135M
python sweep.py
python behavior.py
python make_plots.py
python -m pytest ../tests -q

Closed-form checks

On synthetic data (n=4000, d=64, concept in an 8-dim subspace):

Check	Result
Cross-covariance with z after erasure, fit data	max
Cross-covariance on held-out data	~0.03, sampling noise order 1/sqrt(n)
Linear probe accuracy	0.99 before, at baseline after
Projection structure	idempotent, rank 1
Agreement with concept-erasure package	erased points differ by < 0.5 max abs

results/pytest.log has the run: 5 passed.

Layer sweep on SmolLM2-135M activations

480 prompts of the form {She|He} worked as a {occupation}... (40 occupations, 6 templates), residual stream read at the sentence-final token across all 30 layers, 70/30 train/test split.

Read from results/sweep.json:

• Before erasure, the linear probe reaches ~0.97-1.0 at most middle and late layers. The concept is strongly linearly encoded.
• After erasure, the linear probe drops to the majority baseline (~0.40-0.65, mean near baseline) at every layer.
• Mean edit magnitude is small: ~1-5% of the activation norm.

Falsifiable question: does a nonlinear probe recover the concept?

Yes. A 2-layer MLP probe trained on the erased activations recovers held-out accuracy 0.86-0.99 at most middle and late layers (e.g. 0.95 at layer 29, 0.98 at layer 26). The concept survives linear erasure in nonlinearly-encoded form, which is the documented scope of LEACE. At a few layers (4, 10, 22) recovery is near baseline, those layers also show weak pre-erasure linear signal, so the nonlinear probe finding is layer-dependent, not uniform.

Behavioral effect

A LEACE hook applied to every position at one layer while measuring next-token logits on "The {occupation} said that" (40 occupations):

From results/behavior.json:

• Baseline he-minus-she logit gap: +0.96 (mild he bias). Neutral-text perplexity 3.76.
• Erasing at layer 12 collapses the gap to -0.15, the targeted behavioral effect, at a perplexity cost of 5.26.
• Erasing at layers 16/20 amplifies the gap to +8.4/+12.7. At layer 28 it flips to -9.9. The oblique projection plus affine shift can push the residual toward the other class direction rather than neutralizing it.
• Every hooked layer raises neutral perplexity (5.3-7.0 vs 3.76). Collateral damage is real even for a rank-1 edit.

Limits

• 480 prompts over one pronoun pair is a small probe set. Per-layer accuracies carry a few points of noise (visible as jitter where the pre-erasure probe dips near baseline).
• Erasure was fit and measured at the last token position. Streaming or per-token erasure was not evaluated.
• The behavioral hook erases at one layer at a time. Multi-layer or embedding-space erasure was not tried.
• SmolLM2-135M is small. These numbers do not transfer to instruction-tuned or larger models.
• The nonlinear-recovery result shows erasure of a token-level concept is fragile against adaptive adversaries. LEACE's guarantee is linearity only.