Expectation and Variance in Hierarchical Models

In our last session, we explored how the first and second moments of a distribution—expectation and variance—provide a powerful lens for understanding core deep learning mechanisms like Monte Carlo estimation and adaptive optimizers. We saw that the "noise" in stochastic gradient descent is nothing more than the variance of a gradient estimator.

This lesson builds directly on that foundation. We will introduce two of the most elegant and useful rules in probability theory: the Laws of Total Expectation and Variance. Your objective is to apply these laws to reason about hierarchical and latent-variable models. These laws are the key to analyzing systems with multiple layers of uncertainty, a structure that is at the heart of many sophisticated generative models, including the Variational Autoencoder (VAE) you plan to implement.

1. The Law of Total Expectation: Averaging the Averages

Imagine you want to find the average height of all students in a university. You could survey every single student, or you could take a "hierarchical" approach: first, calculate the average height within each department (the conditional expectation), and then calculate the average of those departmental averages, weighted by the size of each department. This is the intuition behind the Law of Total Expectation, also known as the Law of Iterated Expectations or the tower property.

Formally, for two random variables $X$ and $Y$ , the law states:

E [X] = E_{Y} [E_{X ∣ Y} [X ∣ Y]]

Here, $E [X ∣ Y]$ is the conditional expectation of $X$ given $Y$ . Crucially, this is itself a random variable because it is a function of the random variable $Y$ . The outer expectation, $E_{Y} [\cdot]$ , then takes the expectation of this new random variable over the distribution of $Y$ .

The following video from MIT OpenCourseWare provides a crisp derivation of this law, treating the conditional expectation as the random object it is.

L13.3 The Law of Iterated Expectations

This video formally introduces the Law of Iterated Expectations, showing that the expectation of a conditional expectation is simply the unconditional expectation.

Please watch the following segments: The initial setup, which explains how the conditional expectation \mathbb{E}[X|Y] is defined as a function of the random variable Y. The derivation, which uses the definition of expectation to prove the law. The logic connects what the video calls the "Total Expectation Theorem" to this more abstract formulation. The brief concluding remarks on its generality.

The key insight is that conditioning on an intermediate variable allows you to break a complex expectation calculation into simpler, nested stages.

2. The Law of Total Variance: Decomposing Uncertainty

Just as we can decompose expectations, we can also decompose variance. The Law of Total Variance provides a way to partition the total uncertainty in a random variable $X$ into components related to another variable $Y$ . The formula is:

Var (X) = E [Var (X ∣ Y)] + Var (E [X ∣ Y])

This equation is profoundly important. It states that the total variance of $X$ is the sum of two parts:

Expected Conditional Variance ( $E [Var (X ∣ Y)]$ ): This is the "within-group" variance. It's the average of the variance of $X$ within specific contexts defined by $Y$ . This represents the variability of $X$ that is not explained by $Y$ .
Variance of Conditional Expectation ( $Var (E [X ∣ Y])$ ): This is the "between-group" variance. It measures how much the average value of $X$ changes as $Y$ changes. This represents the variability of $X$ that is explained by $Y$ .

This decomposition is analogous to the partitioning of variance in ANOVA, a concept you may have encountered. To make this tangible, let's consider a classic problem.

The Variance in the Stick Breaking Problem

This MIT OpenCourseWare video solves the "stick breaking" problem, which serves as an excellent and intuitive application of both the Law of Total Expectation and the Law of Total Variance.

Watch the following segments to see how these laws are used to solve a problem in stages: The introduction provides a fantastic mental model: using conditioning to solve a hard problem by breaking it into simpler stages. The problem setup defines the two-step random process. A stick of length l is broken, keeping the left part of length Y. This piece is broken again, keeping the left part of length X. Applying the Law of Iterated Expectations to find \mathbb{E}[X]. Notice how \mathbb{E}[X|Y] = Y/2 is found first, and then the expectation is taken over Y. Applying the Law of Total Variance to find \text{Var}(X). This part clearly calculates the two separate components, \mathbb{E}[\text{Var}(X|Y)] and \text{Var}(\mathbb{E}[X|Y]), and adds them together.

The stick-breaking example elegantly demonstrates the power of this "divide and conquer" strategy. By conditioning on the outcome of the first break ( $Y$ ), the properties of the second break ( $X$ ) become simple. We can then average these simple properties over the randomness of the first break.

3. Application to Hierarchical and Latent-Variable Models

This staged approach is precisely what we need to analyze hierarchical and latent-variable models. These models are defined by cascades of random variables, where the parameters of one distribution are themselves random variables drawn from another.

A Hierarchical Example

