Resampling Conserves Redundancy & Mediation (Approximately) Under the Jensen-Shannon Divergence

[David Lorell]

Around two months ago, John and I published Resampling Conserves Redundancy (Approximately). Fortunately, about two weeks ago, Jeremy Gillen and Alfred Harwood showed us that we were wrong.

This proof achieves, using the Jensen-Shannon divergence (“JS”), what the previous one failed to show using KL divergence (“$D_{KL}$”). In fact, while the previous attempt tried to show only that redundancy is conserved (in terms of $D_{KL}$) upon resampling latents, this proof shows that the redundancy and mediation conditions are conserved (in terms of JS).

Why Jensen-Shannon?

In just about all of our previous work, we have used $D_{KL}$ as our factorization error. (The error meant to capture the extent to which a given distribution fails to factor according to some graphical structure.) In this post I use the Jensen Shannon divergence.

$D_{KL}\left(U||V\right):=E_{U}ln\frac{U}{V}$

$JS\left(U||V\right):=\frac{1}{2}{D}_{KL}\left(U||\frac{U+V}{2}\right)+\frac{1}{2}{D}_{KL}\left(V||\frac{U+V}{2}\right)$

The KL divergence is a pretty fundamental quantity in information theory, and is used all over the place. (JS is usually defined in terms of $D_{KL}$, as above.) We have pretty strong intuitions about what $D_{KL}$ means and it has lots of nice properties which I won’t go into detail about, but we have considered it a strong default when trying to quantify the extent to which two distributions differ.

The JS divergence looks somewhat ad-hoc by comparison. It also has some nice mathematical properties (its square root is a metric, a feature sorely lacking from $D_{KL}$) and there is some reason to like it intuitively: $JS\left(U||V\right)$ is equivalent to the mutual information between X, a variable randomly sampled from one of the distributions, and Z, an indicator which determines the distribution X gets sampled from. So in this sense it captures the extent to which a sample distinguishes between the two distributions.

Ultimately, though, we want a more solid justification for our choice of error function going forward.

This proof works, but it uses JS rather than $D_{KL}$. Is that a problem? Can/​Should we switch everything over to JS? We aren’t sure. Some of our focus for immediate next steps is going to be on how to better determine the “right” error function for comparing distributions for the purpose of working with (natural) latents.

And now, for the proof:

Definitions

Let $P$ be any distribution over $X$ and $\Lambda$.

I will omit the subscripts if the distribution at hand is the full joint distribution with all variables unbound. I.e.$P_{X,\Lambda }$ is the same as P. When variables are bound, they will be written as lower case in the subscript. When this is still ambiguous, the full bracket notation will be used.

First, define auxiliary distributions $Q$, $S$, $R$, and $M$:

$Q:=P_X{P}_{\Lambda |X_1}$, $S:=P_X{P}_{\Lambda |X_2}$, $R:=P_X{Q}_{\Lambda |X_2}=P_X{\sum }_{X_1}\left[P_{X_1|X_2}{P}_{\Lambda |X_1}\right]$, $M:=P_{\Lambda }{P}_{X_1|\Lambda }{P}_{X_2|\Lambda }$

Q, S, and M each perfectly satisfy one of the (stochastic) Natural Latent conditions, with Q and S each satisfying one of the redundancy conditions ($X_2\to X_1\to \Lambda$, and $X_1\to X_2\to \Lambda$, respectively,) and M satisfying the mediation condition ($X_1\leftarrow \Lambda \to X_2$).

R represents the distribution when both of the redundancy factorizations are applied in series to P.

Let $\Gamma$ be a latent variable defined by $P[\Gamma =\gamma |X]:=P[\Lambda =\gamma |X_1]=P[\Gamma =\gamma |X_1]$, with $P^{\Gamma }:=P_{X,\Lambda }{P}_{\Gamma |X}$

Now, define the auxiliary distributions $Q^{\Gamma }$, $S^{\Gamma }$, and $M^{\Gamma }$, similarly as above, and show some useful relationships to P, Q, S, R, and M:

$Q_{X,\gamma }^{\Gamma }:=P_X{P}_{\gamma |X_1}^{\Gamma }=P_{X}Q[\Lambda =\gamma |X_1]=Q[X,\Lambda =\gamma ]$$S_{X,\gamma }^{\Gamma }:=P_X{P}_{\gamma |X_2}^{\Gamma }=P_X{\sum }_{X_1}\left(P_{X_1|X_2}{P}_{\gamma |X_1}\right)=R[X,\Lambda =\gamma ]$, $M_{X,\gamma }^{\Gamma }:=P_{\gamma }^{\Gamma }{P}_{X_1|\gamma }^{\Gamma }{P}_{X_2|\Gamma }^{\Gamma }=P[\Lambda =\gamma ]P[X_1|\Lambda =\gamma ]R[X_2|\Lambda =\gamma ]$

