A Robust Natural Latent Over A Mixed Distribution Is Natural Over The Distributions Which Were Mixed

This post walks through the math for a theorem. It’s intended to be a reference post, which we’ll link back to as-needed from future posts. The question which first motivated this theorem for us was: “Redness of a marker seems like maybe a natural latent over a bunch of parts of the marker, and redness of a car seems like maybe a natural latent over a bunch of parts of the car, but what makes redness of the marker ‘the same as’ redness of the car? How are they both instances of one natural thing, i.e. redness? (or ‘color’?)”. But we’re not going to explain in this post how the math might connect to that use-case; this post is just the math.

Suppose we have multiple distributions $P^1,\dots ,P^k$ over the same random variables $X_1,\dots ,X_n$. (Speaking somewhat more precisely: the distributions are over the same set, and an element of that set is represented by values $(x_1,\dots ,x_n)$.) We take a mixture of the distributions: $P\left[X\right]:={\sum }_j{\alpha }_j{P}^j\left[X\right]$, where ${\sum }_j{\alpha }_j=1$ and $\alpha$ is nonnegative. Then our theorem says: if an approximate natural latent exists over $P\left[X\right]$, and that latent is robustly natural under changing the mixture weights $\alpha$, then the same latent is approximately natural over $P^j\left[X\right]$ for all $j$.

Mathematically: the natural latent over $P\left[X\right]$ is defined by $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]\right)$, and naturality means that the distribution $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]P\right[X=x\left]\right)$ satisfies the naturality conditions (mediation and redundancy).The theorem says that, if the joint distribution $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]{\sum }_j{\alpha }_j{P}^j\right[X=x\left]\right)$ satisfies the naturality conditions robustly with respect to changes in $\alpha$, then $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]P^j\right[X=x\left]\right)$ satisfies the naturality conditions for all $j$. “Robustness” here can be interpreted in multiple ways—we’ll cover two here, one for which the theorem is trivial and another more substantive, but we expect there are probably more notions of “robustness” which also make the theorem work.

Trivial Version

First notion of robustness: the joint distribution $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]{\sum }_j{\alpha }_j{P}^j\right[X=x\left]\right)$ satisfies the naturality conditions to within $ϵ$ for all values of $\alpha$ (subject to ${\sum }_j{\alpha }_j=1$ and $\alpha$ nonnegative).

Then: the joint distribution $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]{\sum }_j{\alpha }_j{P}^j\right[X=x\left]\right)$ satisfies the naturality conditions to within $ϵ$ specifically for ${\alpha }_j={\delta }_{jk}$, i.e.$\alpha$ which is 0 in all entries except a 1 in entry $k$. In that case, the joint distribution is $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]P^k\right[X=x\left]\right)$, therefore $\Lambda$ is natural over $P^k$. Invoke for each k, and the theorem is proven.

… but that’s just abusing an overly-strong notion of robustness. Let’s do a more interesting one.

Nontrivial Version

Second notion of robustness: the joint distribution $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]{\sum }_j{\alpha }_j{P}^j\right[X=x\left]\right)$ satisfies the naturality conditions to within $ϵ$, and the gradient of the approximation error with respect to (allowed) changes in $\alpha$ is (locally) zero.

We need to prove that the joint distributions $(x,\lambda \mapsto P[\Lambda =\lambda |X=x\left]P^j\right[X=x\left]\right)$ satisfy both the mediation and redundancy conditions for each $j$. We’ll start with redundancy, because it’s simpler.

Redundancy

We can express the approximation error of the redundancy condition with respect to $X_i$ under the mixed distribution as

$D_{KL}\left(P\right[\Lambda ,X\left]||P\right[X\left]P\right[\Lambda |X_i\left]\right)=E_X\left[D_{KL}\right(P\left[\Lambda |X\right]||P\left[\Lambda |X_i\right]\left)\right]$

where, recall, $P[\Lambda ,X]:=P\left[\Lambda |X\right]{\sum }_j{\alpha }_j{P}^j\left[X\right]$.

We can rewrite that approximation error as:

$E_X\left[D_{KL}\right(P\left[\Lambda |X\right]||P\left[\Lambda |X_i\right]\left)\right]$

$={\sum }_j{\alpha }_j{P}^j\left[X\right]D_{KL}\left(P\right[\Lambda |X\left]||P\right[\Lambda |X_i\left]\right)$

$={\sum }_j{\alpha }_j{E}_X^j\left[D_{KL}\right(P\left[\Lambda |X\right]||P\left[\Lambda |X_i\right]\left)\right]$

Note that $P^j\left[\Lambda |X\right]=P\left[\Lambda |X\right]$ is the same under all the distributions (by definition), so:

$={\sum }_j{\alpha }_j{D}_{KL}\left(P^j\right[\Lambda ,X\left]||P\right[\Lambda |X_i\left]P^j\right[X\left]\right)$

and by factorization transfer:

$\ge {\sum }_j{\alpha }_j{D}_{KL}\left(P^j\right[\Lambda ,X\left]||P^j\right[\Lambda |X_i\left]P^j\right[X\left]\right)$

In other words: if ${ϵ}_i^j$ is the redundancy error with respect to $X_i$ under distribution $j$, and ${ϵ}_i$ is the redundancy error with respect to $X_i$ under the mixed distribution $P$, then

