Machine Learning for Astroparticle Physics:
A Crash-course in SBI

Lecture 2b - Deep neural networks for density estimation

Christoph Weniger — University of Amsterdam (GRAPPA)

Today's Roadmap

The hand-picked basis problem

Everything so far depends on hand-picked basis functions \(\boldsymbol{\phi}(x)\).

This limits both regression and classification — the bottleneck is always \(\boldsymbol{\phi}\).

Towards Learnable Building Blocks

Idea: Replace fixed templates with parametric functions. A simple building block:

For a 1D input, what does this look like with random \(\mathbf{v}, b\)?

Each coloured curve is one learnable basis function \(\phi_j(\mathbf{x}) = g(\mathbf{v}_j^T \mathbf{x} + b_j)\) with random \(\mathbf{v}_j, b_j\). Here: 1D input.

Common Activation Functions

The nonlinearity \(g\) determines the shape of each building block. Common choices:

The Artificial Neuron

Each building block \(\phi_j(\mathbf{x}) = g(\mathbf{v}_j^T \mathbf{x} + b_j)\) is called a neuron, by analogy with biology:

One hidden layer = a learnable basis

Stack \(H\) of these neurons and let a final linear layer combine them:

This is exactly the linear-basis model \(f(\mathbf{x}) = \mathbf{w}^T\boldsymbol{\phi}(\mathbf{x})\) from earlier, except now the basis functions have learnable parameters.

This is a neural network with one hidden layer. Stack more layers → a multi-layer perceptron (MLP).

Universal Approximation Theorem

With enough neurons, a single hidden layer can approximate any continuous function.

Intuition: each ReLU neuron contributes a "kink". Combine enough kinks → any shape.

Grey: individual neurons \(w_j \cdot g(v_j x + b_j)\). Black: their sum. More neurons = more expressive.

No Closed-Form Solution

When we cannot solve \(\nabla E = 0\) on paper, we need a numerical method to find the minimum of \(E(\mathbf{w})\). The default one is (stochastic) gradient descent.

What the Network Actually Is

Concretely, the one-hidden-layer network from the previous slide, as math (left) and as the PyTorch module you would actually run (right). Colours link each line:

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, H=16):
        super().__init__()
        self.fc1 = nn.Linear(2, H)   # W⁽¹⁾ (H×2), b⁽¹⁾ (H)
        self.g   = nn.ReLU()         # activation g
        self.fc2 = nn.Linear(H, 1)   # W⁽²⁾ (1×H), b⁽²⁾ (1)

    def forward(self, x):
        z = self.fc1(x)              # z = W⁽¹⁾x + b⁽¹⁾
        h = self.g(z)                # h = g(z)
        return self.fc2(h)           # μ = W⁽²⁾h + b⁽²⁾

The learnable parameters are \(\boldsymbol\phi=\{W^{(1)},\mathbf{b}^{(1)},W^{(2)},b^{(2)}\}\). PyTorch collects them in model.parameters(); the optimizer descends the loss over all of them at once.

Gradient Descent

Bundle every trainable quantity into a single parameter vector \(\boldsymbol\phi\):

Start somewhere, then repeatedly step in the direction that decreases the loss:

\(\eta\) = learning rate. The gradient is taken w.r.t. every component of \(\boldsymbol\phi\) at once.

Illustration: Step Size and Curvature

How large should \(\eta\) be? Consider a simplistic quadratic loss \(E(w) = \tfrac{1}{2}\,\kappa\,w^2\) with curvature \(\kappa = E''(w)\).

Practical challenges: the optimal step size is the inverse curvature of the loss, which is typically not known. In many dimensions, curvature differs per direction.

Mini-Batch Stochastic Gradient Descent (SGD)

Computing the full gradient over all \(N\) samples is expensive. Approximate it with a random mini-batch \(\mathcal{B}\) of size \(B \ll N\). The gradient is taken w.r.t. every component of \(\boldsymbol\phi = (\mathbf{W},\mathbf{b},\sigma_\theta,\ldots)\):

