LEACE: Perfect linear concept erasure in closed form

Nora Belrose; David Schneider-Joseph; Shauli Ravfogel; Ryan Cotterell; Edward Raff; Stella Biderman

Source-linked AI summary

LEACE: Perfect linear concept erasure in closed form

Nora Belrose, David Schneider-Joseph, Shauli Ravfogel, Ryan Cotterell, Edward Raff, Stella Biderman

arXiv:2306.03819v4Jun 6, 2023cs.LGcs.CLcs.CY

TL;DR

Concept erasure seeks to remove specified information from embeddings for fairness and interpretability. The paper introduces LEACE, a closed-form least-squares method for preventing linear concept detection, and extends it through concept scrubbing across network layers. Experiments apply these methods to gender information in BERT and POS information in language models, finding substantial POS-related perplexity increases while LEACE causes less collateral damage than SAL.

Problem
Concept erasure needs to remove specified concepts from embeddings to support fairness and interpretability while limiting models’ ability to use them.
Method
LEACE is a closed-form affine transformation that produces linearly guarded features with minimal mean-squared change, while concept scrubbing applies it sequentially across deep-network layers.
Results
LEACE is validated for gender-information removal in BERT, and POS scrubbing substantially increases perplexity across tested language models while causing less collateral damage than SAL.
Takeaways & Limitations
LEACE provides a surgical linear intervention for concept removal, supporting experiments that measure model reliance on POS information and reduce gender bias in embeddings.
Takeaways & Limitations
The paper calls for validation on narrower concepts and behavioral metrics, and notes that nonlinear methods may circumvent linear erasure.

Abstract

from arXiv · show

Concept erasure aims to remove specified features from an embedding. It can improve fairness (e.g. preventing a classifier from using gender or race) and interpretability (e.g. removing a concept to observe changes in model behavior). We introduce LEAst-squares Concept Erasure (LEACE), a closed-form method which provably prevents all linear classifiers from detecting a concept while changing the embedding as little as possible, as measured by a broad class of norms. We apply LEACE to large language models with a novel procedure called "concept scrubbing," which erases target concept information from every layer in the network. We demonstrate our method on two tasks: measuring the reliance of language models on part-of-speech information, and reducing gender bias in BERT embeddings. Code is available at https://github.com/EleutherAI/concept-erasure.

1 Introduction

Concept erasure limits a model’s use of specified concepts for fairness and interpretability. LEACE provides a theory-driven linear erasure method, while concept scrubbing extends erasure across deep-network layers and is evaluated on POS reliance and gender bias.

Concept erasure is motivated by fairness requirements and interpretability methods that seek to remove concepts from internal activations.The motivation concerns preventing protected attributes from affecting predictions and estimating conceptual effects on model behavior.
Concept erasure edits inputs or hidden states to minimize predictive V-information about a concept without finetuning or inspecting model parameters.Guardedness means classifiers in a chosen family cannot outperform a constant predictor at recovering the concept.
LEACE proves that linear guardedness is equivalent to every class having the same mean feature vector, then derives a least-squares affine erasure transformation.The transformation minimizes mean squared changes under all inner-product-induced norms, including Euclidean and Mahalanobis norms.
Concept scrubbing sequentially applies LEACE to activations at every layer of a deep network.The procedure is motivated by evidence that network features can be encoded in linear subspaces.
The paper evaluates LEACE by erasing gender information from BERT embeddings and using concept scrubbing to measure language-model reliance on part-of-speech information.These experiments address profession prediction after BERT intervention and POS information across language-model layers.

2 Preliminaries

The preliminaries formalize concept use through predictive information and define guardedness as the inability of a predictor family to improve beyond constant prediction. They then specialize the framework to convex losses and linear predictors.

Guardedness holds when the conditional feature distribution is among the worst possible for predicting labels with the chosen predictor and loss classes.Equivalently, prediction cannot improve beyond the trivially attainable loss of a constant predictor.
The trivially attainable loss is the lowest expected loss achievable by a constant predictor, with an achieving predictor called the trivial predictor when one exists.The notation Lτ and ητ identifies this loss and predictor in ambiguous cases.
Linear guardedness specializes guardedness to nonnegative losses convex in predictions and predictors of the form η(x)=b+Wx.Here b is a bias vector and W is a weight matrix.