$P_{X,\gamma }^{\Gamma }=P_X{P}_{\gamma |X}=Q[X,\Lambda =\gamma ]$ $P_{\gamma }^{\Gamma }=Q[\Lambda =\gamma ]=P[\Lambda =\gamma ]=P^{\Gamma }[\Lambda =\gamma ]$ $P_{X_1|\gamma }^{\Gamma }=P[X_1|\Lambda =\gamma ]=Q[X_1|\Lambda =\gamma ]$ $P_{X_2|\gamma }^{\Gamma }=\frac{R[X_2,\Lambda =\gamma ]}{P_{\gamma }^{\Gamma }}=R[X_2|\Lambda =\gamma ]$

Next, the error metric and the errors of interest:

Jensen-Shannon Divergence, and Jensen-Shannon Distance (a true metric):

$JS\left(U||V\right):=\frac{1}{2}{D}_{KL}\left(U||\frac{U+V}{2}\right)+\frac{1}{2}{D}_{KL}\left(V||\frac{U+V}{2}\right)$

$\delta (U,V):=\sqrt{JS\left(U||V\right)}=\delta (V,U)$

${ϵ}_1:=JS\left(P||Q\right),{ϵ}_2:=JS\left(P||S\right),{ϵ}_{med}:=JS\left(P||M\right)$

${ϵ}_1^{\Gamma }:=JS\left(P^{\Gamma }||Q^{\Gamma }\right),{ϵ}_2^{\Gamma }:=JS\left(P^{\Gamma }||S^{\Gamma }\right)=JS\left(Q||R\right),{ϵ}_{med}^{\Gamma }:=JS\left(P^{\Gamma }||M^{\Gamma }\right)=JS\left(Q||M^{\Gamma }\right)$

Theorem

Finally, the theorem:

For any distribution $P$ over ($X$, $\Lambda$), the latent $\Gamma \sim P\left[\Lambda |X_i\right]$ has redundancy error of zero on one of it’s factorizations, while the other factorization errors are bounded by small factor of the errors induced by $\Lambda$. More formally:

$\forall P[X,\Lambda ]$, the latent $\Gamma$ defined by $P[\Gamma =\gamma |X]:=P\left[\Lambda |X_1\right]$ has bounded factorization errors ${ϵ}_1^{\Gamma }=0$ and $max({ϵ}_2^{\Gamma },{ϵ}_{med}^{\Gamma })\le 5({ϵ}_1+{ϵ}_2+{ϵ}_{med})$.

In fact, that is a simpler but looser bound than that proven below which achieves the more bespoke bounds of: ${ϵ}_1^{\Gamma }=0$, ${ϵ}_2^{\Gamma }\le (2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_2}{)}^2$, and ${ϵ}_{med}^{\Gamma }\le (2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_{med}}{)}^2$.

Proof

(1) ${ϵ}_1^{\Gamma }=0$

Proof of (1)

$JS\left(P^{\Gamma }||Q^{\Gamma }\right)=0$, since $P_{X,\gamma }^{\Gamma }=Q[X,\Lambda =\gamma ]=Q_{X,\gamma }^{\Gamma }$ and $P_{\Lambda |X}^{\Gamma }=P_{\Lambda |X}$

$■$

(2) ${ϵ}_2^{\Gamma }\le (2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_2}{)}^2$

Lemma 1: $JS\left(S||R\right)\le {ϵ}_1$

$S\left[\Lambda |X_2\right]=P\left[\Lambda |X_2\right]={\sum }_{X_1}P\left[X_1|X_2\right]P\left[\Lambda |X\right]$

$R\left[\Lambda |X_2\right]=Q\left[\Lambda |X_2\right]={\sum }_{X_1}P\left[X_1|X_2\right]P\left[\Lambda |X_1\right]$

$\begin{array}{ccc}JS\left(S||R\right)& =\sum _{X_2}JS\left(S_{\Lambda |X_2}||R_{\Lambda |X_2}\right)& \\ & \le \sum _{X}P\left[X_2\right]P\left[X_1|X_2\right]JS\left(P_{\Lambda |X}||P\right[\Lambda |X_1\left]\right)& \\ & =JS\left(P||Q\right)& \\ & =:{ϵ}_1& \end{array}$^[1]

Lemma 2: $\delta (Q,R)\le \sqrt{{ϵ}_1}+\sqrt{{ϵ}_2}$