SGD = True Gradient + Noise

The mini-batch gradient is a noisy estimate of the true gradient:

MLP in Action: Regression

Train a one-hidden-layer MLP on 1D regression with the Gaussian NLL loss. The grey band is the trained \(\sigma\) — all of \(\boldsymbol\phi = (\mathbf{W},\mathbf{b},\log\sigma)\) is updated by SGD.

• The hidden layer learns its own basis functions — no manual feature engineering.
• Width controls expressiveness: too few neurons underfit, too many can overfit.

Three Functions to Infer

Imagine an unknown smooth field \(\theta(x_1,x_2)\) over the plane. We observe noisy samples at scattered locations and want to infer the whole surface. Three targets, easy → hard:

The prior is uniform over the square: training inputs are drawn uniformly, so the network sees the whole domain. All three are pure regression — map \(\mathbb{R}^2\!\to\!\mathbb{R}\), no classification needed.

MLP in Action: Inferring a 2D Function

A single hidden layer mapping \((x_1,x_2)\to\mu\), trained with mini-batch SGD on the MSE loss:

• Colour = inferred value \(\mu(x_1,x_2)\); scattered points are the noisy training targets on the same scale.
• One hidden layer fits the smooth Blob easily, but struggles with the high-frequency Spiral: motivation for going deep.

Stacking Layers

Apply the same idea repeatedly: the output of one layer becomes the input to the next.

Effective Basis Functions by Depth

Take networks with random weights. Plot what the last hidden layer neurons compute as functions of the 1D input:

What a Deep Network Is

The same object, now with \(L\) hidden layers. Depth is just the one-layer rule repeated, a loop in code and a recursion in math:

import torch.nn as nn

class DeepMLP(nn.Module):
    def __init__(self, L=4, H=16):
        super().__init__()
        dims = [2] + [H]*L
        self.layers = nn.ModuleList(
            nn.Linear(dims[l], dims[l+1]) for l in range(L))
        self.g    = nn.ReLU()              # activation g
        self.head = nn.Linear(H, 1)        # W⁽ᴸ⁺¹⁾, b⁽ᴸ⁺¹⁾

    def forward(self, x):
        h = x                              # h⁽⁰⁾ = x
        for layer in self.layers:        # stack L hidden layers
            h = self.g(layer(h))           # h⁽ˡ⁾ = g(W⁽ˡ⁾h⁽ˡ⁻¹⁾ + b⁽ˡ⁾)
        return self.head(h)             # μ = W⁽ᴸ⁺¹⁾h⁽ᴸ⁾ + b⁽ᴸ⁺¹⁾

How Do We Train Deep Networks?

We need \(\nabla_{\mathbf{W}^{(l)}} L\) for every layer. The chain rule propagates gradients backwards:

MLP in Action: Stacking Hidden Layers

The same demo, now with a slider for the number of hidden layers

• On the Spiral, one hidden layer underfits the high-frequency centre; adding layers lets the network compose features and resolve it.
• Depth earns its keep precisely when the target is sharp or structured — the smooth Blob barely needs it.

Key Takeaways

• Learnable basis functions: each neuron is a simple bump or ramp whose shape the network tunes itself, so you never hand-pick features again.
• Universal approximation: one hidden layer, made wide enough, can approximate any continuous function to any accuracy you want.
• No closed form: a learned basis makes the loss non-convex, so we descend it with gradient descent rather than solving it directly.
• Depth buys richness: stacking layers composes simple features into complex ones, so deep nets crack problems a single wide layer struggles with, like the spiral.
• One network, many tasks: changing only the input dimension or the output head turns the same MLP from 1D curve-fitting into 2D surface inference.
• The bridge ahead: swap the hand-built target for a simulator's parameter and data pairs, and this becomes neural posterior estimation in Lecture 3.