Max-info factorizations of a conditional distribution

Setup

Let $X,Y,Z$ be finite random variables taking values in finite spaces $\mathcal X,\mathcal Y,\mathcal Z$. Let $p(y\mid x)$ be a fixed conditional law on $\mathcal Y$ given $X$. Let $\mathcal R$ be a class of conditionals $r(z\mid x)$, and $\mathcal G$ a class of conditionals $g(y\mid x,z)$.

Define the feasible factorization set

$$F(p):=\Bigl\{(r,g)\in \mathcal R\times \mathcal G:\; p(y\mid x)=\sum_{z\in\mathcal Z} r(z\mid x)\,g(y\mid x,z)\ \ \forall x,y \Bigr\}.$$

For our setting with expressive router-generator classes (e.g., Transformer neural networks), there is not a unique factorization for a given LLM distribution $p(y\mid x)$. That is, $|F(p)|$ is large. We will consider the following max-info factorization criterion:

$$\arg\max_{(r,g) \in F(p)} I(Y; Z \mid X)$$

I.e., among all feasible factorizations, select on that maximizes the conditional mutual information between $Y$ and $Z$ given $X$. In other words, select a factorization that maximizes the information $Z$ encodes about $Y$ beyond the input $X$.

To define $I(Y;Z\mid X)$ one must also fix an input law $\mu$ on $\mathcal X$. Given $(r,g)\in F(p)$, define the induced joint law

$$P_{\mu,r,g}(x,y,z):=\mu(x)\,r(z\mid x)\,g(y\mid x,z).$$

All entropies and mutual informations below are computed under $P_{\mu,r,g}$. All pointwise statements hold for each $x$, and averaging with respect to $\mu$ yields the global statements.

Define the weighted generalized Jensen–Shannon divergence ("Divergence Measures Based on the Shannon Entropy", Lin 2002)

$$\operatorname{JSD}_{r(\cdot\mid x)}(\{g(\cdot\mid x,z)\}_z) := \sum_z r(z\mid x) \, \mathrm{KL}(g(\cdot\mid x,z) \,\Vert\, M) = H(M)-\sum_z r(z\mid x)H(g(\cdot\mid x,z)),$$

where

$$M = \sum_z r(z\mid x) \, g(\cdot\mid x,z).$$

This acts as a measure of the separation or divergence between a collection of distributions $\{g(\cdot\mid x,z)\}_z$.

Exact objective identity

The following result relates the conditional mutual information between $Y$ and $Z$ given $X$ (how much information $Z$ encodes about $Y$ beyond $X$) to the separation between the collection of $z$-conditioned generators $\{g(\cdot\mid x,z)\}$ as measured by the weighted generalized Jensen-Shannon divergence. In turn, this implies that maximizing the conditional mutual information among feasible factorizations, $\max_{(r, g) \in \mathcal{F}(p)} I(Z; Y | X)$ is equivalent to maximizing the separation between the $z$-conditioned generators. This is the sense in which this selection of factorization is "disentangled".

^max-info-jsd-identity

\begin{proof} By definition,

$$\begin{align*} I(Y;Z\mid X) &:= \sum_x \mu(x) \sum_z \sum_y P(z, y \mid x) \log \frac{P(z, y \mid x)}{P(z\mid x) P(y\mid x)} \\ &= \sum_x \mu(x)\sum_z r(z\mid x)\sum_y g(y\mid x,z) \log\frac{g(y\mid x,z)}{p(y\mid x)}. \end{align*}$$

This is exactly

$$\mathbb E_X\sum_z r(z\mid X)\,\mathrm{KL}\!\bigl(g(\cdot\mid X,z)\,\|\,p(\cdot\mid X)\bigr).$$

Now use

$$\sum_z r(z\mid x) \,\mathrm{KL}(g(\cdot\mid x,z) \| p(\cdot\mid x)) = H(p(\cdot\mid x))-\sum_z r(z\mid x) H(g(\cdot\mid x,z)),$$

because $\sum_z r(z\mid x)g(\cdot\mid x, z) = p(\cdot\mid x)$. Averaging over $x\sim\mu$ gives

$$I(Y;Z\mid X)=\mathbb E_X\operatorname{JSD}_{r(\cdot\mid X)}(\{g(\cdot\mid X,z)\}_z).$$

\end{proof}

Another way to interpret this result is that the set of factorizations that maximize the conditional mutual information $I(Y; Z \mid X)$ are exactly the set of factorizations that minimizes the entropies of the $z$-conditioned generators $H(g(\cdot\mid x,z))$ in expectation of $X, Z$. That is, the max-info criterion finds a factorization that maximizes the concentration of the distribution of $y$ after conditioning on the latent variable $z$.

^entropy-minimization-form

\begin{proof} Under $P_{\mu,r,g}$, the conditional law of $Y$ given $(X,Z)=(x,z)$ is $g(\cdot\mid x,z)$. Therefore

$$H(Y\mid X,Z)=\mathbb E_{X,Z}H(g(\cdot\mid X,Z)),$$

and $I(Y;Z\mid X)=H(Y\mid X)-H(Y\mid X,Z)$. \end{proof}