3 Theoretical Results

The theoretical results establish several equivalent criteria for linear concept guarding, centered on equality of class-conditional means. They also connect this criterion to convex-loss optimality and zero feature-label covariance.

Theoretical Results: Linear guardedness, optimal trivial loss for convex losses, equal class-conditional means, zero cross-covariance, and statistical parity of linear classifiers are equivalent.The equal-means condition is E[X|Z=i]=E[X], while zero cross-covariance means ΣXZ is the zero matrix.
3.1 Equality of Class Centroids Implies Linear Guardedness: Equal class-conditional means prevent linear predictors from improving on the trivially attainable loss under convex losses.Jensen’s inequality lower-bounds a linear predictor’s loss by the loss of an appropriate constant predictor.
3.1 Equality of Class Centroids Implies Linear Guardedness: When all class centroids coincide, replacing each point by the global mean leaves linear classifiers unable to distinguish classes through feature means.In this setting, the data are indistinguishable to the linear predictor and W=0 is sufficient.
3.2 Linear Guardedness Implies Equality of Class Centroids: For losses with bounded, category-independent nonvanishing off-category derivatives, optimality of a constant linear predictor implies equal class-conditional means.The lemma applies to cross-entropy under nonzero class probabilities.
3.3 Linearly Guarded Labels Have Zero Covariance with the Features: Equal class-conditional means are equivalent to every feature component having zero covariance with every one-hot label component.This equivalence follows because each label component identifies a class with positive probability.
Theoretical Results: The equivalence of the first four conditions, together with the appendix result for statistical parity, completes the stated theoretical characterization.The paper identifies the statistical-parity condition as the fifth equivalent criterion.

4 Least-Squares Concept Erasure

LEACE characterizes affine transformations that remove all linear information about Z, then selects the unique least-squares-optimal transformation across a broad class of norms. The resulting eraser is an oblique projection that preserves more useful variance than orthogonal-projection alternatives while supporting categorical and real-valued concepts.

Characterizing guarded transformations: Linear guardedness after an affine transformation is equivalent to placing the columns of ΣXZ in the transformation’s null space.The bias does not affect cross-covariance; the condition reduces to PΣXZ = 0.
Least-squares optimization: LEACE minimizes mean squared edit distance subject to zero cross-covariance, with a bias b* = E[X] − P*E[X] in the uncentered case.The centered solution uses a positive-semidefinite matrix M, including potentially degenerate inner products.
Norm-independent optimality: A single erasure function minimizes all norms induced by positive-definite inner products, including Euclidean and Mahalanobis norms.These norms can serve as local quadratic approximations to divergence measures such as Kullback–Leibler divergence.
LEACE construction: The LEACE projection whitens the data, removes the subspace correlated with Z, and unwhitens it to restore the original covariance structure.The final operation subtracts the removed component from x, surgically eliminating linearly available information about Z.
Oblique projections: Unlike orthogonal projections, LEACE is generally oblique because guardedness constrains the projection’s null space rather than its range.Choosing the range freely lets LEACE minimize mean squared distance while preserving more original variance than an orthogonal projection.
Extension beyond categorical concepts: The same optimal affine eraser extends from categorical concepts to arbitrary Z in R^k under ordinary least-squares regression.The relevant equivalences and theorems do not require a categorical assumption on Z.

5 Evaluation

The evaluation tests LEACE for removing gender information while preserving profession prediction and reducing downstream gender bias, and examines concept erasure’s effect on language-model behavior. LEACE achieves near-random gender prediction with a small edit and substantially reduces bias, while POS erasure affects BERT most strongly at a later layer.