${ϵ}_i\ge {\sum }_j{\alpha }_j{ϵ}_i^j$

The redundancy error of the mixed distribution is at least the weighted average of the redundancy errors of the individual distributions.

Since the ${\alpha }_j{ϵ}_i^j$ terms are nonnegative, that also means

${ϵ}_i^j\le \frac{1}{{\alpha }_j}{ϵ}_i$

which bounds the approximation error for the $i^{th}$ redundancy condition under distribution $j$. Also note that, insofar as the latent is natural across multiple $\alpha$ values, we can use the $\alpha$ value with largest ${\alpha }_j$ to get the best bound for ${ϵ}_i^j$.

Mediation

Mediation relies more heavily on the robustness of naturality to changes in $\alpha$. The gradient of the mediation approximation error with respect to $\alpha$ is:

$\frac{\partial }{\partial {\alpha }_j}{D}_{KL}\left(P\right[\Lambda ,X\left]||P\right[\Lambda \left]{\prod }_{i}P\right[X_i|\Lambda \left]\right)$

$={\sum }_{X,\Lambda }P\left[\Lambda |X\right]P^j\left[X\right]\text{ln}\frac{P[\Lambda ,X]}{P\left[\Lambda \right]{\prod }_{i}P\left[X_i|\Lambda \right]}$

(Note: it’s a nontrivial but handy fact that, in general, the change in approximation error of a distribution $P\left[Y\right]$ over some DAG $d{D}_{KL}\left(P\right[Y\left]||{\prod }_{i}P\right[Y_i|Y_{pa\left(i\right)}\left]\right)$ under a change $dP$ is ${\sum }_{Y}dP\left[Y\right]\text{ln}\frac{P\left[Y\right]}{{\prod }_{i}P\left[Y_i|Y_{pa\left(i\right)}\right]}$.)

Note that this gradient must be zero along allowed changes in $\alpha$, which means the changes must respect ${\sum }_j{\alpha }_j=1$. That means the gradient must be constant across indices $j$:

$\text{constant}={\sum }_{X,\Lambda }P\left[\Lambda |X\right]P^j\left[X\right]\text{ln}\frac{P[\Lambda ,X]}{P\left[\Lambda \right]{\prod }_{i}P\left[X_i|\Lambda \right]}$

To find that constant, we can take a sum weighted by ${\alpha }_j$ on both sides:

$\text{constant}={\sum }_j{\alpha }_j{\sum }_{X,\Lambda }P\left[\Lambda |X\right]P^j\left[X\right]\text{ln}\frac{P[\Lambda ,X]}{P\left[\Lambda \right]{\prod }_{i}P\left[X_i|\Lambda \right]}$

$=D_{KL}\left(P\right[\Lambda ,X\left]||P\right[\Lambda \left]{\prod }_{i}P\right[X_i|\Lambda \left]\right)$

So, robustness tells us that the approximation error under the mixed distribution can be written as

$D_{KL}\left(P\right[\Lambda ,X\left]||P\right[\Lambda \left]{\prod }_{i}P\right[X_i|\Lambda \left]\right)=\text{constant}={\sum }_{X,\Lambda }P\left[\Lambda |X\right]P^j\left[X\right]\text{ln}\frac{P[\Lambda ,X]}{P\left[\Lambda \right]{\prod }_{i}P\left[X_i|\Lambda \right]}$

for any $j$.

Next, we’ll write out $P[\Lambda ,X]$ as a mixture weighted by $\alpha$, and use Jensen’s inequality on that mixture and the logarithm:

$=E^j\left[\text{ln}\frac{{\sum }_j{\alpha }_{j}P\left[\Lambda |X\right]P^j\left[X\right]}{P\left[\Lambda \right]{\prod }_{i}P\left[X_i|\Lambda \right]}\right]$

$\ge E^j\left[{\sum }_j{\alpha }_j\text{ln}\frac{P\left[\Lambda |X\right]P^j\left[X\right]}{P\left[\Lambda \right]{\prod }_{i}P\left[X_i|\Lambda \right]}\right]$

$={\sum }_j{\alpha }_j{D}_{KL}\left(P^j\right[\Lambda ,X\left]||P\right[\Lambda \left]{\prod }_{i}P\right[X_i|\Lambda \left]\right)$

Then factorization transfer gives:

$\ge {\sum }_j{\alpha }_j{D}_{KL}\left(P^j\right[\Lambda ,X\left]||P^j\right[\Lambda \left]{\prod }_i{P}^j\right[X_i|\Lambda \left]\right)$

Much like redundancy, if ${ϵ}_i^j$ is the mediation error with respect to $X_i$ under distribution $j$ (note that we’re overloading notation, $ϵ$ is no longer the redundancy error), and ${ϵ}_i$ is the mediation error with respect to $X_i$ under the mixed distribution $P$, then the above says

${ϵ}_i\ge {\sum }_j{\alpha }_j{ϵ}_i^j$

Since the ${\alpha }_j{ϵ}_i^j$ terms are nonnegative, that also means

${ϵ}_i^j\le \frac{1}{{\alpha }_j}{ϵ}_i$

which bounds the approximation error for the $i^{th}$ mediation condition under distribution $j$.