SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Somnath Basu Roy Chowdhury  
UNC Chapel Hill
Nicholas Monath Google DeepMind
Avinava Dubey  
Google Research
Amr Ahmed Google Research
Snigdha Chaturvedi
{somnath, snigdha}@cs.unc.edu
{nmonath, avinavadubey, amra}@google.com
Distributed representations provide a vector space that captures meaningful relationships between data instances. The distributed nature of these representations, however, entangles together multiple attributes or concepts of data instances (e.g., the topic or sentiment of a text, characteristics of the author (age, gender, etc), etc). Recent work has proposed the task of concept erasure [50, 52], in which rather than making a concept predictable, the goal is to remove an attribute from distributed representations while retaining other information from the original representation space as much as possible. In this paper, we propose a new distance metric learning-based objective, the Kernelized Rate-Distortion Maximizer (KRaM), for performing concept erasure. KRaM fits a transformation of representations to match a specified distance measure (defined by a labeled concept to erase) using a modified rate-distortion function. Specifically, KRaM's objective function aims to make instances with similar concept labels dissimilar in the learned representation space while retaining other information. We find that optimizing KRaM effectively erases various types of concepts—categorical, continuous, and vector-valued variables—from data representations across diverse domains. We also provide a theoretical analysis of several properties of KRaM's objective. To assess the quality of the learned representations, we propose an alignment score to evaluate their similarity with the original representation space. Additionally, we conduct experiments to showcase KRaM's efficacy in various settings, from erasing binary gender variables in word embeddings to vector-valued variables in GPT-3 representations.
Learned representations, particularly distributed representations [30], are at the core of machine learning with applications in natural language [39], images [21], biology [4], physics [3], and several other domains [40, 48]. These vector-based representations of data instances create an inner product space where similarities and nearest-neighbor relationships are "meaningful". However, due to the distributed nature of these representations, the definition of "meaningful" is often not easily discernible. In other words, shared properties of data instances nearby in the vector space are not always evident. These shared properties are often referred to as concepts [31]. For instance, concepts in representations of images include objects in the image, whether it is indoor or outdoor, etc. Similarly, concepts in text representations include the topics, and characteristics of the author (e.g., geographic location, gender, etc.). For applications that necessitate conditioning on specific attributes to make or explain predictions, these distributed representations can pose challenges. Consequently,
a significant body of work has focused on jointly learning representations and disentangling their underlying concepts [28, 29].
However, state-of-the-art representations for tasks across various domains often come from pretrained models (e.g., ViTs [23] for images, GPT [14] for text, among others). In many cases, these pre-trained representations are utilized directly without fine-tuning the original model, due to factors such as computational burden or limited model API access [24]. This presents a challenge for fitting disentangled representations of data instances and their concepts since the model parameters are frozen. Rather, it presents an opportunity for learning a transformation (or similarly learning a distance metric) for the pre-trained representations.
While learning such a transformation of pre-trained representations can be used for many applications (e.g., classification, regression, etc.) that involve disentangling specific concepts from representations, we focus on the recently proposed task of concept erasure [50]. The objective of concept erasure is twofold: (1) to learn representations that minimize the classification accuracy (or mean squared error) for a specific concept variable and (2) to retain as much other information from the original representations as possible. This becomes possible when the representations are not correlated with the concept variable. We should note that there are indeed some trivial methods to reduce the correlation – such as generating random representations or making all representations identical. However, these solutions fail to retain any information from the original representation space. This highlights one of the key challenges in concept erasure – retaining information from the original space while removing a given concept. This challenge is accentuated by the lack of a pre-specified downstream task for which the representations could be optimized. Instead, objectives for concept erasure use no supervision (apart from the labels of the concept to erase). This makes it difficult to use adversarial learning [60] or mutual information estimation [17, 43] methods for concept erasure.
The independence of concept erasure from down-stream tasks makes it amenable to a variety of applications that use the modified representations. For example, when developing a toxicity classifier for online comments, an organization may seek to ensure that content from different religious backgrounds is treated equitably. This can be achieved using concept erasure to remove information about religion from the text representations [60]. Concept erasure has also shown promise in interpreting the decision-making of large models by studying counterfactual scenarios where certain properties of the input are erased from intermediate layers [26].
In this paper, we present a new distance metric learning-based objective for concept erasure. We refer to our objective, Kernelized Rate-Distortion Maximizer (KRaM). The objective fits a transformation of representations to match a specified distance measure using a kernelized rate-distortion function, where the kernel is constructed using concept labels. Specifically, KRaM's objective function tries to make instances with similar concept labels dissimilar in the learned representation space (Figure 1). Empirically, we find that optimizing KRaM results in representations that are uncorrelated with the concept variable, effectively leading to its erasure. To evaluate the quality of the representations, we propose a  $k$ -nearest neighbour based measure to capture the alignment of the learned representations with the original representation space. We conduct extensive experiments to demonstrate that KRaM is capable of erasing various types of concepts—categorical, continuous, and vector-valued variables—from data representations across a wide range of domains. We also theoretically analyze several properties of the proposed objective function. Our primary contributions are:
- We introduce a novel framework - KRaM, which uses a kernelized formulation of the rate-distortion function that is able to delete a range of concepts - categorical, continuous, or vector-valued variables from representations (Section 3).
- We propose a computationally efficient alignment measure, to evaluate how informative the learned representations are about the original representation space (Section 4).
- We perform a theoretical analysis of KRaM's objective and alignment measure. We conduct extensive experiments to showcase the efficacy of KRaM in a range of settings, from erasing binary gender variables in word embeddings to vector-valued variables in GPT-3 representations (Section 5).
2 Preliminaries & Background
In this section, we first formally describe the concept erasure setup, then discuss some prior concept erasure techniques, and finally introduce the fundamentals of rate-distortion theory.
Problem Setup. In concept erasure [50, 52], we consider the input representations  $x \in \mathcal{X}$  and the concept  $a \in \mathcal{A}$  as random variables. We assume access to samples  $[(x_1, a_1), (x_2, a_2), \ldots]$  drawn from the joint distribution  $P(\mathcal{X}, \mathcal{A})$ . The goal of concept erasure is to learn a function  $f(\cdot)$  that generates representations  $[f(x_1), f(x_2), \ldots] \in \mathcal{Z}$ , such that it is infeasible to predict the concept labels  $a \in \mathcal{A}$  from  $\mathcal{Z}$ . In addition to erasing the concept variable,  $f(x) \in \mathcal{Z}$  should retain as much information about  $x \in \mathcal{X}$  as possible. We do not impose any constraints on the nature of the concept. It can be: categorical ( $a \in [k]$ ), continuous ( $a \in \mathbb{R}$ ), or vector-valued ( $a \in \mathbb{R}^d$ ) random variable.
Concept Erasure is also closely related to learning invariant representations with respect to an attribute through adversarial learning [8, 25, 32, 60]. However, concept erasure differs from adversarial learning in two key aspects: (a) the input representations  $x \in \mathcal{X}$  remain frozen during the concept erasure process (only erasure function  $f$  is updated), and (b) it does not rely on a specific downstream task. This setup is beneficial in situations where we can access representations but lack the necessary resources or infrastructure to train or fine-tune the model that generated them. In the following section, we describe the details of the proposed concept erasure framework, KRaM.
Prior Work. Concept erasure was initially introduced by [13] in the context of removing binary gender labels from GloVe embeddings [46]. Initial works on this problem [13, 49] performed concept erasure by projecting representations onto the null space of the optimal separating linear subspace for the categorical concept. Recent work [50] has introduced a generalized objective for this solution, presenting it as a minimax game between concept identification and nullspace projection, and further provided a closed-form solution for its relaxed convex version. Nonetheless, these techniques are limited by two main assumptions: (a) the erasure function  $f(\cdot)$  is linear, and (b) the concept variable is categorical. A linear erasure function ensures that a linear subspace, which could identify the concept label for instances, does not exist in the learned representation space,  $\mathcal{Z}$ . Consequently, it prevents any linear network from extracting the concept labels from the learned representations. However, it can still be possible for a non-linear network to predict the concept labels by identifying a non-linear concept subspace. Given that most modern ML architectures rely on non-linear networks, it is crucial to ensure that the concept is inaccessible to non-linear networks. More recent works have made progress towards non-linear concept erasure. This has been done either by using a linear concept erasure after projecting the input into a non-linear feature space [52], or directly utilizing a non-linear erasure function  $f$  using a rate-distortion objective [18]. Despite their potential, these concept erasure techniques require access to categorical concept labels for all instances. Hence, it is not possible to erase other forms of concept variables (continuous or vector-valued) using these techniques without discretization of the concept labels, which often leads to information loss. In contrast to these techniques, our erasure framework KRaM is able to handle a variety of concept variables (categorical, continuous, or vector-valued) while performing non-linear concept erasure. This becomes possible as KRaM presumes access to a kernel matrix that is defined by the concept labels, and does not impose any additional constraints on the nature of the concept variable. Next, we discuss the fundamentals of the rate-distortion function that forms a building block of our framework.
Rate Distortion. In information theory [20], the compactness of a distribution is measured by their coding length - the number of binary bits required to encode it. In lossy data compression, a set of vectors  $\mathcal{Z} = \{z_1,\dots ,z_n\} \in \mathbb{R}^{n\times d}$ , sampled from a distribution  $P(\mathcal{Z})$ , is encoded using a coding scheme, such that the transmitted vectors  $\{\hat{z}_i\}_{i = 1}^n$  can be recovered up to a distortion  $\epsilon$ . The minimal number of bits required per vector to encode the sequence  $\mathcal{Z}$  is defined by the rate-distortion function  $R(\mathcal{Z},\epsilon)$ . The optimal  $R(\mathcal{Z},\epsilon)$  for vectors  $\mathcal{Z}$  sampled from a multivariate Gaussian  $\mathcal{N}(0,\Sigma)$  is:
$$
R (\mathcal {Z}, \epsilon) = \frac {1}{2} \log_ {2} \det  \left(I + \frac {d}{n \epsilon^ {2}} \mathcal {Z} \mathcal {Z} ^ {T}\right), \tag {1}
$$
where  $n$  is the number of vectors and  $d$  is the dimension of individual vectors. Equation 1 provides a tight bound even in cases where the underlying distribution  $P(\mathcal{Z})$  is degenerate [38]. The rate-distortion function is also closely tied to the sphere packing problem [38] and represents the volume (or intrinsic dimension) of a representation set. Recent works like MCR² [59] have built on the rate-distortion function to learn discriminative representations for classification tasks. Concept erasure techniques [18, 19], have also used the rate-distortion function to erase categorical variables. Even though KRaM uses the rate-distortion function similar to these techniques, it is more versatile and capable of erasing different types of concept variables. KRaM also imposes additional constraints on the feature space to ensure robust concept erasure, which we discuss in the following section.
Figure 1: An illustration of concept erasure using KRaM. Input representations are retrieved from a large language model. The original representations (a) encodes a binary concept variable (the two classes are shown in  $\bullet$  and  $\bullet$ ), which we aim to erase.  $R(\mathcal{Z}|\mathbf{K})$  term forces instances from the same class to move apart. However, for robust concept erasure the size of the representation space  $R(Z)$  matters, which we illustrate visually in (b), (c), and (d). KRaM enforces the constraint  $R(Z) = b$  to erase the concept while retaining information from original representations.
3 Kernelized Rate-Distortion Maximizer (KRaM)
We proceed by discussing how a concept variable can be erased from a representation set. A concept cannot be extracted by a predictive network (which is equivalent to erasure) if there is minimal correlation between the representation set,  $\mathcal{Z}$ , and the concepts,  $\mathcal{A}$  [58]. Note that distances in the representation space can be indicative of the concept variable. For example, in Figure 1 (a) we observe that instances in the original representation space with similar concept labels (shown by their color  $\bullet$  and  $\bullet$ ) appear close to each other. Here, the distance between instances is correlated with the concept label, thereby making it feasible to identify the concept labels via a linear or non-linear boundary. Ideally, we want a representation space where distances are not reflective of concept labels, where instances with different concept labels appear together (e.g., Figure 1 (c) & (d)).
The intuition behind our approach,  $\mathsf{KRA}\mathsf{M}$ , is to make the distances in the learned representation space,  $\mathcal{Z}$ , uncorrelated with the concept variable,  $\mathcal{A}$ . Specifically, we try to make instance pairs similar in the concept space to be distant (or dissimilar) in the learned representation space,  $\mathcal{Z}$ . We achieve this by learning an erasure function  $f$  (parameterized by a neural network) to transform the input  $x \in \mathcal{X}$  into  $\mathcal{Z}$ . We propose a kernelized formulation of the rate-distortion function to train  $f$ :
$$
R (\mathcal {Z} | \mathbf {K}) = \frac {1}{2} \log_ {2} \det  \left(I + \frac {d}{n \epsilon^ {2}} \mathcal {Z} \mathcal {Z} ^ {T} \odot \mathbf {K}\right), \tag {2}
$$
where  $\mathcal{Z} = f(\mathcal{X})\in \mathbb{R}^{n\times d}$  and the kernel matrix,  $\mathbf{K}\in \mathbb{R}^{n\times n}$ , captures similarities between concept labels. The entries of the kernel matrix are inversely proportional to the distance between concept labels  $\mathbf{K}_{ij}\propto 1 / \mathrm{d}(a_i,a_j)$ , where  $\mathrm{d}(\cdot ,\cdot)$  can be an arbitrary symmetric distance function, such that  $\mathrm{d}(x,x) = 0$ . We observe that  $R(\mathcal{Z}|\mathbf{K})$  is sensitive to the scale of the representations  $f(x)\in \mathcal{Z}$ . Therefore, we fix the Frobenius norm of the representations using a layer normalization layer [5] to make  $f(x)\in \mathbb{S}^d$ , ensuring that individual instances have an equal impact on the loss.
Next, we discuss how maximizing  $R(\mathcal{Z}|\mathbf{K})$  (Equation 2) helps in concept erasure. We proceed by noting that maximizing the standard version of the rate-distortion function (Equation 1) is equivalent to increasing the covariance of the representations,  $\mathcal{Z}\mathcal{Z}^T$ . In the kernelized rate-distortion function (Equation 2), we observe that the kernel matrix  $\mathbf{K}$  assigns higher weights to instance pairs that have similar concept labels ( $\mathbf{K}_{ij} \propto 1 / \mathrm{d}(a_i, a_j)$ ). Intuitively, this means that maximizing  $R(\mathcal{Z}|\mathbf{K})$  results in an increased dissimilarity between instance pairs with similar concept labels, as illustrated by the arrows in the center of Figure 1(a). This gradually leads to the distances in the learned representation space  $\mathcal{Z}$  being unrelated to the concept labels.
However, simply maximizing  $R(\mathcal{Z}|\mathbf{K})$  may not guarantee robust concept erasure without imposing additional constraints on the overall feature space. We explain why this happens by considering a few scenarios. First, consider the scenario where we do not impose any constraints on the feature space, shown in Figure 1(b). In this scenario, the volume (or intrinsic dimension) of the feature space (analogous to  $R(\mathcal{Z})$  term) also expands as  $R(\mathcal{Z}|\mathbf{K})$  is maximized (Lemma 1). Here, we observe that even though the intra-group distances have increased it is still possible to separate the two groups
using a non-linear decision boundary. Second, we consider the scenario where we try to minimize the volume of the feature space, which is equivalent to minimizing  $R(\mathcal{Z})$ . This is illustrated in Figure 1(c), where all instances are pushed together making it hard to predict the concept labels. However, it also results in a significant loss of information from the original representations (as the volume or intrinsic dimension collapses). As different instances become almost similar it destroys the unique features present in Figure 1(a), potentially rendering them ineffective for downstream tasks. Thus, it appears that the optimal approach is to maintain a constant size of the feature space as illustrated in Figure 1(d). We verify these scenarios empirically in Section 5.2. With this consideration, we present the following objective:
$$
\max  _ {f} R (\mathcal {Z} | \mathbf {K}), \text {s u b j e c t t o} R (\mathcal {Z}) = b, \tag {3}
$$
where  $b = R(\mathcal{X})$  is the initial number of bits required to encode the data and  $\mathcal{Z} = f(\mathcal{X})$ . In practice, we found that satisfying the equality  $R(\mathcal{Z}) = b$  using a Lagrangian function hinders the maximization of  $R(\mathcal{Z}|\mathbf{K})$ . For concept erasure, we only want the feature space volume to be constant and do not require it to exactly be  $b$ . Therefore, we optimize a relaxed version of the objective (Equation 3):
$$
\max  _ {f} R (\mathcal {Z} | \mathbf {K}) - \lambda | R (\mathcal {Z}) - b |, \tag {4}
$$
where  $\lambda$  is a hyperparameter. The second term in Equation 4 penalizes the network,  $f$ , if the overall volume  $R(\mathcal{Z})$  deviates too far from  $b$ . Depending on the nature of the attribute (categorical, continuous, or vector-valued), the user can define the kernel matrix  $\mathbf{K}$  between concept labels. For categorical concept variables, in our experiments, we use the kernel matrix whose values are  $\mathbf{K}_{ij} \in \{0,1\}$ , where  $\mathbf{K}_{ij} = 1$  if  $a_i = a_j$  otherwise  $\mathbf{K}_{ij} = 0$ . For continuous and vector-valued variables, the kernel matrix can be derived from the concept labels by using a suitable kernel function (e.g., Gaussian, Laplacian, or Cauchy kernels) if it is not specified by the user. In our experiments, we use a Gaussian (RBF) kernel function for continuous and vector-valued concepts.
Lemma 1 (General Bounds for  $R(\mathcal{Z}|\mathbf{K})$ ). For any set of representations  $\mathcal{Z} \in \mathbb{R}^{n \times d}$ , a kernel matrix  $\mathbf{K} \in \mathbb{R}^{n \times n}$  using a kernel function satisfying  $k(x, x) = 1$  and  $\epsilon > 0$ , it holds that:
$$
R (\mathcal {Z}) \leq R (\mathcal {Z} | \mathbf {K}) \leq \frac {n}{2} \log_ {2} \left(1 + d / n \epsilon^ {2}\right), \tag {5}
$$
where the first equality is satisfied when  $\mathbf{K} = \mathbf{1}\mathbf{1}^T$  and the second equality when  $\mathcal{Z}\mathcal{Z}^T = I$ .
The detailed proof is provided in Appendix A.1. This result shows that  $R(\mathcal{Z}|\mathbf{K})$  has a lower bound equal to the rate-distortion function of the representations with the upper bound being independent of the kernel matrix. We empirically show that maximizing  $R(\mathcal{Z}|\mathbf{K})$  also results in an increase in  $R(\mathcal{Z})$  in Section 5.2. Using the results of the above lemma, we can show that the proposed objective (Equation 4) is bounded in the following corollary.
Corollary 1. Using assumptions in Lemma 1, for  $\lambda \in [0,1]$  the objective function (Equation 4) is bounded between  $[- \lambda b, \max \{(1 + \lambda)U - \lambda b, (1 - \lambda)U + \lambda b\}]$ , where  $U = \frac{n}{2} \log_2(1 + d / n\epsilon^2)$ .
4 Measuring Alignment
A limitation of prior works [18, 49, 50] is the lack of evaluation beyond proxy tasks of how information preserved by the learned representations  $\mathcal{Z}$  about the original representations  $\mathcal{X}$ , which we refer to as alignment. An erasure framework that generates random representations is able to erase concept  $\mathcal{A}$  perfectly but does not retain any information from  $\mathcal{X}$ . Therefore, it is important to measure the alignment as well while optimizing the erasure function  $f$ . To this end, we propose a computationally efficient measure by computing the overlap between  $k$ -nearest neighbour sets of  $\mathcal{X}$  and  $\mathcal{Z}$ . The average nearest neighbour overlap across all representations is the alignment score  $(A_k)$  for a given concept erasure function  $f$ :
$$
A _ {k} (f) = \underset {x \sim \mathcal {X}} {\mathbb {E}} \left[ | \operatorname {k n n} (x) \cap \operatorname {k n n} (f (x)) | ] / k, \right. \tag {6}
$$
where  $\mathrm{knn}(\cdot)$  function computes the  $k$ -nearest neighbour set of a representation. Alignment scores lie between  $A_{k}(f) \in [k/n, 1]$  (Lemma 2). Notice that this measure is quite similar to non-parametric approaches for mutual information (MI) estimation [41, 35]. These methods, while leveraging nearest neighbour information, compute data statistics within a hypercube. However, these MI
estimates are often biased for high-dimensional data [27]. In contrast to these methods, we utilize the bijective mapping (erasure function  $f$ ) between  $\mathcal{X}$  and  $\mathcal{Z}$  (which is not available in the general case of MI estimation) to compute the overlap between nearest neighbour sets. The computation of  $A_{k}$  can be made faster using an efficient nearest neighbour data structure like kd-tree [10]. Note that a similar measure was used to determine the stability of word embeddings [57].