^cloning-invariance

\begin{proof} Feasibility is unchanged since one mixture component was merely duplicated and its mass split. Also,

$$\alpha H(g(\cdot\mid x,z_0))+(1-\alpha)H(g(\cdot\mid x,z_0))=H(g(\cdot\mid x,z_0)),$$

so $\mathbb E_{X,Z}H(g(\cdot\mid X,Z))$ is unchanged. Apply #^entropy-minimization-form. \end{proof}

This already shows that even optimal solutions need not have pairwise-distinct decoders.

The basic upper bound and its optimizer set

Define the deterministic feasible factorization subset

$$F_{\mathrm{det}}(p) := \Bigl\{(r,g)\in F(p): H_{\mu,r,g}(Y\mid X,Z)=0\Bigr\}.$$

Equivalently, for $\mu(x)r(z\mid x)$-almost every $(x,z)$,

$$g(y\mid x,z)=\mathbf 1\{y=\phi(x,z)\}$$

for some map $\phi:\mathcal X\times \mathcal Z\to \mathcal Y$.

The following result states that if there exists a factorization with where the generator is deterministic once conditioned on $z$, then the max-info factorization criterion is exactly the set of all deterministic factorizations.

^deterministic-maximizers

\begin{proof} The bound follows from

$$I(Y;Z\mid X)=H(Y\mid X)-H(Y\mid X,Z),$$

with $H(Y\mid X,Z)\ge 0$. Equality holds if and only if $H(Y\mid X,Z)=0$, which for discrete $Y$ is equivalent to $g(\cdot\mid x,z)$ being Dirac for $\mu(x)r(z\mid x)$-almost every $(x,z)$. \end{proof}

This result says two things:

1. The max-info factorization criterion selects deterministic generators when possible
2. The max-info criterion has no preference among deterministic generators. I.e., latent semantics may not be unique, and may not be consistent across different inputs. In this way, $Z$ can act more like a random seed than a semantically-meaningful hidden cause.

Expressive classes: why the criterion is often degenerate

Here degenerate means that, in expressive classes, the optimization often attains the trivial upper bound $H(Y\mid X)$, and the argmax is a large non-identifiable family.

The following result says that deterministic factorizations are the max-info factorization if and only if the size of the latent space is large enough compared to the size of the support space.

\begin{proof} If $|\mathcal Z|\ge s_*$, then for each $x$ with $\mu(x)>0$, enumerate

$$\operatorname{supp} p(\cdot\mid x)=\{y_1,\dots,y_{m_x}\}, \qquad m_x\le s_*\le |\mathcal Z|.$$

Choose distinct latent states $z_1,\dots,z_{m_x}$, set

$$\phi(x,z_i)=y_i,\qquad r(z_i\mid x)=p(y_i\mid x),$$

and assign zero mass to all other states. Then

$$p(y\mid x)=\sum_z r(z\mid x)\,\mathbf 1\{y=\phi(x,z)\},$$

so a deterministic feasible factorization exists, and #^deterministic-maximizers yields value $H(Y\mid X)$.

Conversely, if $|\mathcal Z|<s_*$, then for some $x_0$ with $\mu(x_0)>0$,

$$|\operatorname{supp} p(\cdot\mid x_0)|>|\mathcal Z|.$$

Any deterministic factorization at $x_0$ is a convex combination of at most $|\mathcal Z|$ Dirac masses, hence its support size is at most $|\mathcal Z|$, contradiction. Therefore no deterministic feasible factorization exists, so #^deterministic-maximizers forbids achieving $H(Y\mid X)$. \end{proof}

This is the trivial simulator $Z\equiv Y$ in distribution given $X$.

Even imposing $Z\perp X$ does not resolve the degeneracy

The following result states that when $|\mathcal Z|$ is (very) large, even imposing $Z \perp X$ does not resolve the degeneracy.

\begin{proof} Represent the channel $p(y\mid x)$ as a row-stochastic matrix

$$W=[p(y\mid x)]_{x,y}\in \prod_{x\in \mathcal X}\Delta_{\mathcal Y}.$$

The set $\prod_{x\in \mathcal X}\Delta_{\mathcal Y}$ is a convex polytope of dimension $|\mathcal X|(|\mathcal Y|-1)$. Its extreme points are exactly the deterministic channels $D_f$ whose $x$-th row is the Dirac mass at $f(x)$. Therefore,

$$W=\sum_{i=1}^m q_i D_{f_i}$$

for some deterministic maps $f_i:\mathcal X\to \mathcal Y$. By Carathéodory's theorem, every element of $\prod_{x\in \mathcal X}\Delta_{\mathcal Y}$ can be represented as a convex sum of at most $|\mathcal X|(|\mathcal Y|-1)+1$ points. This implies that there exists $(q_i)_{i=1}^{m}, (f_i)_{i=1}^{m}$ such that $W$ can be written as $W=\sum_{i=1}^m q_i D_{f_i}$ with $m\le |\mathcal X|(|\mathcal Y|-1)+1.$

