flowchart LR w --> u1 --> u2 --> u3 --> u5 --> L w --> u4 --> u5
Differentiable Probabilistic Programming
static support · a differentiable density U(X) = -log γ(X) · automatic differentiation · MAP · random-walk MH · Hamiltonian Monte Carlo
Javier Burroni
2026-06-29
What Chapter 7 changes
\[\newcommand\grad{\nabla}\]
In Chapters 4–6, the evaluator and the message interface can run programs whose execution structure may depend on random choices. That flexibility is a problem for gradient-based inference.
To run MAP or HMC we need a fixed-dimensional object \(X = (x_1, \ldots, x_d)\) and a differentiable function \(X \mapsto -\log\gamma(X)\). So Chapter 7 makes a compromise:
- deterministic computation may still be rich and dynamic;
- but the latent random variables must have static support;
- here, the addresses collected by the evaluator give a fixed vector of continuous latents.
Where this lecture goes
Chapter 7 turns a probabilistic program with fixed continuous latents into a differentiable function \(X \mapsto -\log\gamma(X)\). MAP optimizes that function; HMC uses its gradient to sample from the posterior.
Making static support operational
We first run the program once in address-discovery mode, fixing the latent address set and the vector it indexes:
\[\mathcal{A} = (a_1,\ldots,a_d), \qquad X = (x_{a_1},\ldots,x_{a_d}) \in \mathbb{R}^d.\]
Then every evaluation of the potential must visit exactly those addresses. If an evaluation visits a different set, the program is outside the static-support fragment that MAP and HMC handle, and the runtime rejects it.
A dynamic check, not a static analysis
We do not magically make arbitrary programs static-support. We restrict to programs whose sample-address sequence is stable, and the runtime checks that assumption on every potential evaluation.
Activity: discover the static support
Before MAP or HMC, we need to turn the program into a fixed-dimensional problem.
The first pass runs the program once in discovery mode:
- at every
sample, record the address; - instead of drawing randomly, send back a deterministic dummy value, here
0; - at every
observe, send back the observed value; - when the program terminates, return the ordered list of sample addresses.
This gives us the latent vector
\[ X = (x_{a_1}, \ldots, x_{a_d}). \]
Complete _discover.
Potential evaluation
At a sample site \(a_i\): read \(x_{a_i}\) from \(X\) and add \(-\log p_i(x_{a_i})\) to \(U\). At an observe site: read the fixed datum \(y_j\) and add \(-\log p_j(y_j \mid X)\) to \(U\).
After execution, require that the visited sample addresses equal \(\mathcal{A}\). The program then defines
\[ \begin{aligned} U \colon \mathbb{R}^d &\to \mathbb{R}, \\ X &\mapsto -\log\gamma(X) = -\sum_{i=1}^d \log p_i(x_{a_i}) - \sum_{j=1}^m \log p_j(y_j \mid X). \end{aligned} \]
Why automatic differentiation appears
Execution at a fixed \(X\) hands us this single scalar, the potential \(U(X)\).
Everything in the lecture now turns on one question:
The question AD answers
If running the program computes \(U(X)\), how do we obtain its gradient \(\grad_X U(X)\) with respect to the whole latent vector?
Question 1
Suppose we have
\[\text{i.e.,} \qquad f(x) = 64x(1-x)(1-2x)^2(1-8x+8x^2)^2.\]
We need to evaluate \(f'(x)\).
What does reverse mode automatic differentiation do?
a) It differentiates the closed-form formula symbolically and simplifies it using SymPy.
b) It approximates the derivative numerically with \([f(x+\epsilon)-f(x-\epsilon)]/2\epsilon\) for a tiny \(\epsilon\).
c) It applies the chain rule to the operations the code actually runs, giving a derivative without forming the closed-form expression.
d) It just automates evaluating an \(f'\) you derived by hand.
Vote: https://pe.app/ppls
Four routes to the same derivative
To get \(f'(x)\) you could:
- By hand: code \(f'\) yourself. Exact, but error-prone.
- Symbolically (SymPy): expand into a (larger) derivative formula. Exact, but balloons.
- Numerically: \([f(x+\epsilon)-f(x-\epsilon)]/2\epsilon\). Approximate, and one pass per parameter.
- Automatic differentiation is none of these: chain rule on the operations the code runs. One sweep, every parameter.
“Back propagation is an algorithm that computes the chain rule, with a specific order of operations that is highly efficient.” (Goodfellow, Bengio & Courville)
The answer: reverse-mode automatic differentiation
loss.backward() runs reverse-mode AD: the efficient, adjoint-accumulating implementation of the chain rule on the operations the code runs. Up to ordinary floating-point arithmetic, with no finite-difference approximation, one backward sweep gives \(\partial L/\partial\phi\) for every parameter.
This is what lets us express models of arbitrary complexity (like the full GPT) and optimize with gradient-based methods such as SGD.
Two passes
Forward: run the code, record each value and the operation that produced it.
Backward: traverse the recorded operations in reverse topological order, accumulating adjoints.
backward() call differentiates all of them.A caveat about the book’s Algorithm 17
The book presents a recursive backward procedure that propagates derivatives from the output back to the inputs. That procedure expresses the chain rule, but it is not quite what we usually mean by backpropagation.
Important
Backpropagation is the dynamic-programming version of the chain rule: it accumulates one adjoint per intermediate value and processes the recorded operations in reverse order.
Why this matters:
- a recursive path-sum may revisit the same shared subcomputation many times;
- true reverse-mode AD accumulates all contributions at that node once;
- each primitive operation/edge is processed once in the backward sweep.
Without that accumulation, the recursive path-sum can be exponential in the number of operations.
The book’s Algorithm 17 is therefore useful as a chain-rule explanation, but the efficient algorithm is the tape/adjoint implementation used by PyTorch, JAX, TensorFlow, etc.
…but it wasn’t always obvious we’d need it
AD is the engine under every framework.
Maximum a posteriori (MAP) estimation
MAP uses exactly that object, \(X^\star = \arg\min_X U(X)\). With \(\grad_X U\) in hand, just descend.
It is a good warm-up because it checks that the program can be evaluated at supplied latent values, that the score is differentiable, and that PyTorch can backpropagate through the execution. On the regression model it lands on the posterior mode \((a,b)=(1.75,-0.62)\).
A point, not the posterior
MAP returns a single high-density point. It is not the main inference method; it is the proof that the machinery works.
Question 2
MAP maximizes the unnormalized log-posterior \(\log\gamma(X)\), never the true posterior \(\log\pi(X) = \log\gamma(X) - \log Z\).
Why is it safe to ignore the normalizing constant \(Z\)?
a) Because \(Z = 1\) for any model built from sample and observe.
b) Because \(Z\) does not depend on \(X\), so it shifts \(\log\gamma\) by a constant and changes neither the gradient nor the \(\arg\max\).
c) Because the optimizer estimates \(Z\) from the points it visits along the way.
d) Because the prior is uniform, so \(Z\) cancels.
Random walk Metropolis–Hastings (RWMH)
To get the posterior we need to sample. The simplest controller uses only \(U\). Propose
\[X' = X + \epsilon, \quad \epsilon \sim \mathcal{N}(0, \sigma^2 I), \qquad \alpha = \min\!\left\{1, \exp\!\left(-U(X') + U(X)\right)\right\}.\]
It samples, and its mean is right.
But this posterior is correlated: the slope and intercept trade off, so the mass is a thin diagonal ridge, and isotropic steps fight that geometry.
Even tuned, it produces far fewer effective draws than HMC.
Why we are not done
MAP finds one high-density point. RWMH samples, but struggles with ridges. We still want a sampler that uses the geometry of \(U\), not just its values.
Hamiltonian Monte Carlo (HMC): move with the gradient
Read \(U = -\log\gamma\) as a potential energy. Give the position \(X\) a momentum \(R\), integrate with leapfrog using \(\grad U\), and accept on the change in \(H = U(X) + \tfrac12\lVert R\rVert^2\).
The gradient bends each trajectory according to the local geometry, while the momentum keeps it from moving purely uphill, so the chain explores the ridge instead of collapsing to a point. In the notebook HMC accepts about 0.75 and reaches far more effective draws than the random walk.
Extra reading
Radford M. Neal, MCMC using Hamiltonian dynamics: a thorough treatment of HMC, leapfrog integration, and tuning.
Question 3
If HMC uses the gradient of \(\log\gamma(X)\), why does it not simply climb to the MAP and stay there?
a) Because the gradient is used to move (ideally) along curves of equal density on the joint \((X,R)\).
b) Because the gradient changes \(X\) directly, but the random momentum refresh sometimes pushes \(X\) away from the MAP.
c) Because the gradient is used only to choose a local Gaussian approximation around the current point.
d) Because the gradient is used only after proposing, to correct the acceptance probability.
Question 4
If the ideal HMC trajectory follows curves of equal joint density in \((X,R)\), why do we still need a Metropolis–Hastings accept/reject step?
a) Because the momentum refresh changes the density, and the accept/reject step corrects for drawing a new \(R\).
b) Because leapfrog preserves the joint density exactly, but the accept/reject step is needed to make the proposal symmetric.
c) Because leapfrog is a numerical approximation: it is reversible and volume-preserving, but it does not preserve the joint density exactly.
d) Because the trajectory follows an iso-density curve only in \(X\)-space; the accept/reject step restores the correct density on \((X,R)\).
Three controllers over one potential
Every method scores with \(U\); they differ in whether they use its gradient, and in what they return.
| method | uses \(U(X)\)? | uses \(\grad U(X)\)? | output |
|---|---|---|---|
| MAP | yes | yes | one point (the mode) |
| RWMH | yes | no | posterior samples, slow on ridges |
| HMC | yes | yes | posterior samples, moves with the geometry |
A note on variational inference
A fourth use, not covered today: fit a guide \(q_\phi(X)\) by maximizing the ELBO \(\mathbb{E}_{q_\phi}[\log\gamma(X) - \log q_\phi(X)]\), whose gradient again comes from AD.
Activity
In the notebook, MAP, RWMH, and HMC are controllers over the message interface. They all call one function, the potential controller, which returns \(U = -\log\gamma\), here with both message branches removed:
Complete both branches in your notebook. Hint: each branch subtracts a log-density from U; they differ only in where the scored value comes from.
The .sum() collapses a possibly vector-valued log_prob to a scalar. Because d.log_prob is differentiable, U carries a gradient back to the latents, which is what makes MAP and HMC work over the unchanged evaluator.
Next class
Reading
- Chapter 8.1: programs as deep generative models
- Chapter 8.2: programs as inference models, guide programs
- Chapter 8.3: stochastic-gradient learning and inference, the ELBO, amortized inference
Focus while reading
- where a neural network enters a probabilistic program
- what an inference model (\(q_\phi\)) proposes, and how it is scored
- which objective the gradient now climbs
- how a proposal can be learned instead of fixed
Answer key
Solutions
- Q1: c. AD applies the chain rule to the operations the code runs, so it returns an exact derivative without a symbolic formula (a) or a finite-difference approximation (b), and it does the differentiating, not just evaluation of a hand-derived \(f'\) (d).
- Q2: b. \(Z\) is constant in \(X\), so it shifts \(\log\gamma\) without moving the gradient or the \(\arg\max\). It is not generally \(1\) (a), the optimizer never sees \(Z\) (c), and the prior here is Gaussian, not uniform (d).
- Q3: a. HMC is not gradient ascent on \(\log\gamma(X)\). The gradient is coupled with momentum \(R\), so the ideal motion follows joint iso-density curves in \((X,R)\) instead of climbing to the MAP.
- Q4: c. The exact trajectory would preserve joint density, but leapfrog is only a numerical approximation. It is reversible and volume-preserving, so a Metropolis–Hastings step can correct the remaining energy/density error.
- Activity.
x = theta[a]; U = U - d.log_prob(x).sum(); m.send(x), andobserveis the same ony. Each branch only subtracts a log-density; the latent’s value is whatever the optimizer or sampler currently holds.
Addendum: reverse mode AD
Question 2
A regularized linear unit, with \(x\) and \(y\) fixed constants:
\[L(w) = (w x - y)^2 + w^2\]
Rewrite this as a list of elementary operations, one per line, naming each intermediate.
Hint: start with \(u_1(w, x) = w \cdot x\).
\[\begin{aligned} u_1(w, x) &= w \cdot x \\ u_2(u_1, y) &= u_1 - y \\ u_3(u_2) &= u_2 \cdot u_2 \\ u_4(w) &= w \cdot w \\ u_5(u_3, u_4) &= u_3 + u_4 \\ L(w) &= u_5 \end{aligned}\]
You just wrote the forward pass
Each \(u_i\): one operation on earlier values. Note \(w\) appears twice, in \(u_1\) and \(u_4\).
First, see \(L\) as a composition
Naming each output \(u_i\) abstracts it into one step from earlier values, so \(L\) is a nested composition:
\[L \;=\; u_5\big(\,\underbrace{u_3(u_2(u_1(w)))}_{(wx-y)^2},\;\; \underbrace{u_4(w)}_{w^2}\,\big),\]
\[u_1=wx,\quad u_2=u_1-y,\quad u_3=u_2^2,\quad u_4=w^2,\quad u_5=u_3+u_4.\]
Multi-argument ops branch the graph: \(w\) feeds \(u_1\) and \(u_4\), reaching \(L\) by two paths.
Differentiate the composition with the chain rule.
Recall: the chain rule
For \(L=f(u(w))\), \(\frac{\partial L}{\partial w}=\frac{\partial L}{\partial u}\frac{\partial u}{\partial w}\).
Apply the chain rule and reuse
\(L = u_5\big(\,u_3(u_2(u_1(w))), u_4(w)\,\big),\)
\(\frac{\partial L}{\partial w} =\) \(\color{#984ea3}{\frac{\partial L}{\partial u_1}}\frac{\partial u_1}{\partial w} + \color{#984ea3}{\frac{\partial L}{\partial u_4}}\frac{\partial u_4}{\partial w}\)
Compute each \(\partial L/\partial u_i\) once from \(L\) and save it. Each later gradient reuses an incoming factor:
through \(u_1\)
\(\dfrac{\partial L}{\partial u_3} = \underbrace{\color{#984ea3}{\dfrac{\partial L}{\partial u_5}}}_{\text{incoming}} \cdot \underbrace{\dfrac{\partial u_5}{\partial u_3}}_{\text{local}}\)
\(\dfrac{\partial L}{\partial u_2} = \dfrac{\partial L}{\partial u_5}\dfrac{\partial u_5}{\partial u_3}\dfrac{\partial u_3}{\partial u_2}\) \(=\; \color{#984ea3}{\dfrac{\partial L}{\partial u_3}}\,\dfrac{\partial u_3}{\partial u_2}\)
\(\dfrac{\partial L}{\partial u_1} = \dfrac{\partial L}{\partial u_5}\dfrac{\partial u_5}{\partial u_3}\dfrac{\partial u_3}{\partial u_2}\dfrac{\partial u_2}{\partial u_1}\) \(=\; \color{#984ea3}{\dfrac{\partial L}{\partial u_2}}\,\dfrac{\partial u_2}{\partial u_1}\)
through \(u_4\)
\[\frac{\partial L}{\partial u_4} = \color{#984ea3}{\frac{\partial L}{\partial u_5}}\,\frac{\partial u_5}{\partial u_4}\]
flowchart LR w --> u1 --> u2 --> u3 --> u5 --> L w --> u4 --> u5
The backward pass in one line
Sweep from \(L\) back to \(w\), seed \(\partial L/\partial L = 1\).
Every gradient has the same shape: an incoming factor times a local one:
\[\boxed{\;\frac{\partial L}{\partial u_i} \;=\; \sum_{j\,:\, u_i \to u_j}\underbrace{\color{#984ea3}{\frac{\partial L}{\partial u_j}}}_{\text{incoming}}\cdot\underbrace{\frac{\partial u_j}{\partial u_i}}_{\text{local}}\;}\]
Sum over ops where \(u_i\) is an input.
One edge: one term. Several: a sum.
Local now, the rest later
Local \(\partial u_j/\partial u_i\) is fixed when \(u_j\) is computed (e.g. \(u_1 = wx\) gives \(\partial u_1/\partial w = x\)).
Incoming \(\partial L/\partial u_j\) comes later, from the backward pass.
The idea: let the program record itself
Wrap every number in a Var; overload +, *, exp, …, so each op records itself as Python runs.
Built at runtime, the graph is whatever the code ran, so if, loops, and calls need no special handling.
Each Var stores:
- its value
- the inputs it was built from
- a function
_backwardthat pushes gradient to those inputs
class Var:
def __init__(self, value, _inputs=()):
self.value, self.grad = value, 0.0
self._inputs = tuple(_inputs)
self._backward = lambda: None # filled in by each op
def __add__(self, other):
out = Var(self.value + other.value, _inputs=(self, other))
def _backward():
pass # ← filled in next
out._backward = _backward
return outWhat each operation records
Each operation stores on its output out a _backward closure: a local rule for sending gradient to its inputs.
The rule it records
Given the incoming gradient out.grad (\(= \partial L/\partial u_j\)), it adds each input’s local contribution: \[\underbrace{\partial L/\partial u_i}_{\text{input's }\texttt{.grad}} \mathrel{+}= \underbrace{\partial L/\partial u_j}_{\texttt{out.grad},\ \text{read}}\;\cdot\;\underbrace{\partial u_j/\partial u_i}_{\text{local, captured now}}\]
flowchart LR w --> u1 --> u2 --> u3 --> u5 --> L w --> u4 --> u5
Walking the graph: one reverse sweep
Each op’s _backward pushes to its own inputs (left). The driver (right) orders the nodes, seeds the output, fires each rule once.
What one op leaves behind
The order
build lists each node after its inputs. reversed(topo) then reaches every node only after all its consumers: out.grad is final before out._backward() runs. Reverse-topological order.
Three ops, only the partial changes
Multiplication
Exponential
From toy ops to real layers
Multiplication’s scaled-by-the-other-input rule is the weight gradient of every linear layer. Write the partials; backward() does the rest.
Extra Question
A friend sends VarAlt: each node pushes gradient straight to its inputs the instant it’s reached. “Saves memory, no slowdown.” Same net on both:
class VarAlt:
def __init__(self, value):
...
def _backward(g): self.grad += g
self._backward = _backward
def __add__(self, other):
...
def _backward(g): self._backward(g); other._backward(g)
out._backward = _backward
return out
def sigmoid(self):
...
out._backward = lambda g: self._backward(g * s * (1 - s))
return out
def backward(self): self._backward(1)
Which version finishes first?
a) VarAlt finishes first.
b) They finish at about the same time.
c) Your engine wins; VarAlt lags far behind.
d) Neither finishes.
tiny_resnet wasn’t a toy
Skip connections, everywhere
Each transformer block = 2 sub-layers (attention, MLP), each in its own skip. 2 skips/block.
96 × 2 = 192 skips, all layer(h) + h, just deeper. VarAlt pays per path to the input.
Not the skip: multiple consumers
A value feeding 2+ ops blows up with VarAlt.
- skip: \(2^{n}\) paths
- \(n\) MLPs (width \(d\)): \(d^{n}\)
AD is not just the chain rule
Both apply the chain rule and reach the identical gradient. They differ only in how much work they repeat.
Var: once per edge
One push per edge, regardless of paths.
Cost: \(O(\#\text{edges})\).
The stored inputs and ordering buy exactly this.
VarAlt: once per path
Pushes the instant a node is reached, so edges below a shared node are re-traversed once per path.
h.sigmoid() + h reuses h twice: \(P(k) = 2P(k-1)\), so \(2^n\) paths to the input.
Cost: exponential.
flowchart RL h3["L = σ₃ + h₂<br/>(output)"] --> s3["σ₃ = sigmoid(h₂)"   ] h3 --> h2["h₂ = σ₂ + h₁"] s3 --> h2 h2 --> s2["σ₂ = sigmoid(h₁)"   ] h2 --> h1["h₁ = σ₁ + x"] s2 --> h1 h1 --> s1["σ₁ = sigmoid(x)"   ] h1 --> x["x<br/>(input)"] s1 --> x
Introduction to Probabilistic Programming