Information Theory in Representation Learning

As we continue to build our understanding of what constitutes an effective representation, information theory provides a powerful quantitative framework. This branch of mathematics allows us to measure uncertainty and information content, offering precise tools to analyze and design representation learning algorithms, including the VAEs that are central to this course. Understanding these principles will help clarify why certain objective functions are used and how we can evaluate the quality of learned representations.

Quantifying Information and Uncertainty

At the heart of information theory are a few fundamental quantities that help us reason about data and models.

Entropy

Entropy, denoted as $H(X)$ for a random variable $X$, measures the average amount of uncertainty or "surprise" associated with the outcomes of $X$. For a discrete random variable $X$ with probability mass function $P(x)$, its entropy is:

$$H(X) = - \sum_{x \in \mathcal{X}} P(x) \log_2 P(x)$$

The logarithm is typically base 2, in which case entropy is measured in bits. A distribution that is sharply peaked (i.e., one outcome is very likely) has low entropy, while a uniform distribution (all outcomes equally likely) has maximum entropy for a given number of states. In representation learning, entropy can characterize the diversity or complexity of data features or latent variables.

Mutual Information

Mutual Information (MI) is a measure of the amount of information that one random variable contains about another. For two random variables $X$ and $Z$, their mutual information $I(X;Z)$ quantifies the reduction in uncertainty about $X$ that results from knowing $Z$, or vice versa. It's defined as:

$$I(X;Z) = H(X) - H(X|Z) = H(Z) - H(Z|X)$$

where $H(X|Z)$ is the conditional entropy of $X$ given $Z$. MI can also be expressed using KL divergence (discussed next):

$$I(X;Z) = D_{KL}(P(x,z) || P(x)P(z))$$

This shows that MI measures the dependency between $X$ and $Z$. If $X$ and $Z$ are independent, $I(X;Z) = 0$.

In representation learning, we are often interested in a latent representation $Z$ that captures significant information about the input $X$. Thus, a high $I(X;Z)$ is generally desirable. For instance, an encoder in an autoencoder aims to produce a $Z$ that retains as much information as possible about $X$ to allow for accurate reconstruction. MI is also a foundation for understanding and promoting disentanglement, where we might want different components of a latent vector $Z = (Z_1, ..., Z_d)$ to be informative about distinct, independent factors of variation in the data, implying low $I(Z_i; Z_j)$ for $i \neq j$.

Kullback-Leibler (KL) Divergence

The Kullback-Leibler (KL) divergence, or relative entropy, measures how one probability distribution $P$ diverges from a second, expected probability distribution $Q$. For discrete distributions $P$ and $Q$ defined over the same probability space $\mathcal{X}$, it's given by:

$$D_{KL}(P || Q) = \sum_{x \in \mathcal{X}} P(x) \log \frac{P(x)}{Q(x)}$$

For continuous distributions, the sum is replaced by an integral. Important properties of KL divergence include:

• $D_{KL}(P || Q) \ge 0$.
• $D_{KL}(P || Q) = 0$ if and only if $P = Q$.
• It is not symmetric: $D_{KL}(P || Q) \neq D_{KL}(Q || P)$ in general.

In VAEs, as we'll see in detail in Chapter 2, the KL divergence plays a significant role. It often appears in the VAE objective function as a regularization term, encouraging the learned distribution of latent variables $q(z|x)$ (the approximate posterior) to be close to a chosen prior distribution $p(z)$ (e.g., a standard normal distribution). This regularization is important for ensuring that the latent space has good properties for generation.

The Information Bottleneck Principle

The Information Bottleneck (IB) principle provides a formal framework for learning representations that are both compressed and informative. Given an input variable $X$ and a target variable $Y$ (which could be class labels in a supervised task, or $X$ itself for reconstruction), the goal is to learn a stochastic mapping to a representation $Z$, $p(z|x)$, that acts as a "bottleneck." This $Z$ should be maximally informative about $Y$ while being minimally informative about $X$.

This trade-off is formalized by the objective:

$$\mathcal{L}_{IB} = I(Z;Y) - \beta I(X;Z)$$

We aim to maximize this Lagrangian, where $\beta$ is a Lagrange multiplier that controls the trade-off between the informativeness of $Z$ about $Y$ and the compression of $X$ into $Z$.

• Maximizing $I(Z;Y)$ means $Z$ should retain as much information as possible about the target $Y$.
• Minimizing $I(X;Z)$ (by maximizing $-I(X;Z)$) means $Z$ should compress $X$, discarding information in $X$ that is not relevant to $Y$.

The Information Bottleneck framework. The representation $Z$ is learned to be a compressed version of input $X$ while retaining information relevant to a target $Y$.

The IB principle is highly relevant to VAEs. While not always explicitly framed this way, the VAE objective encourages learning a compressed latent representation $Z$ (via the KL divergence term, which can be related to $I(X;Z)$ under certain conditions) that is sufficient for reconstructing $X$ (which relates to $I(X;Z)$ or $I(Z;X)$ where $Y=X$). Understanding IB helps motivate the structure of VAE objectives and the properties desired in learned latent spaces.

Information Theory in the Context of VAEs

As hinted, information-theoretic quantities are not just analytical tools; they are deeply embedded in the mechanics of VAEs.

1. Reconstruction and $I(X;Z)$: The reconstruction term in the VAE objective, typically $-\mathbb{E}_{q(z|x)}[\log p(x|z)]$, encourages the decoder $p(x|z)$ to accurately reconstruct $X$ from $Z$. This implicitly pushes $Z$ to retain information about $X$, thus relating to $I(X;Z)$. A $Z$ that is highly informative about $X$ will allow for better reconstruction.
2. KL Divergence and Regularization: The term $D_{KL}(q(z|x) || p(z))$ in the VAE objective forces the approximate posterior $q(z|x)$ for each input $x$ to be close to a prior $p(z)$. This term can be interpreted in several ways:
   • As a regularizer, it prevents $q(z|x)$ from becoming too complex or too specific to individual $x_i$, promoting a smoother, more organized latent space.
   • It encourages encoding efficiency, as $z$ samples can be thought of as being drawn from $p(z)$.
   • It has connections to $I(X;Z)$. For instance, if $p(z)$ is factorial (i.e., components are independent) and $q(z|x)$ is also encouraged to be factorial, it helps in achieving some form of disentanglement. Minimizing $D_{KL}(q(z|x) || p(z))$ under an evidence lower bound (ELBO) maximization implicitly controls the channel capacity between $X$ and $Z$.

Evaluating Representations

Guiding the learning process, information theory provides tools for evaluating the quality of learned representations. For example, mutual information can be used to assess:

• Informativeness: How much does a latent variable $Z_i$ tell us about a known underlying factor of variation in the data?
• Disentanglement: Are different latent variables $Z_i$ and $Z_j$ statistically independent? $I(Z_i; Z_j)$ can quantify this. We will revisit these evaluation aspects in Chapter 5 when discussing metrics for disentanglement, many of which are derived from information-theoretic principles.

In summary, information theory offers a precise language and a set of tools to analyze the flow of information in probabilistic models like VAEs. It helps us understand what it means for a representation to be "good" (e.g., informative, compressed, disentangled) and provides mechanisms to build these properties into our models. This foundation will be valuable as we move into the mathematical specifics of VAEs and their advanced variants.