Now let let $|\mathcal Z| = m$ and $Z\sim \boldsymbol{q} = (q_1, \ldots, q_m)$ independently of $X$, and define $Y=f_Z(X)$. Then

$$p(y\mid x)=\sum_z q(z)\,\mathbf 1\{y=f_z(x)\},$$

so the factorization is feasible and deterministic given $(X,Z)$. Hence $H(Y\mid X,Z)=0$, and #^deterministic-maximizers yields $I(Y;Z\mid X)=H(Y\mid X)$. \end{proof}

This implies that, when $|\mathcal Z|$ is large-enough, even the strong restriction $Z\perp X$ is not enough by itself if the decoder may depend on $(X,Z)$ arbitrarily. One must impose additional structure such as small $|\mathcal Z|$, $g(y\mid x,z)=g(y\mid z)$, parametric restrictions, or explicit penalties to get a more structured selection principle.

The structured regime $g(y\mid x,z)=g(y\mid z)$: connection to nonnegative rank

Now impose the shared-decoder restriction

$$g(y\mid x,z)=g(y\mid z)=:g_z(y).$$

This is implicitly assuming a type of information bottleneck where all $y$-relevant information about $x$ is encoded into $z$ by the router. We continue to assume that the router and generator classes are unrestricted subject to this constraint. I.e., $\mathcal{R}$ includes all $r(z\mid x)$ and $\mathcal{G}$ includes all $g(y\mid z)$.

Let $W\in \mathbb R_+^{|\mathcal X|\times |\mathcal Y|}$ be the channel matrix $W_{xy}=p(y\mid x)$, and let $|\mathcal Z|=m$.

We begin by recalling the definition of the non-negative rank.

The following result states that if we restrict the generator class to depend only on $z$, then exact factorization of $p(y\mid x)$ is possible if and only if the nonnegative rank of the channel $W_{xy} = p(y\mid x)$ is at most $|\mathcal Z|$. Moreover, this is equivalent to a nonnegative matrix factorization problem, and the max-info criterion becomes a minimum-entropy nonnegative factorization problem.

\begin{proof} Let $|\mathcal Z|= m$. If there exists a factorization $(r,g)$ of $p(y\mid x)$ such that $p(y\mid x) = \sum_z r(z\mid x) g(y\mid z)$, then the row-stochastic matrices $R, G$ defined as

$$R_{xz}=r(z\mid x),\qquad G_{zy}=g_z(y),$$

form a nonnegative factorization $W=RG$ with inner dimension $m$, so $\operatorname{rank}_+(W)\le m$.

Conversely, if $\operatorname{rank}_+(W)\le m$, there exist $V\in\mathbb R_+^{|\mathcal X|\times m}$, $U\in\mathbb R_+^{m\times |\mathcal Y|}$ such that

$$W=VU.$$

Let

$$s_z:=\sum_y U_{zy}.$$

Discard any zero rows if needed, and define

$$G_{zy}:=\frac{U_{zy}}{s_z},\qquad R_{xz}:=V_{xz}s_z.$$

Then $W=RG$, each row of $G$ sums to $1$, and each row of $R$ sums to

$$\sum_z R_{xz}=\sum_z V_{xz}s_z=\sum_y W_{xy}=1.$$

So $R$ and $G$ are row-stochastic and define the required factorization.

Finally, #^entropy-minimization-form gives

$$H(Y\mid X,Z)=\mathbb E_{X,Z}H(g_Z)=\sum_z P(Z=z)H(g_z),$$

hence the stated objective formula. \end{proof}

This regime defines a restriction on the complexity of the generator class (or can also be read as defining a restriction on the size of $Z$). Unlike the fully expressive case, one cannot usually make $Y$ deterministic given $(X,Z)$ unless the channel matrix admits a sufficiently small deterministic or nonnegative-rank decomposition.

---

References and literature connections

1. J. Lin (1991), Divergence measures based on the Shannon entropy. Introduces the Jensen–Shannon divergence and its weighted generalized form, which is exactly the quantity appearing in #^max-info-jsd-identity. URL: https://www.cise.ufl.edu/~anand/sp06/jensen-shannon.pdf
2. J. E. Cohen and U. G. Rothblum (1993), Nonnegative ranks, decompositions, and factorizations of nonnegative matrices. Standard reference for nonnegative rank and NMF, used in the shared-decoder proposition above. URL: https://lab.rockefeller.edu/cohenje/assets/file/211CohenRothblumNonnegativeRanksDecompositionsFactorsNonnegativeMatricesLAA1993.pdf
3. C. T. Li and A. El Gamal (2017/2018), Strong Functional Representation Lemma and Applications to Coding Theorems. Formalizes the independent-noise representation $Y=\phi(X,Z)$ with $Z\perp X$, connected to the independent-noise theorem above. URL: https://arxiv.org/pdf/1701.02827
4. S. Yang, R. Xu, J. Chen, J. Zhang (2017), Intrinsic Capacity. Discusses decompositions of a channel into deterministic channels and the induced non-uniqueness of intrinsic states, closely related to the degeneracy above. URL: https://arxiv.org/pdf/1706.06858