This post assumes some familiarity with the idea of concept erasure and our LEACE concept erasure method. We encourage the reader to consult our arXiv paper for background.
In our paper titled LEACE: Perfect linear concept erasure in closed form, we introduced a novel concept erasure method that achieves least squares optimality within the framework of affine transformations. Building on this foundational work, we now delve deeper into the topic by presenting an extension that focuses on deriving the least squares optimal edit while constraining the edited representation as a function of the unedited representation. This refined approach allows us to address critical aspects of concept erasure without the need for label access during inference.
Understanding the Need for Free Form LEACE
In previous discussions, we tackled the problem of concept erasure under the assumption that the transformation could depend on both the representation $mathbf{x}$ and the label $mathbf{z}$. While this approach permits a more nuanced edit, it also introduces complexities, particularly the necessity of accessing labels at inference time. Moreover, this method risks injecting non-linearly represented information into the representation itself, a concern we explored in our earlier blog post. By shifting our focus away from label dependency, we can navigate these challenges more effectively.
Theorem 1: Introducing Free Form LEACE
This brings us to Theorem 1, where we introduce Free Form LEACE (“FF-LEACE”). The essence of this theorem is to derive a closed-form formula for the function $r : mathbb{R}^n rightarrow mathbb{R}^n$ that remains closest to the identity function while ensuring that no linear classifier can outperform random chance when predicting $mathrm{Z}$ from $mathrm{r(X)}$. In mathematical terms, this translates to ensuring that $mathrm{Cov}(mathrm{r(X)}, mathrm{Z}) = mathbf{0}$. This formulation is pivotal as it allows us to achieve optimality while adhering to the constraints set forth in our analysis.
Derivation of the Free Form LEACE
We start by considering a more generalized scenario where we assume $Omega = mathbb{R}^n$ and $h(x) = x$. The derivation begins with the following setup:
Theorem 1.
Let $X$ be a random object taking values in the set $Omega$, and let $Z$ be a random vector in $mathbb{R}^k$. Define a measurable function $h: Omega rightarrow mathbb{R}^n$ and $f : Omega rightarrow mathbb{R}^k$ such that $f(x) = mathbb{E}[Z | X=x]$. Under the assumption that $h(X)$ and $f(X)$ possess finite second moments, we can express the objective:
$$
mathop{mathrm{inf hspace{0.5em}}}_{substack{mathrm r : Omega rightarrow mathbb{R}^n}} mathbb{E} big| mathrm r(X) – mathrm h(X) big|^2_{mathbf M} quad mathrm{s.t.} hspace{0.5em} mathrm{Cov}(mathrm r(X), mathrm Z) = mathbf{0}
$$
This objective is minimized by the function:
$$
r^*(x) = h(x) – mathbf{Sigma}_{h(X)Z} mathbf{Sigma}_{f(X)f(X)}^+ big(mathrm f(x) – mathbb{E}[mathrm Z]big)
$$
In this notation, $mathbf{A}^{+}$ represents the Moore-Penrose pseudoinverse of a matrix $mathbf{A}$, an essential component in our derivation.
Proof of Theorem 1
To establish the validity of our theorem, we first recognize that $mathrm{Cov}(r(X), Z) = mathrm{Cov}(r(X), f(X))$, allowing us to reformulate the objective:
$$
P_1 = mathop{mathrm{inf hspace{0.5em}}}_{substack{mathrm r : Omega rightarrow mathbb{R}^n}} mathbb{E} big| mathrm r(X) – mathrm h(X) big|^2_{mathbf M} quad mathrm{s.t.} hspace{0.5em} mathrm{Cov}(mathrm r(X), mathrm f(X)) = mathbf{0}
$$
Next, we introduce a related objective defined in the Hilbert space of square-integrable real-valued random variables, leading us to:
$$
P_2 = mathop{mathrm{inf hspace{0.5em}}}_{substack{mathrm Y in mathcal H}^n} mathbb{E} big| mathrm Y – mathrm h(X) big|^2_{mathbf M} quad mathrm{s.t.} hspace{0.5em} mathrm{Cov}(mathrm Y, mathrm f(X)) = mathbf{0}
$$
By establishing a correspondence between feasible functions for $P_1$ and random variables for $P_2$, we demonstrate that $P_1 leq P_2$. Given that both $h(X)$ and $f(X)$ have finite second moments, we can leverage results from our previous work, showing that:
$$
mathrm Y_{mathrm{LEACE}} = mathrm h(X) – mathbf{Sigma}_{h(X)f(X)} mathbf{Sigma}_{f(X)f(X)}^+ big(mathrm f(X) – mathbb{E}[mathrm f(X)]big)
$$
Ultimately, we find that $Y_{mathrm{LEACE}}$ aligns with our optimal function $r^*(X)$, confirming that $P_1 = P_2$ with minimizer $r^*$. This result is pivotal in linking our theoretical framework to practical applications in concept erasure.
Practical Challenges in Implementing FF-LEACE
One of the primary hurdles in applying FF-LEACE lies in estimating the conditional expectation $f(x) = mathbb{E}[Z | X=x]$. This task can be particularly daunting, especially when dealing with high-dimensional data. However, with access to a comprehensive dataset, one potential solution is to use a neural network to approximate $f(x)$. Once we have a learned function, we can effectively implement FF-LEACE, leveraging its optimality to enhance our concept erasure capabilities.
For those interested in practical applications, a PyTorch implementation of FF-LEACE is available in our GitHub repository. In this resource, we demonstrate its utility on a toy problem where the conditional expectation can be computed in closed form, providing a hands-on opportunity to explore the method’s effectiveness.
Inspired by: Source
