SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 1: LEACE projection in 3 steps. First the data is whitened, ensuring equal variance in all directions. It is then orthogonally projected onto  $\mathrm{colsp}(\mathbf{W}\Sigma_{\mathrm{XZ}})^{\perp}$ , guaranteeing linear guardedness. Finally, we unwhiten the data so that its covariance structure mimics the original.
Putting together Theorems 4.2 and 4.3 and rearranging, we arrive at the LEACE formula:
$$
r _ {\text {L E A C E}} (\boldsymbol {x}) = \boldsymbol {x} - \mathbf {W} ^ {+} \mathbf {P} _ {\mathbf {W} \boldsymbol {\Sigma} _ {\mathrm {X Z}}} \mathbf {W} (\boldsymbol {x} - \mathbb {E} [ \mathrm {X} ]) \tag {1}
$$
Intuitively, LEACE de-means and whitens  $\pmb{x}$ , projects onto the subspace responsible for correlations between  $X$  and  $Z$ , then unwhitens the result. Finally, it subtracts this value from  $\pmb{x}$ , thereby surgically removing the linearly available information about  $Z$ .
4.2 Oblique Projections are Least-Squares Optimal
Prior work on linear concept erasure has assumed that erasure functions should be orthogonal projections [29, 33, 37], appealing to the well-known fact that an orthogonal projection of a point  $\pmb{x}$  onto a subspace  $U$  yields the nearest point in  $U$  to  $\pmb{x}$ . But even in the case where  $X$  is centered,  $r_{\mathrm{LEACE}}$  is not an orthogonal projection in general. Orthogonal projection matrices are symmetric, and  $\mathbf{I} - \mathbf{W}^{+}\mathbf{P}_{\mathbf{W}\Sigma_{\mathrm{XZ}}}\mathbf{W}$  is only symmetric in the special case where  $\mathbf{P}_{\mathbf{W}\Sigma_{\mathrm{XZ}}}$  and  $\mathbf{W}$  commute. It is an oblique projection however, since applying  $\mathbf{P}^*$  twice yields the same result as applying it once:  $(\mathbf{P}^*)^2 = \mathbf{I} - 2\mathbf{W}\mathbf{P}_{\mathbf{W}\Sigma_{\mathrm{XZ}}}\mathbf{W}^{+} + \mathbf{W}^{+}\mathbf{P}_{\mathbf{W}\Sigma_{\mathrm{XZ}}}\mathbf{W}\mathbf{W}^{-1}\mathbf{P}_{\mathbf{W}\Sigma_{\mathrm{XZ}}}\mathbf{W} = \mathbf{P}^*$ .
Orthogonal projections are generally not least-squares optimal for concept erasure because the necessary and sufficient condition for linear guardedness,  $\mathbf{P}\Sigma_{\mathrm{XZ}} = \mathbf{0}$ , is a constraint on the nullspace of  $\mathbf{P}$ , and not on its range. We may freely choose the range of the projection to minimize the mean squared distance, as long as we zero out  $\operatorname{colsp}(\Sigma_{\mathrm{XZ}})$ . In Figure 1, an orthogonal projection would map all points onto the dashed line, thereby preserving less of the variance of the original data than LEACE does (green line). See Appendix F for a concrete example.
4.3 Extension to Continuous Z
While not a focus of this work, it's worth noting that LEACE can also be applied to the setting where  $Z$  takes arbitrary values in  $\mathbb{R}^k$ , as long as we restrict ourselves to the ordinary least squares regression loss  $\mathcal{L}(\eta, \mathbf{z}) = \| \eta - \mathbf{z} \|_2^2$ . In particular, the proofs of equivalence between conditions 1 and 2 given in Appendix B make no categorical assumption on  $Z$ , and the equivalence between the optimality of a zero weight matrix (condition 2) and zero cross-covariance (condition 4) is well known in the OLS setting. We can then apply Theorems 4.2 and 4.3, which also make no categorical assumption, to derive the same optimal affine eraser as in the categorical case.
5.1 Intrinsic Evaluation
Following Ravfogel et al. [32] we evaluate the ability of our method to remove gender information from the last hidden layer of a frozen BERT model. We use the biographies dataset of De-Arteaga et al. [6], composed of short biographies annotated by both binary gender and profession. We embed each biography with the [CLS] representation in the last layer of BERT, enforce the same-conditional-mean constraint to remove gender information from the [CLS], and then evaluate the performance of the model, after the intervention, on the main task of profession prediction. We
Figure 3: The correlation between  $GAP_{\text{female},y}^{TPR}$  and the relative proportion of women in profession  $y$ , for BERT representation, before (left;  $R = 0.867$ ) and after (right;  $R = 0.392$ ) the projection.
compare our intervention with RLACE [32], which uses gradient-based optimization to solve a linear concept-erasure adversarial game.
Concept erasure results. First, we evaluate the ability of logistic regression classifiers to recover the removed information. The results, presented in Fig. 2, show that our method is the only to achieve random accuracy (perfect erasure) with a small edit, although RLACE (but not INLP) comes close. At the same time, our method is around 2 orders of magnitude faster, and does not require gradient-based optimization.
5.2 Downstream Fairness
How does our intervention affect the behavior of the model on the main classification task of profession prediction? We fit a logistic regression profession-prediction classifier over the projected [CLS] representations.
To measure the bias in a classifier, we follow De-Arteaga et al. [6] and use the TPR-GAP measure, which quantifies the bias in a classifier by considering the difference (GAP) in the true positive rate (TPR) between individuals with different protected attributes (e.g. race or gender). We use the notation  $\mathrm{GAP}_{z,y}^{\mathrm{TPR}}$  to denote the TPR-gap in some main-class label  $y$  (e.g. "nurse" prediction) for some protected group  $z$  (e.g. "female"), we also consider  $\mathrm{GAP}_{z}^{\mathrm{TPR,RMS}}$ , the RMS of the TPR-gap across all professions for a protected group  $z$ :
$$
G A P _ {z} ^ {\mathrm {T P R , R M S}} = \sqrt {\frac {1}{| C |} \sum_ {y \in C} (G A P _ {z , y} ^ {\mathrm {T P R}}) ^ {2}}
$$
To calculate the relation between the bias the model exhibits and the bias in the data, we also calculate  $\sigma_{\mathrm{(GAP^{TPR},\% Women)}}$  the correlation between the TPR gap in a given profession and the percentage of women in that profession.
Figure 2: Gender prediction accuracy after bias-removal projection versus the mean squared distance from the original representation for INLP, RLACE, and LEACE on BERT representations.
Results. The main-task classifier achieves profession-prediction accuracy of  $77.3\%$  on the projected representations (compared with  $79.3\%$  over the original representations), indicating that the intervention minimally affects the ability to predict the profession of a person from the representation of their biography. At the same time, the TPR gap drops significantly from 0.198 to 0.084, indicating a sharp drop in the biased behavior of the profession classifier. Indeed, inspecting the correlation  $\sigma_{(\mathrm{GAP^{TPR}},\% Women)}$  between the gap (per profession) and the representation of women in
Figure 4: Amnesic probing results on bert-base-uncased.
this profession, we see that this correlation plummets from 0.867 to 0.392 after erasure. Re-fitting the main-task logistic regression classifier over the projected representations yields a slightly higher main-task accuracy of  $78.1\%$ , at the price of significantly increasing the TPR gap to 0.158.
5.3 Revisiting Amnesic Probing
Elazar et al. [10] have introduced the idea of amnesic probing as a causal intervention that aims to test the importance of a given concept (e.g. part-of-speech tag) to some main task (e.g. language modeling). They applied Iterative Nullspace Projection (INLP) to remove different concepts from the hidden representations of the model, and assessed the degree to which its behavior changed when performing masked language modeling. Since INLP often requires dozens of iterations to completely erase the concept, its usage in this context raises concerns of collateral damage due to magnitude of the intervention and the non-exhaustive nature of INLP removal. Here, we replicate their experiments on the bert-base-uncased model with our interventions.
Experimental setup. We use part-of-speech (POS) tags as our concept of interest. We collect sentences and their coarse POS tags ("Noun", "Verb" etc.; 18 in total) from the English Universal Dependencies dataset [27]. We tokenize the sentences with the BERT tokenizer and map each wordpiece to the POS tag of the word to which it belongs. We collect the unmasked BERT representations for each layer, intervene to linearly erase the POS concept from that layer, and continue the forward pass until the last layer, from which we compute the distribution of the MLM over the vocabulary. Note that in each experiment we intervene on a single layer. We quantify the decrease in accuracy following the intervention, as well as the increase in the loss. We compare with a baseline intervention of a random orthogonal projection whose null space has the same rank as the label space (18). For INLP, we perform 20 iterations. This is needed because INLP does not effectively remove the concept; even after 20 iterations, classification accuracy is above majority accuracy. As a result, INLP reduces the rank of the representation by 360. By contrast, our method decreases the rank just by 17.
Results. The results are shown in Fig. 4b. Our intervention only mildly changes BERT LM accuracy and loss until layer 8, with the highest drop recorded in layer 11. INLP, in contrast, shows maximum effect at layer 6. Since it removes hundreds of dimensions, it is difficult to attribute this effect to the erasure of the concept. These results suggest that the causal effect of the POS concept on the language model is concentrated in layer 11. Interestingly, this stands in contrast with POS linear probing results, which are optimal at earlier layers [38]. As Elazar et al. [10] have noted, probing does not generally correlate with intervention-based analysis techniques.
The softmax probabilities of a multiclass logistic regression classifier can leak the removed information if another classifier is stacked on top of it [31], though this setup is not linear.
6 Concept Scrubbing
Unfortunately, Elazar et al. [10] were forced to limit their interventions to a single layer due to the limitations of INLP. INLP often requires the deletion of several dozen dimensions before linear guarding is achieved—as demonstrated in Figure 2. Kumar et al. [21] show empirically and theoretically that INLP causes needless "collateral damage" to useful parts of the representation that are orthogonal to the concept being erased. Because of this collateral damage, it's impossible to apply INLP to multiple layers of a transformer without causing its outputs to collapse into gibberish.
Algorithm 1 Concept scrubbing
Require: Model with ℓ layers f = fℓ o ... o f1
Require: Design matrix X ∈ ℝn×d
Require: Label matrix Z ∈ ℝn×k
Ensure: LEACE parameters for each layer in f
1: H1← Embed(X)
2: L ← list()
3: for l ∈ 1... ℓ do
4: Fit (P, b) on Hl and Z
5: Append (P, b) to L
6: Hl← P(Hl - μHl) + μHl (Eq. 1)
7: Hl+1 ← fl(Hl)
8: return L
Instead, we would like to erase all linear information about a concept in every intermediate representation, which we term concept scrubbing. LEACE makes concept scrubbing possible and eminently practical. It causes minimal collateral damage, induces little computational overhead, and the covariance statistics it relies on can be computed in a streaming fashion, without ever storing all the hidden states in memory or on disk.
Algorithm. Any intervention on the model at layer  $\ell$  changes the distribution of hidden states at layers  $\ell' > \ell$ . Because of this, the naive approach of independently fitting LEACE parameters  $(\mathbf{P}, \mathbf{b})$  for all layers of the clean model, then applying them all at once, may fail to fully erase the target concept. Instead, we fit LEACE parameters sequentially, starting from the first layer and proceeding to the final layer. After we compute  $(\mathbf{P}, \mathbf{b})$  for a layer, we immediately use them to scrub the hidden states for that layer, then feed these scrubbed representations to the next layer (Algorithm 1).
Condition	LLaMA	Pythia
7B	13B	30B	160M	1.4B	6.9B	12B
No intervention	0.69	0.66	0.62	0.90	0.70	0.64	0.62
Random erasure	0.69	0.66	0.62	0.99	0.72	0.66	0.63
LEACE	1.73	1.84	1.96	2.79	2.25	3.57	3.20
SAL	3.24	3.26	3.16	3.53	3.44	4.17	4.69
unigram entropy	2.90	2.90	2.90	2.66	2.66	2.66	2.66
Table 1: Perplexity in autoregressive language models when removing linearly available part-of-speech information from the input to each transformer layer. Units are bits per UTF-8 byte. The unigram baseline assigns probabilities to tokens based only on their frequency and not on the context.
6.1 Experimental Details
Dataset. For each model family, we use a sample from the respective pretraining distribution: the validation split of the Pile [13] for the Pythia models [2], and the RedPajama replication of the LLaMA pretraining corpus for the LLaMA family [39]. sampling a slice of  $2^{22}$  tokens for fitting the LEACE parameters and another slice of  $2^{22}$  tokens for evaluation. Since neither corpus comes with part-of-speech tags, we use the model from the SpaCy library [19] to automatically generate Universal Dependency tags [23].
Baseline method. We also run concept scrubbing using full-rank SAL [37], which is similar to our method but lacks a bias term and does not adjust for correlations between features (Appendix D).
Architecture. We focus on autoregressive language models. We evaluate our method on EleutherAI's Pythia 160M, 1.4B, 6.9B, and 12B models [2], and Meta's LLaMA 7B, 13B, and 30B [39]. We apply concept erasure to the input of each transformer block, immediately after normalization is applied (LayerNorm or RMSNorm).
Randomized erasure. Almost any intervention on a neural network will cause its performance to degrade to some extent. Following Elazar et al. [10], we isolate the effect of the concept erasure by comparing it to a control condition in which we orthogonally project onto a random linear subspace of the same rank as the cross-covariance matrix. To reduce the variance of our results, we sample a fresh subspace for each minibatch, and erase that subspace at each layer, reporting the cross-entropy loss averaged over subspaces.
Training efficiency. Algorithm 1 avoids redundant computation by caching the layer  $i$  hidden states for every data point, then using them to run layer  $i + 1$ . This approach has the downside of requiring a large amount of memory or disk space during training (up to 500GB in our experiments). It's possible to avoid caching any hidden states and instead recompute them as needed, at the expense of increasing the total compute cost from  $O(\ell)$  to  $O(\ell^2)$ .
We find strong evidence that autoregressive language models heavily rely on linearly encoded part-of-speech information. While erasing a randomly selected subspace has little to no effect on language modeling performance, scrubbing away part-of-speech information induces a large increase in perplexity across all models (Table 1).
The specific numbers, however, depend on the erasure method used: SAL induces significantly larger increases in perplexity for all models we tested. We take this to mean that SAL inflicts more collateral damage on other useful features in the representation than LEACE does. In other words, interventions made with LEACE are more surgical than those made with prior work; they more closely approximate the ideal of a perfect intervention which only erases the target concept and keeps everything else fixed [40, 15]. If this experiment were conducted with SAL alone, we would have overestimated the causal effect of part-of-speech.
Much work remains to be done to validate concept scrubbing. Specifically, we'd like to see experiments that target concepts much narrower than part-of-speech, and use behavioral metrics to determine whether scrubbing changes the network in the ways we'd intuitively expect. If these experiments succeed, an exciting next step would be the incorporation of concept scrubbing into the pretraining and/or finetuning process. This may make it possible to train deep neural networks subject to conceptual constraints. It remains to be seen if gradient-based optimizers will be able to "circumvent" such constraints by learning completely nonlinear representations of protected attributes.
In this work, we focused exclusively on linear concept erasure due to its simplicity and tractability. Some authors have proposed nonlinear concept erasure techniques based on kernel methods, but have found that erasure functions fit using one kernel do not generalize well to other kernels [33, 37]. We conjecture that it is intractable to nondestructively edit X so as to prevent a general nonlinear adversary from recovering Z, unless the data generating process for X is known in detail.<sup>6</sup>
A major motivation of concept erasure is that it promises to prevent models from using a concept in a post hoc, model-agnostic fashion. But if our concept scrubbing procedure turns out to yield unsatisfactory results in practical use cases, the most promising research direction might then be to improve model-specific techniques, such as those that modify the training procedure [8, 9, 14].
We are grateful to CoreWeave for providing the compute resources used in Section 6. Shauli Ravfogel is grateful to be supported by the Bloomberg Data Science PhD Fellowship.
<sup>6</sup>We suspect erasing a concept is at least as hard as extracting it from the original representation. But in the worst case, information about Z could be encoded cryptographically in X, which would be intractable to decode given standard computational complexity assumptions. If the data is generated by a known algorithm, however, it may be possible to efficiently eliminate mutual information between Z and X by simply breaking the links in the causal graph that connect them.
The problem of linear concept erasure is an instance of the general problem of information removal. Information removal methods generally divide into adversarial methods, which are applied during training, and the post-hoc linear methods considered in this paper. Adversarial methods [8, 42, 4, 9, 44] use a gradient-reversal layer during training to induce representations that do not encode the protected attribute. However, Elazar and Goldberg [9] have shown that these methods fail in exhaustively removing all the information associated with the protected attribute: it is often possible to train new adversaries that successfully recover the removed information. Linear methods have been proposed as a tractable alternative, where one identifies a linear subspace that captures the concept of interest, and neutralizes it using algebraic techniques. Different methods have been proposed for the identification of the subspace, e.g. PCA and variants thereof [3, 20], orthogonal-rotation [7], classification-based [29], spectral [37, 36] and adversarial approaches [32].
Few works theoretically characterize the condition of linear guardedness. Haghhighatkhah et al. [16] extensively analyzed the problem of preventing linear classification, with the focus on decreasing accuracy. They provide a constructive proof of an optimal intervention for an SVM classifier. Ravfogel et al. [31] have proposed a formal definition of linear guardedness based on  $\mathcal{V}$  information, and characterized the fairness implications of guardedness; we show the relations with our definition above. Ravfogel et al. [32] provide an adversarial formulation of the problem, derive a closed-formed solution to certain cases, and propose an SGD-based optimization for others. While they seek an orthogonal projection, we empirically showed that their solution is very close to ours. Sadeghi et al. [35] and Sadeghi and Boddeti [34] both study an adversarial formulation of concept erasure for linear regression, and they trade-off with main-task performance. In contrast to Ravfogel et al. [32], they consider a general linear adversary, i.e. not necessarily a projection matrix. Closest to our work are Kleindessner et al. [20], Haghhighatkhah et al. [17], Shao et al. [37]. As we showed above (§ 4), those methods do achieve the goal of linear guardedness though they are unable to prove this fact. At the same time, they are not optimal in terms of damage to the original representation space.
B Equivalence of Guardedness with the Optimality of Constant Predictors
The following two theorems establish the equivalence of conditions 1 and 2 (indeed, they do so in the general setting, with no assumption of convex loss or linear predictors).
Theorem B.1. Suppose  $\mathrm{X}(\mathcal{V},\mathfrak{L})$ -guards  $\mathbf{Z}$ . Then for every loss  $\mathcal{L}\in \mathfrak{L}$ , the corresponding trivially attainable loss  $L_{\tau}^{(\mathrm{Z},\mathcal{L})}$  cannot be improved upon by any predictor  $\eta (\cdot ;\pmb {\theta})\in \mathcal{V}$ , i.e.  $L_{\tau} = \inf_{\pmb{\theta}}\mathbb{E}[\mathcal{L}(\eta (\mathrm{X};\pmb {\theta}),\mathrm{Z})]$ .
Proof. Consider the null random vector  $\mathrm{X}'(\omega) = 0$ . Since all predictors are constant on  $\mathrm{X}'$ , and the trivially attainable loss gives the best available expected loss among constant predictors, we must have:
$$
L _ {\tau} = \inf  _ {\boldsymbol {\theta}} \mathbb {E} [ \mathcal {L} (\eta \left(\mathrm {X} ^ {\prime}; \boldsymbol {\theta}\right), \mathrm {Z}) ] \tag {4}
$$
The right side of equation (4) is the best possible loss achievable by a function  $\eta (\cdot ;\theta)$  on the joint distribution of  $(\mathrm{X}^{\prime},\mathrm{Z})$ , which by the definition of guardedness is upper bounded by the best possible loss achievable on the joint distribution of  $(\mathrm{X},\mathrm{Z})$ :
$$
\inf  _ {\boldsymbol {\theta}} \mathbb {E} [ \mathcal {L} (\eta \left(\mathrm {X} ^ {\prime}; \boldsymbol {\theta}\right), \mathrm {Z}) ] \leq \inf  _ {\boldsymbol {\theta}} \mathbb {E} [ \mathcal {L} (\eta (\mathrm {X}; \boldsymbol {\theta}), \mathrm {Z}) ] \tag {5}
$$
Combining equations (4) and (5), and the fact that all constant functions exist in our function class  $\mathcal{V} = \{\eta (\cdot ;\pmb {\theta})\}$ , we arrive at our desired result:
$$
L _ {\tau} = \inf  _ {\boldsymbol {\theta}} \mathbb {E} [ \mathcal {L} (\eta (\mathrm {X}; \boldsymbol {\theta}), \mathrm {Z}) ]
$$
Theorem B.2. Suppose that for every loss  $\mathcal{L} \in \mathfrak{L}$ , the corresponding trivially attainable loss  $L_{\tau}^{(\mathrm{Z},\mathcal{L})}$  cannot be improved upon by any predictor  $\eta (\cdot ;\pmb {\theta}) \in \mathcal{V}$ , i.e.  $L_{\tau} = \inf_{\pmb{\theta}}\mathbb{E}[\mathcal{L}(\eta (\mathrm{X};\pmb {\theta}),\mathrm{Z})]$ . Then  $\mathrm{X}$ $(\mathcal{V},\mathfrak{L})$ -guards  $\mathbf{Z}$ .
Proof. Let  $\mathrm{X}' : \Omega \to \mathbb{R}^d$  be any other random data vector with finite first moment.
Since all constant predictors exist in our predictor class  $\mathcal{V} = \{\eta (\cdot ;\pmb {\theta})\}$ , the best loss achievable on  $(\mathrm{X}^{\prime},\mathrm{Z})$  by functions in  $\mathcal{V}$  must be at least as good as the trivially attainable loss (the best loss available by such constant predictors):
$$
\inf  _ {\boldsymbol {\theta}} \mathbb {E} [ \mathcal {L} (\eta (\mathrm {X} ^ {\prime}; \boldsymbol {\theta}), \mathrm {Z}) ] \leq L _ {\tau}
$$
By assumption, the trivially attainable loss cannot be improved upon over  $(\mathrm{X},\mathrm{Z})$  by predictors in  $\mathcal{V}$ :
Since our choice of  $\mathrm{X}^{\prime}$  was arbitrary, this shows that  $\mathrm{X}$  maximizes the minimal achievable loss, so  $\mathrm{X}(\mathcal{V},\mathfrak{L})$ -guards  $\mathrm{Z}$ .
C Linear Guardedness is Equivalent to Linear Statistical Parity
To measure the effect of linear guardedness on main-task classifiers, we use the following minimal definition of "fairness" with respect to an attribute, adapted from Edwards and Storkey [8].
Definition C.1 (Statistical Parity). Let  $\mathrm{X}$  and  $\mathrm{Z}$  be defined as above, and let  $f$  be a function with domain  $\mathbb{R}^d$ . Then  $f$  exhibits statistical parity with respect to  $\mathrm{Z}$  when evaluated on  $\mathrm{X}$  if
$$
\forall z \in \mathcal {Z}: \mathbb {E} [ f (\mathrm {X}) | \mathrm {Z} = z ] = \mathbb {E} [ f (\mathrm {X}) ].
$$
We now prove the equivalence of conditions 3 and 5.