$\text{Let}{d}_x:=\delta (P_{\Lambda |x_1},P_{\Lambda |x_2}),a_x:=\delta (P_{\Lambda |x},P_{\Lambda |x_1})\text{, and}{b}_x:=\delta (P_{\Lambda |x},P_{\Lambda |x_2})$

$\begin{array}{ccc}\delta (Q,S)=\sqrt{JS(Q,S)}& =\sqrt{E_{P_X}JS\left(P_{\Lambda |X_1}||P_{\Lambda |X_2}\right)}& \\ & =\sqrt{E_{P_X}(d_X{)}^2}& \\ & \le \sqrt{E_{P_X}(a_X+b_X{)}^2}\text{by the triangle inequality of metric}\delta & \\ & \le \sqrt{E_{P_X}(a_X{)}^2}+\sqrt{E_{P_X}(b_X{)}^2}\text{via the Minkowski Ineqality}& \\ & =\sqrt{JS\left(P||Q\right)}+\sqrt{JS\left(P||S\right)}=\sqrt{{ϵ}_1}+\sqrt{{ϵ}_2}& \end{array}$

Proof of (2)

$\sqrt{{ϵ}_2^{\Gamma }}=\sqrt{JS\left(P^{\Gamma }||S^{\Gamma }\right)}=\sqrt{JS\left(Q||R\right)}=:\delta (Q,R)$

$\begin{array}{ccc}\delta (Q,R)& \le \delta (Q,S)+\delta (S,R)\text{by the triangle inequality of metric}\delta & \\ & \le \delta (Q,R)+\sqrt{{ϵ}_1}\text{by Lemma 1}& \\ & \le 2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_2}\text{by Lemma 2}& \end{array}$

$■$

(3) ${ϵ}_{med}^{\Gamma }\le (2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_{med}}{)}^2$

Proof of (3)

$\begin{array}{ccc}JS\left(M||M^{\Gamma }\right)& =\sum _{\gamma }P[\Lambda =\gamma ]JS\left(P\right[X_2|\Lambda =\gamma \left]||R\right[X_2|\Lambda =\gamma \left]\right)& \\ & =E_{P_{\Lambda }}JS\left(S_{X_2|\Lambda }||R_{X_2|\Lambda }\right)& \\ & \le JS\left(S||R\right)\text{by the Data Processing Inequality}& \end{array}$

$\begin{array}{ccc}\sqrt{{ϵ}_{med}^{\Gamma }}& =\delta (P^{\Gamma },M^{\Gamma })& \\ & =\delta (Q,M^{\Gamma })& \\ & \le \delta (Q,P)+\delta (P,M)+\delta (M,M^{\Gamma })\text{by the triangle inequality of metric}\delta & \\ & =\sqrt{{ϵ}_1}+\sqrt{{ϵ}_{med}}+\sqrt{JS(M,M^{\Gamma })}& \\ & \le \sqrt{{ϵ}_1}+\sqrt{{ϵ}_{med}}+\sqrt{JS(M,M^{\Gamma })}& \\ & \le 2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_{med}}\text{by Lemma 1}& \end{array}$

$■$

Results

So, as shown above, (using Jensen-Shannon Divergence as the error function,) resampling any latent variable according to either one of its redundancy diagrams (just swap ${ϵ}_1$ and ${ϵ}_2$ for the bounds when resampling from $X_2$) produces a new latent variable which satisfies the redundancy and mediation diagrams approximately as well as the original, and satisfies one of the redundancy diagrams perfectly.

The bounds are:
$\begin{array}{ccc}& {ϵ}_1^{\Gamma }=0& \\ & {ϵ}_2^{\Gamma }\le (2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_2}{)}^2& \\ & {ϵ}_{med}^{\Gamma }\le (2\sqrt{{ϵ}_1}+\sqrt{{ϵ}_{med}}{)}^2& \end{array}$

Where the epsilons without superscripts are the errors corresponding to factorization via the respective naturality conditions of the original latent $\Lambda$ and X.

Bonus

For $a,b>0$, $(2\sqrt{a}+\sqrt{b}{)}^2\le 5(a+b)$ by Cauchy-Schwartz with vectors $[2,1],[\sqrt{a},\sqrt{b}]$Thus the simpler, though looser, bound: $max\{{ϵ}_1^{\Gamma },{ϵ}_2^{\Gamma },{ϵ}_{med}^{\Gamma }\}\le 5({ϵ}_1+{ϵ}_2+{ϵ}_{med})$

[Edit: Here is a collab session where I numerically tested the bounds on a system with 3 binary variables, both with random sampling and with a simple gradient ascent test aiming to break the bound. All numerical checks passed.]

1. ^

   The joint convexity of $JS\left(U||V\right)$, which justifies this inequality, is inherited from the joint convexity of KL Divergence.