5.1 Intrinsic Evaluation: LEACE is the only evaluated method to achieve random gender-prediction accuracy with a small embedding edit, while RLACE comes close and INLP does not.LEACE is also around 2 orders of magnitude faster and does not require gradient-based optimization.
5.2 Downstream Fairness: 77.3% profession-prediction accuracy remains after projection, compared with 79.3% on original embeddings, indicating a small effect on the main task.The evaluation fits a logistic-regression profession classifier over projected [CLS] embeddings.
5.2 Downstream Fairness: 0.084 TPR gap after projection, down from 0.198, indicates substantially reduced biased behavior in the profession classifier.TPR-GAP measures the difference in true-positive rates between individuals with different protected attributes.
5.2 Downstream Fairness: 0.392 correlation between profession-level TPR gaps and women’s representation follows erasure, down from 0.867 before projection.The reported correlation measures the relation between classifier bias and the percentage of women in each profession.
5.3 Revisiting Amnesic Probing: POS erasure changes BERT language-model accuracy and loss only mildly through layer 8, with the largest drop at layer 11, whereas INLP has its maximum effect at layer 6.These results suggest the causal effect of POS information is concentrated in layer 11, unlike POS linear-probing optima at earlier layers.
5.3 Revisiting Amnesic Probing: A stacked classifier can recover removed information from multiclass logistic-regression softmax probabilities, although that setup is not linear.This caveat qualifies the linear-erasure evaluation.

6 Concept Scrubbing

Concept scrubbing sequentially applies LEACE across a transformer’s layers to remove linearly encoded POS information while limiting collateral damage. Across autoregressive language models, POS scrubbing sharply increases perplexity, whereas random-subspace erasure has little effect and SAL causes larger disruptions.

6 Concept Scrubbing: INLP’s collateral damage prevents applying it across multiple transformer layers without causing outputs to collapse into gibberish.The passage attributes this to deleting several dozen dimensions before linear guarding is achieved.
6 Concept Scrubbing: LEACE concept scrubbing minimizes collateral damage, adds little computational overhead, and computes covariance statistics in a streaming fashion without storing all hidden states.These properties make scrubbing practical across every layer of a deep network.
6 Concept Scrubbing: Concept scrubbing applies LEACE parameters sequentially, immediately scrubbing each layer’s hidden states before feeding them to the next layer.Sequential fitting accounts for the fact that interventions change downstream hidden-state distributions.
6 Concept Scrubbing: The evaluation removes linearly available POS information from every transformer-layer input in autoregressive language models, using Pythia and LLaMA model families.The models range from Pythia 160M to 12B and LLaMA 7B to 30B; interventions occur after normalization.
6.2 Results: Large perplexity increases occur across all tested models after POS scrubbing, while erasing a random subspace has little to no effect on language-model performance.The random-subspace condition controls for performance degradation caused by generic interventions.
6.2 Results: SAL produces significantly larger perplexity increases than LEACE for every tested model, indicating greater disruption to other useful embedding features.The authors therefore characterize LEACE interventions as more surgical than SAL interventions.

7 Limitations and Future Work

The paper identifies limits of concept erasure and concept scrubbing, including the unresolved scope of nonlinear adversaries and the need for broader behavioral validation. It also situates linear erasure among adversarial and algebraic information-removal methods and relates guardedness to fairness.

7 Limitations and Future Work: Concept scrubbing still needs validation on narrower concepts and with behavioral metrics that test whether network changes match expectations.The authors propose these experiments as necessary future work before incorporating scrubbing into pretraining or finetuning.
7 Limitations and Future Work: The paper focuses exclusively on linear concept erasure because of its simplicity and tractability, while nonlinear kernel-based erasure may not generalize across kernels.The authors conjecture that nondestructive erasure against a general nonlinear adversary may be intractable without detailed knowledge of the data-generating process.
7 Limitations and Future Work: Concept erasure is motivated as a post hoc, model-agnostic alternative to training-time methods, but unsatisfactory scrubbing could shift attention toward model-specific training interventions.
A Additional Related Work: Adversarial information-removal methods can leave protected-attribute information recoverable by newly trained adversaries.
Theory and Fairness Implications: For linear predictors, guardedness is equivalent to all class-conditional means being identical, and therefore to statistical parity for every linear predictor.The appendix states the equivalence between guardedness and equal class-conditional means, and separately proves the corresponding statistical-parity characterization.

E Derivation of LEACE

The derivation states the optimization problem underlying LEACE and presents its solution through Theorem 4.2. The theorem is supported by two independent proofs.