Consider a statistical model for data with a nested structure, like modeling student test scores across different schools in different districts. The scores in a school depend on a school-specific effect, which in turn depends on a district-specific effect.

The following post from Cross Validated discusses exactly such a 3-level model and uses our laws to find the marginal (overall) properties of an observation.

bayesian - How to write down a 3-level hierarchical model with a level-1 equation (reduced form)? - Cross Validated

This Stack Exchange thread presents a 3-level hierarchical model and provides a clear derivation of its marginal mean and variance using the laws of total expectation and variance.

First, read the question itself, starting from the first paragraph, to understand the model setup: y_{rct} \sim N(\mu_{ct}, \sigma_{rct}) \mu_{ct} \sim N(\mu_c, \sigma_{ct}) \mu_c \sim N(0, 1) This defines a clear hierarchy of dependencies. Next, read the first part of the top-voted answer, which applies the laws. Focus on the derivations under the headings "law of total expectation" and "law of total variance". Observe how the expectation is calculated by "towering" up the hierarchy, and the variance is decomposed into the sum of variances at each level.

As the derivation shows, $E [y_{r c t}] = 0$ and $Var (y_{r c t}) = σ_{c}^{2} + σ_{c t}^{2} + σ_{r c t}^{2}$ . This demonstrates a critical result for hierarchical models: the marginal variance of an observation is the sum of the variances at each level of the hierarchy. The model can attribute the total observed variance to different sources of randomness.

Connection to Latent-Variable Models (VAEs)

A latent-variable model like a VAE has a similar two-level hierarchical structure:

A latent code $z$ is sampled from a prior distribution $p (z)$ (e.g., $N (0, I)$ ).
An observation $x$ is sampled from a conditional distribution $p_{θ} (x ∣ z)$ (the "decoder").

We can use the laws to understand the marginal distribution $p_{θ} (x) = \int p_{θ} (x ∣ z) p (z) d z$ , which represents the model's ability to generate data.

Let's assume the decoder outputs the mean of a Gaussian, so $p_{θ} (x ∣ z) = N (x ∣ f_{θ} (z), σ^{2} I)$ , where $f_{θ}$ is the decoder network.

Marginal Mean: Using the Law of Total Expectation:
$E [X] = E_{Z} [E [X ∣ Z = z]] = E_{Z} [f_{θ} (z)]$
The expected output of the VAE is the average of the decoder network's output, taken over the entire latent prior.
Marginal Variance: Using the Law of Total Variance:
$Var (X) = E_{Z} [Var (X ∣ Z = z)] + Var_{Z} (E [X ∣ Z = z])$
Substituting our VAE components:
$Var (X) = E_{Z} [σ^{2} I] + Var_{Z} (f_{θ} (z)) = σ^{2} I + Var_{Z} (f_{θ} (z))$
This is a beautiful result. It tells us that the total variance of the data generated by the VAE comes from two sources: the fixed observation noise $σ^{2}$ from the decoder's output distribution, and the variance of the decoder's output itself as it maps different latent codes $z$ to the data space. A powerful decoder will have a high $Var_{Z} (f_{θ} (z))$ , meaning it can generate a diverse manifold of data.

These expectations (like $E_{Z} [f_{θ} (z)]$ ) and marginals ( $p_{θ} (x)$ ) are often intractable to compute directly, which is precisely why we resort to Monte Carlo estimation (as we saw in the last lesson) and variational inference (which we will cover in detail soon).

Conclusion

In this lesson, we've added two indispensable rules to our probabilistic toolkit. These laws provide the formal framework for analyzing systems with layered randomness, moving from abstract rules to concrete applications in the models you are working towards.

Key Takeaways:

The Law of Total Expectation, $E [X] = E [E [X ∣ Y]]$ , allows us to compute a complex expectation by breaking it down into a nested sequence of simpler, conditional expectations.
The Law of Total Variance, $Var (X) = E [Var (X ∣ Y)] + Var (E [X ∣ Y])$ , decomposes total uncertainty into two sources: the average variance within groups and the variance between group averages.
These laws are the natural way to analyze hierarchical and latent-variable models. They allow us to "collapse" the layers of a model to understand the marginal properties of the data it describes or generates.
For a VAE, this decomposition reveals that the model's generative diversity stems from the variance of its decoder network's outputs across the latent space.

In our next lesson, we will begin our survey of the specific probability distributions that serve as the fundamental building blocks for nearly all deep learning models. We'll start with the distributions for discrete outcomes: the Bernoulli and Categorical distributions, and connect them directly to the binary and categorical cross-entropy loss functions used in classification.