E Derivation of LEACE: Theorem 4.2 formulates LEACE using centered random vectors, covariance and cross-covariance matrices, a positive semidefinite inner product, and matrix pseudoinverses.
E Derivation of LEACE: Theorem 4.2 states that the objective has a closed-form solution.
E Derivation of LEACE: The appendix provides two independent proofs of Theorem 4.2.

E.1 Algebraic Proof

The algebraic proof decomposes the optimization into independent row problems in a suitable basis, then removes components correlated with the concept while preserving non-varying components.

E.1 Algebraic Proof: Choosing a basis that diagonalizes the inner product makes each row of the transformation an independent optimization problem.The resulting row-wise weights become irrelevant, so the objective can be treated as Euclidean in a convenient basis.
E.1 Algebraic Proof: The basis separates the nonzero-variance components of X from the subspace whose whitened components have nonzero covariance with Z.
E.1 Algebraic Proof: Each component of X can be represented as a linear combination of the nontrivial components of its whitened vector.
E.1 Algebraic Proof: The covariance constraint forces coefficients on whitened components correlated with Z to zero.The proof identifies these as the first ℓ components in the chosen basis.
E.1 Algebraic Proof: The particular choice B = I yields P* = I − W+P_WΣ_XZW and leaves non-varying components of X intact.
E.1 Algebraic Proof: The solution is unique except for zero-variance columns and zero-weighted rows of the pseudoinner product.

E.2 Covector Proof

The covector proof recasts the problem in a covariance inner-product space, where erasure becomes projection onto the subspace of linear forms uncorrelated with Z.

E.2 Covector Proof: The proof diagonalizes the inner product so each row of P becomes a separate optimization problem.
E.2 Covector Proof: The objective depends on each row P_i through the scalar random variable ξ = P_i^T X, subject to Cov(ξ, Z) = 0.
E.2 Covector Proof: Linear forms uncorrelated with Z form a subspace Z⊥ under the covariance inner product.
E.2 Covector Proof: Changing coordinates to the whitened vector WX makes covariance-space orthogonality correspond to orthogonality of covectors.
E.2 Covector Proof: The solution projects each component of X onto Z⊥ and assembles the resulting components into the final transformation.
E.2 Covector Proof: The final solution is almost surely equivalent to Eq. 6 while preserving non-varying components of X.The proof notes that singular covariance cases permit arbitrary choices in the kernel without changing the inner product.

F The Optimality of Oblique Projections

LEACE’s optimal erasure need not be an orthogonal projection: an oblique projection can preserve more useful variance while achieving complete linear concept erasure with less mean squared distortion.

In a two-feature example, the concept information resides in X2, so orthogonal erasure nullifies that component entirely.
X1 contains information about X2 without containing information about Z, enabling oblique erasure to partially reconstruct X2 while preserving full concept erasure.
Both orthogonal and oblique methods fully erase prediction of Z, but the oblique method achieves lower mean squared edit distance.
Concept scrubbing can cause embedding norms to diverge at some layers, while a convex combination with an orthogonal projection corrects this in practice without worsening MSE.

I Oracle LEACE

Oracle LEACE removes covariance with a known concept using the nearest transformed embedding under broad positive-semidefinite inner products. Its solution is an appropriately shifted OLS residual, with Figure 5 providing the Hilbert-space projection intuition.

Oracle LEACE assumes oracle access to each concept label and computes the nearest X′ satisfying Cov(X′, Z) = 0.
LEACE is nearest to X under every positive-semidefinite inner product aT Mb by decomposing the optimization into independent component-wise projection problems.
The Oracle LEACE solution is the ordinary least-squares residual from regressing X on Z, plus a constant shift preserving X’s mean.
The derivation separates constant and varying components because covariance constraints are invariant to constant shifts.
In Hilbert-space terms, each centered component is projected onto the subspace orthogonal to the centered concept variables, producing the closest uncorrelated residual.
The method’s notation includes one-hot labels, random-vector components, and square-integrability assumptions for the Hilbert-space formulation.
For ordinary LEACE, whitening, orthogonal projection onto colsp(WΣXZ)⊥, and unwhitening form a three-step construction that guarantees linear guardedness.
Figure 2 compares gender-prediction accuracy with mean squared embedding displacement, while Figure 3 corrects BERT’s no-intervention MLM accuracy to 0.928.