Module 7: Neural Networks—Learning Functions from Data

Learning Objectives¶

By the end of this module, you will be able to:

• Explain why neural networks are universal function approximators and what this means for scientific emulation

• Construct a multi-layer perceptron (MLP) architecture from first principles

• Derive how forward propagation computes predictions through layers of transformations

• Connect backpropagation to the automatic differentiation you learned in Module 6

• Implement gradient descent training using loss functions appropriate for regression

• Apply data normalization and recognize why it’s essential for neural network training

• Use ensemble methods to quantify predictive uncertainty

• Build neural network emulators in Equinox and train them with Optax

• Evaluate emulator accuracy and understand failure modes

• Integrate trained emulators into NumPyro for Bayesian inference

7.1: The Emulation Problem¶

Priority: 🔴 Essential

Why Emulators Matter¶

Consider your N-body simulator. Each run takes seconds to minutes. But scientific inference requires thousands of model evaluations—MCMC needs to evaluate the likelihood at every step, optimization needs gradients computed repeatedly. What if we could train a fast approximation that captures the essential physics?

The computational bottleneck: Your final project asks you to infer what initial conditions $(Q_0, a)$ produced a given cluster outcome. Using your actual N-body code inside MCMC would require:

That’s one inference run. Unacceptable.

The emulator solution: Train a neural network on ~100 N-body runs. The network learns to predict outputs from inputs in microseconds. Now inference takes minutes, not months.

Emulators in Modern Astrophysics¶

This isn’t a toy problem. Emulation is transforming computational science:

Cosmology: The Dark Energy Survey uses neural network emulators to predict matter power spectra. Full $N$-body simulations take GPU-hours; emulators return predictions in milliseconds, enabling proper Bayesian parameter estimation.

Stellar evolution: MESA simulations of stellar interiors are expensive. Emulators predict stellar properties (luminosity, radius, lifetime) from initial conditions, enabling population synthesis studies with millions of virtual stars.

Galaxy formation: IllustrisTNG-scale hydrodynamic simulations are computationally prohibitive to run thousands of times. Emulators trained on existing runs enable exploration of the parameter space.

Your final project: You’re learning the same workflow used in frontier research—just on a tractable problem (star clusters instead of cosmological volumes).

7.2: From Linear Regression to Neural Networks¶

Priority: 🟡 Important

Before diving into neural network details, let’s see how they emerge naturally from ideas you already know.

The Progression of Function Approximation¶

Level 1: Linear Regression

The simplest model: $\hat{y} = \mathbf{w}^T \mathbf{x} + b$

You find weights $\mathbf{w}$ and bias $b$ that minimize squared error. This works beautifully when the true relationship is linear—but nature rarely cooperates.

Level 2: Polynomial/Basis Regression

To capture nonlinearity, you engineer features: $\hat{y} = w_0 + w_1 x + w_2 x^2 + w_3 x^3 + \ldots$

This is still linear in the parameters (the $w_i$), but nonlinear in the input. The problem: you must choose which features to include. For your emulator, should you use $Q_0^2$? $\log a$? $Q_0 \times a$? You’d have to guess.

Level 3: Neural Networks

Here’s the key insight: what if we learned the features?

A neural network with one hidden layer computes:

The hidden layer creates learned nonlinear features of the input. The output layer then does linear regression on these features. Training optimizes both the feature extraction and the final regression simultaneously.

Where Does Your Emulator Fit?¶

Your task: learn the mapping $(Q_0, a) \to (f_{\rm bound}, \sigma_v, r_h)$.

Is this relationship linear? Almost certainly not. The physics of gravitational collapse, violent relaxation, and stellar escape involves complex nonlinear dynamics. You could try to engineer features based on virial theorem scaling relations—but why guess when you can learn?

A neural network will discover whatever nonlinear transformations of $(Q_0, a)$ are useful for predicting cluster outcomes. If it turns out linear features suffice, the network can learn that too (it includes linear models as a special case).

7.3: The Computational Neuron¶

Priority: 🔴 Essential

Anatomy of a Neuron¶

The artificial neuron is the atomic unit of neural networks. It takes $d$ inputs $x_1, x_2, \ldots, x_d$ and produces one output:

where:

• $\mathbf{w} = (w_1, \ldots, w_d)$ are weights controlling how much each input matters

• $b$ is a bias allowing the neuron to shift its activation threshold

• $\sigma$ is a nonlinear activation function

Think of it as “how strongly should this neuron fire, given these inputs?” The weights determine sensitivity to each input; the bias sets the baseline; the activation function introduces nonlinearity.

The quantity $z = \mathbf{w}^T \mathbf{x} + b$ is called the pre-activation—the neuron’s response before the nonlinear squashing. The activation function then decides the actual output: $a = \sigma(z)$.

Why Nonlinearity Is Essential¶

Activation Functions¶

Several activation functions are commonly used:

ReLU (Rectified Linear Unit)—the modern default:

Beautifully simple: pass positive values unchanged, zero out negative values. Computationally cheap and works remarkably well. The gradient is 1 for $z > 0$ and 0 for $z < 0$, which helps training (no “vanishing gradients” for positive inputs).

Sigmoid (the “Fermi function” you know from statistical mechanics):

Squashes outputs to $(0, 1)$. Useful for probabilities, but can cause vanishing gradients in deep networks (derivative is small when $|z|$ is large).

Tanh:

Squashes to $(-1, 1)$. Zero-centered, which can help optimization. Same vanishing gradient concern as sigmoid.

For your emulator: Use ReLU for hidden layers. It’s the standard choice for regression and works well in practice.

7.4: Layering Neurons into Networks¶

Priority: 🔴 Essential

The Multi-Layer Perceptron¶

A single neuron has limited expressiveness. The power comes from layering neurons into a network:

Input layer: Your raw features $\mathbf{x} \in \mathbb{R}^{d_{\text{in}}}$. Not really “neurons”—just the data entering the network.

Hidden layer(s): Where the magic happens. Each layer transforms its input through weights, biases, and activations. A hidden layer with $n$ neurons computes:

where $\mathbf{W} \in \mathbb{R}^{n \times d_{\text{in}}}$ is the weight matrix, $\mathbf{b} \in \mathbb{R}^n$ is the bias vector, and $\sigma$ applies element-wise.

Output layer: Produces final predictions $\hat{\mathbf{y}} \in \mathbb{R}^{d_{\text{out}}}$. For regression, typically no activation (linear output).

Your Emulator Architecture¶

For the final project, you’ll build:

• Input: 2 features $(Q_0, a)$

• Hidden 1: 64 neurons, ReLU activation

• Hidden 2: 64 neurons, ReLU activation

• Output: 3 predictions $(f_{\text{bound}}, \sigma_v, r_h)$, linear (no activation)

Counting Parameters¶

How many learnable parameters does your network have? This matters for understanding model complexity.

Layer 1 (input → hidden 1):

• Weight matrix $\mathbf{W}_1 \in \mathbb{R}^{64 \times 2}$: $64 \times 2 = 128$ weights

• Bias vector $\mathbf{b}_1 \in \mathbb{R}^{64}$: 64 biases

• Subtotal: 192 parameters

Layer 2 (hidden 1 → hidden 2):

• Weight matrix $\mathbf{W}_2 \in \mathbb{R}^{64 \times 64}$: $64 \times 64 = 4096$ weights

• Bias vector $\mathbf{b}_2 \in \mathbb{R}^{64}$: 64 biases

• Subtotal: 4,160 parameters

Layer 3 (hidden 2 → output):

• Weight matrix $\mathbf{W}_3 \in \mathbb{R}^{3 \times 64}$: $3 \times 64 = 192$ weights

• Bias vector $\mathbf{b}_3 \in \mathbb{R}^{3}$: 3 biases

• Subtotal: 195 parameters

Total: 4,547 parameters

7.5: The Universal Approximation Theorem¶

Priority: 🟡 Important

The Theoretical Foundation¶

Here’s a remarkable mathematical result that justifies using neural networks as flexible function approximators:

Theorem (Universal Approximation): Let $f: K \to \mathbb{R}$ be any continuous function defined on a compact set $K \subset \mathbb{R}^d$. For any $\epsilon > 0$, there exists a neural network $\hat{f}$ with a single hidden layer and sufficiently many neurons such that:

In words: any continuous function can be approximated arbitrarily well by a neural network with one hidden layer, if you use enough neurons.

What this means for emulation: Your N-body simulator defines a continuous function from initial conditions to summary statistics (assuming the summary statistics vary continuously with inputs, which they do for your problem). The theorem guarantees that some neural network can approximate this mapping to any desired accuracy.

The Intuition: Building Functions from Simple Pieces¶

How can a network approximate arbitrary functions?

For sigmoid activations: Each neuron creates a “soft step”—a smooth transition from low to high output. By positioning many soft steps at different locations with different heights, you can approximate any shape. It’s like building a sculpture from many small clay pieces, each contributing a local bump or dip.

For ReLU activations: Each neuron creates a “hinge”—a bend in the function at some location. A network with $n$ ReLU neurons can create a piecewise linear function with up to $n$ hinges. With enough hinges, piecewise linear functions can approximate any continuous curve.

The Crucial Caveats¶

7.6: Forward Propagation¶

Priority: 🔴 Essential

Computing Predictions¶

Forward propagation is how we compute predictions: apply each layer’s transformation in sequence, feeding outputs forward.

For your 2-hidden-layer network, the forward pass computes:

Each layer takes the previous layer’s output, applies a linear transformation (matrix multiply + bias), and passes through an activation (except the final layer).

The Computational Graph Perspective¶

Forward propagation defines a computational graph—exactly the structure JAX uses for automatic differentiation (Module 6!).

This graph shows exactly how the output $\hat{\mathbf{y}}$ depends on the input $\mathbf{x}$ and all parameters $\boldsymbol{\theta} = (\mathbf{W}_1, \mathbf{b}_1, \mathbf{W}_2, \mathbf{b}_2, \mathbf{W}_3, \mathbf{b}_3)$.

When we train, we need gradients $\nabla_{\boldsymbol{\theta}} \mathcal{L}$. The computational graph tells JAX exactly how to compute them via the chain rule—that’s what jax.grad does automatically.

🔬 Conceptual Checkpoint: Architecture Foundations¶

Before moving to training, verify your understanding:

---

## 7.7: Training as Optimization

**Priority: 🔴 Essential**

### The Loss Function

:::{margin}
**Loss function** (or **cost function**, **objective function**): A scalar function measuring how wrong the network's predictions are. Training minimizes this function.

Training requires a measure of “how wrong” the network is. For regression, the standard choice is mean squared error (MSE):

where $\boldsymbol{\theta}$ denotes all network parameters, $N$ is the number of training examples, and $\hat{\mathbf{y}}_i = f_{\text{network}}(\mathbf{x}_i; \boldsymbol{\theta})$ is the prediction for input $\mathbf{x}_i$.

The loss is a function of the parameters: different weights produce different predictions, which produce different errors. Training finds parameters that minimize the loss.

Gradient Descent¶

We want $\boldsymbol{\theta}^* = \arg\min_{\boldsymbol{\theta}} \mathcal{L}(\boldsymbol{\theta})$. Gradient descent iteratively steps toward the minimum:

where $\eta > 0$ is the learning rate—how large a step we take.

The gradient $\nabla_{\boldsymbol{\theta}} \mathcal{L}$ points in the direction of steepest increase in loss. We move in the opposite direction to decrease the loss.

Backpropagation = Chain Rule = Autodiff¶

Backpropagation computes gradients of the loss with respect to all network parameters. The key insight: it’s just the chain rule applied systematically.

Consider how the loss depends on first-layer weights:

The chain rule propagates gradients backward from the loss through each layer to the parameters.

The Adam Optimizer¶

Plain gradient descent uses a fixed learning rate for all parameters. This can be suboptimal—some parameters benefit from larger steps, others from smaller.

Adam adapts the learning rate per parameter based on gradient history:

• First moment (exponential moving average of gradients): Which direction have we been moving?

• Second moment (exponential moving average of squared gradients): How variable have gradients been?

Parameters with consistently large gradients in the same direction get boosted. Parameters with noisy, inconsistent gradients get damped.

In practice: Adam is the default optimizer for neural networks. It works well across problems with minimal tuning. Start with learning rate $\eta = 10^{-3}$.

Optax provides Adam and many other optimizers.

7.8: Weight Initialization—Breaking Symmetry¶

Priority: 🟡 Important

Why Initialization Matters¶

Before training begins, you must set initial values for all weights and biases. This choice profoundly affects training dynamics.

Scale Matters Too¶

Even with random initialization, the scale of initial weights affects training:

Weights too large: Pre-activations $z = \mathbf{W}\mathbf{x} + \mathbf{b}$ become large. For sigmoid/tanh, this saturates the activation (gradient ≈ 0). For ReLU, values may explode through layers.

Weights too small: Signals shrink as they pass through layers. By the time they reach the output, everything is near zero. Gradients vanish.

Xavier/Glorot initialization (what Equinox uses by default) sets the right scale:

This keeps the variance of activations roughly constant across layers, enabling stable gradient flow.

Why Different Seeds → Different Solutions¶

Neural network loss landscapes are non-convex—they have many local minima and saddle points. Gradient descent finds a minimum, not the global minimum. Which minimum you find depends on where you start.

Different random seeds produce different initial weights → different optimization trajectories → different final solutions.

These solutions may all have similar training loss but differ in their predictions, especially in regions with sparse data. This is the foundation for ensemble uncertainty (Section 7.10): train multiple networks with different seeds and measure how much they disagree.

7.9: Practical Training¶

Priority: 🔴 Essential

Data Normalization¶

Neural networks are sensitive to input/output scales. If $Q_0 \in [0.5, 1.5]$ but $a \in [50, 200]$, gradients with respect to $a$-related weights will dominate, causing slow and unstable learning.

Standardization transforms each feature to zero mean and unit variance:

where $\mu_x$ and $\sigma_x$ are computed from your training set only.

Denormalization: To convert predictions back to physical units:

Multi-Output Normalization¶

Your emulator predicts three quantities: $(f_{\text{bound}}, \sigma_v, r_h)$. The MSE loss sums over all outputs:

Problem: If outputs have different scales, this implicitly weights them unequally. If $r_h$ varies from 20–150 AU (range 130) while $f_{\text{bound}}$ varies from 0.3–1.0 (range 0.7), squared errors for $r_h$ are ~35,000× larger. The network will focus on $r_h$ at the expense of $f_{\text{bound}}$.

Solution: Normalize each output independently to zero mean, unit variance. Then all three contribute roughly equally to the loss.

The normalization pattern (you’ll implement this):

NORMALIZATION PROCEDURE
=======================

1. COMPUTE STATISTICS (from training set only!)
   μ = mean of training data, per feature    # Shape: (3,) for outputs
   σ = std of training data, per feature     # Shape: (3,) for outputs

2. NORMALIZE any data (train, test, or new):
   x̃ = (x - μ) / σ
   
3. DENORMALIZE predictions back to physical units:
   x = x̃ · σ + μ

Critical: Compute μ and σ from training data, then use those same values everywhere—for normalizing test data, for denormalizing predictions, and inside your NumPyro model. Store these statistics alongside your trained model.

Hyperparameters¶

Key hyperparameters for your emulator:

Train/Test Split¶

With limited data (~100 simulations), use an 80/20 split: 80 examples for training, 20 for testing.

Convergence Criteria¶

Training is “done” when:

1. Loss has plateaued: Not decreasing meaningfully over ~100 epochs

2. Loss is reasonably small: MSE $\lesssim 0.01$ on normalized data means typical errors < 10% of standard deviation

3. No pathologies: No NaN values, no loss increasing

The Training Loop¶

A typical training loop:

# Pseudocode
losses = []
for epoch in range(num_epochs):
    # Forward pass: compute predictions
    predictions = model(params, x_train)
    
    # Compute loss
    loss = jnp.mean((predictions - y_train)**2)
    losses.append(loss)
    
    # Backward pass: compute gradients
    gradients = jax.grad(loss_fn)(params)
    
    # Update parameters
    params = update_with_adam(params, gradients)
    
    if epoch % 100 == 0:
        print(f"Epoch {epoch}: loss = {loss:.6f}")

# ALWAYS plot the loss curve
plt.plot(losses)
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.yscale("log")

Always plot the loss curve. Its shape reveals problems that printed numbers hide.

Debugging Checklist¶

7.10: Uncertainty via Ensembles¶

Priority: 🔴 Essential

The Problem with Single Networks¶

A trained neural network gives point predictions—single numbers with no uncertainty. But scientific applications need uncertainty quantification:

• How confident is this prediction?

• Are we extrapolating outside the training distribution?

• How should prediction uncertainty propagate to inference?

A single network provides none of this. Worse, networks can be confidently wrong, especially when extrapolating beyond training data.

Deep Ensembles¶

The simplest approach to neural network uncertainty: train multiple networks with different random seeds.

Why does this work? From Section 7.8: different initializations lead to different local minima. These solutions may give similar predictions where training data is dense but diverge where data is sparse.

Procedure:

1. Train $M$ networks with different random seeds (typically $M = 3$–5)

2. For any input $\mathbf{x}$:

   • Get predictions from each: $\hat{y}_1(\mathbf{x}), \ldots, \hat{y}_M(\mathbf{x})$

   • Mean prediction: $\bar{y}(\mathbf{x}) = \frac{1}{M} \sum_{m=1}^{M} \hat{y}_m(\mathbf{x})$

   • Uncertainty: $\sigma_y(\mathbf{x}) = \sqrt{\frac{1}{M-1} \sum_{m=1}^{M} (\hat{y}_m(\mathbf{x}) - \bar{y}(\mathbf{x}))^2}$

The spread $\sigma_y$ estimates epistemic uncertainty—uncertainty from limited training data.

Interpreting Ensemble Spread¶

Uncertainty should be high:

• Near edges of training distribution (extrapolation begins)

• In gaps between training points

• Where the function is complex/rapidly varying

Uncertainty should be low:

• In regions densely covered by training data

• Where all ensemble members agree

• Where the function is simple/slowly varying

Ensemble in Practice¶

The ensemble workflow is conceptually simple:

ENSEMBLE WORKFLOW
=================

Training:
1. Split random key into M independent keys
2. Train M models, each with a different key
3. Store all trained models

Prediction:
1. Get predictions from all M models
2. Stack predictions into array of shape (M, num_outputs)
3. Mean = jnp.mean(predictions, axis=0)
4. Std = jnp.std(predictions, axis=0, ddof=1)  # ddof=1 for unbiased

Implementation hints:

• jax.random.split(key, M) creates M independent keys

• Store models in a Python list

• For single input: preds = jnp.array([model(x) for model in models])

• For batched inputs: use jax.vmap inside the list comprehension

For your final project, $M = 5$ ensemble members is a good balance between computational cost and uncertainty quality. You’ll implement this yourself—the pattern above tells you what to compute, not how to code it.

🔬 Conceptual Checkpoint: Training and Uncertainty¶

Before moving to implementation, verify your understanding:

---

## 7.11: The JAX Ecosystem—Equinox and Optax

**Priority: 🔴 Essential**

### Equinox: Neural Networks as PyTrees

:::{margin}
**Equinox**: A JAX library for neural networks by Patrick Kidger. Models are PyTrees (nested containers JAX understands), enabling seamless use of `jax.grad`, `jax.jit`, `jax.vmap`.

Equinox makes neural networks fit naturally into JAX’s functional paradigm:

Models are PyTrees: Parameters live in a tree structure. JAX transformations (grad, jit, vmap) work seamlessly because they know how to traverse PyTrees.

Models are callables: An Equinox model is a function. You call it: output = model(input).

Explicit state: No hidden mutable variables. All parameters are explicit. This makes debugging easier and enables JAX’s transformations.

Learning from Documentation¶

Glass-Box Exercise: Manual Forward Pass¶

Before using any library, make sure you understand what’s happening inside:

import jax.numpy as jnp

def manual_forward(params, x):
    """Forward pass through a 2-64-64-3 MLP.
    
    params: dict with keys 'W1', 'b1', 'W2', 'b2', 'W3', 'b3'
    x: input array of shape (2,)
    Returns: output array of shape (3,)
    """
    # Layer 1: linear transformation + ReLU
    z1 = params['W1'] @ x + params['b1']
    h1 = jnp.maximum(0, z1)
    
    # Layer 2: linear transformation + ReLU
    z2 = params['W2'] @ h1 + params['b2']
    h2 = jnp.maximum(0, z2)
    
    # Output layer: linear only (no activation)
    y = params['W3'] @ h2 + params['b3']
    return y

def loss_fn(params, x, y_true):
    y_pred = manual_forward(params, x)
    return jnp.mean((y_pred - y_true)**2)

grad_fn = jax.grad(loss_fn)  # This works!

A Teaching Example: Learning $\sin(x)$¶

To illustrate Equinox patterns without solving your final project, let’s build a tiny network that learns the sine function. This is NOT your emulator architecture—adapt the concepts, don’t copy the code.

import jax
import jax.numpy as jnp
import equinox as eqx
import optax

# A SIMPLE 1D example - your emulator will differ!
class TinySinNet(eqx.Module):
    layer1: eqx.nn.Linear
    layer2: eqx.nn.Linear
    
    def __init__(self, key):
        k1, k2 = jax.random.split(key)
        self.layer1 = eqx.nn.Linear(1, 16, key=k1)  # 1 input
        self.layer2 = eqx.nn.Linear(16, 1, key=k2)  # 1 output
    
    def __call__(self, x):
        x = jax.nn.relu(self.layer1(x))
        return self.layer2(x)

Key patterns to notice (these transfer to your emulator):

1. Inherit from eqx.Module — makes your class a PyTree

2. Declare layers as typed attributes — Equinox tracks these

3. Split keys for each layer — ensures independent initialization

4. __call__ defines forward pass — apply layers and activations

Essential Equinox Patterns¶

These small patterns are the building blocks. You’ll combine them for your emulator.

Pattern 1: Creating layers

# eqx.nn.Linear(in_features, out_features, key=...)
layer = eqx.nn.Linear(64, 32, key=some_key)

Pattern 2: Applying activations

# Activations are just JAX functions
x = jax.nn.relu(layer(x))  # ReLU
x = jax.nn.tanh(layer(x))  # Tanh (if you needed it)

Pattern 3: Filtering arrays from models

# Equinox models contain arrays (parameters) and non-arrays (structure)
# Many operations need just the arrays:
params = eqx.filter(model, eqx.is_array)

Pattern 4: Computing gradients

# eqx.filter_value_and_grad works like jax.value_and_grad
# but handles models with non-array components
loss, grads = eqx.filter_value_and_grad(loss_fn)(model)

Pattern 5: Applying updates

# After computing updates from optimizer:
new_model = eqx.apply_updates(model, updates)

Pattern 6: JIT-compiling with models

# Use eqx.filter_jit instead of jax.jit for functions involving eqx.Module
@eqx.filter_jit
def train_step(model, ...):
    ...

The Training Loop Structure¶

Here’s the structure of a training loop—you’ll fill in the details:

TRAINING LOOP PSEUDOCODE
========================

1. INITIALIZE
   - Create model with random key
   - Create optimizer (e.g., optax.adam)
   - Initialize optimizer state from model parameters

2. FOR EACH EPOCH:
   a. FORWARD PASS
      - Compute predictions: pred = model(x_train)  
      - (Use jax.vmap to handle batches)
   
   b. COMPUTE LOSS
      - MSE: loss = mean((pred - y_train)²)
   
   c. BACKWARD PASS  
      - Compute gradients w.r.t. model parameters
      - (Use eqx.filter_value_and_grad)
   
   d. UPDATE
      - Get updates from optimizer
      - Apply updates to model
   
   e. LOG
      - Store loss for plotting
      - Print progress periodically

3. PLOT LOSS CURVE (always!)

4. RETURN trained model

Your task: Translate this pseudocode into working Equinox/Optax code. The documentation tutorials show exactly how.

Training an Ensemble: The Idea¶

For uncertainty quantification, you’ll train multiple networks:

ENSEMBLE PSEUDOCODE
===================

1. Split your random key into M keys (one per model)

2. FOR EACH key:
   - Train a model using that key
   - (Different key → different initialization → different solution)
   - Store the trained model

3. TO PREDICT with uncertainty:
   - Get prediction from each model
   - Mean = average of predictions
   - Uncertainty = standard deviation of predictions

Implementation hints:

• jax.random.split(key, n) gives you n independent keys

• Store models in a Python list

• For ensemble prediction, iterate over models and stack results with jnp.array([...])

• Use jnp.mean(..., axis=0) and jnp.std(..., axis=0, ddof=1) to aggregate

7.12: From Emulator to Inference¶

Priority: 🔴 Essential

The Goal¶

You’ve trained a fast emulator. Now use it for Bayesian inference: given observed cluster properties, infer the initial conditions that produced them.

The emulator replaces the expensive N-body simulation inside your inference loop. NumPyro’s NUTS sampler will call it thousands of times—this is why speed matters.

The Conceptual Framework¶

In NumPyro, you write a generative model—code describing how observations arise from parameters. The structure directly mirrors Bayes’ theorem:

GENERATIVE MODEL STRUCTURE
==========================

1. PRIOR: What do we believe about parameters before seeing data?
   - Sample Q₀ from some distribution (e.g., Uniform over training range)
   - Sample a from some distribution
   
2. FORWARD MODEL: Given parameters, what do we predict?
   - Normalize the sampled (Q₀, a) using TRAINING statistics
   - Pass through your emulator to get predicted (f_bound, σᵥ, rₕ)
   - Denormalize predictions back to physical units
   
3. LIKELIHOOD: How probable are actual observations given predictions?
   - Compare predictions to observations
   - Account for observation uncertainty (σ_obs)

The key insight: Your emulator serves as the forward model. Instead of running an expensive N-body simulation for each $(Q_0, a)$ the sampler proposes, you call your fast neural network.

What You Need to Figure Out¶

The Normalization Challenge¶

Your emulator was trained on normalized data. Your NumPyro model must handle this correctly:

NORMALIZATION FLOW
==================

Physical parameters (Q₀, a) from prior
          ↓
    Normalize using TRAINING statistics
          ↓
Normalized inputs to emulator
          ↓
    Emulator predicts normalized outputs
          ↓
    Denormalize using TRAINING statistics
          ↓
Physical predictions (f_bound, σᵥ, rₕ)
          ↓
    Compare to physical observations in likelihood

Choosing Observation Uncertainty¶

The likelihood requires $\sigma_{\rm obs}$—how uncertain are our “observations”? For synthetic observations from your simulations, this represents:

1. Emulator error: The network doesn’t perfectly predict simulations

2. Stochastic variation: Different random seeds produce scatter even for fixed $(Q_0, a)$

Practical approach: Use your test-set RMSE for each output.

COMPUTING σ_obs
===============

1. Get emulator predictions on test set (normalized)
2. Denormalize to physical units
3. Compare to true test values
4. RMSE = sqrt(mean((pred - true)²)) for each output
5. Use these three RMSE values as σ_obs

Validation: Parameter Recovery¶

The critical test of your inference pipeline:

PARAMETER RECOVERY TEST
=======================

1. Pick a simulation from your TEST set
   - You know the true (Q₀, a) that generated it
   - You have its summary statistics (f_bound, σᵥ, rₕ)

2. Treat summary statistics as "observations"

3. Run inference to get posterior over (Q₀, a)

4. Check: Do 95% credible intervals contain true values?
   - Compute 2.5th and 97.5th percentiles of posterior samples
   - True value should fall within this range
   
5. Repeat for several test cases
   - If ~95% contain truth, your pipeline is well-calibrated
   - If fewer, σ_obs may be too small (overconfident)
   - If more, σ_obs may be too large (uninformative)

Visualization: Create corner plots showing the joint posterior over $(Q_0, a)$ with true values marked. The true values should fall within the 95% contours.

Connecting the Pieces¶

Here’s how all the components fit together in your final project:

COMPLETE INFERENCE PIPELINE
===========================

┌─────────────────────────────────────────────────────────┐
│  TRAINING PHASE (done before inference)                 │
├─────────────────────────────────────────────────────────┤
│  1. Generate N-body training data                       │
│  2. Compute and store normalization statistics          │
│  3. Train ensemble of emulators                         │
│  4. Evaluate on test set → get σ_obs                    │
└─────────────────────────────────────────────────────────┘
                          ↓
┌─────────────────────────────────────────────────────────┐
│  INFERENCE PHASE                                        │
├─────────────────────────────────────────────────────────┤
│  1. Define NumPyro model:                               │
│     - Priors on (Q₀, a)                                 │
│     - Emulator as forward model (with normalization!)   │
│     - Gaussian likelihood with σ_obs                    │
│                                                         │
│  2. Run NUTS:                                           │
│     - Warmup: adapts step size and mass matrix          │
│     - Sampling: collects posterior samples              │
│                                                         │
│  3. Analyze posterior:                                  │
│     - Summary statistics (mean, std, credible intervals)│
│     - Corner plots                                      │
│     - Parameter recovery validation                     │
└─────────────────────────────────────────────────────────┘

7.13: Synthesis¶

Priority: 🔴 Essential

The Complete Workflow¶

You’ve learned a powerful workflow for modern computational science:

Connections Across the Course¶

What Makes a Good Emulator?¶

Accuracy: Predictions match simulations in the training region

• Test RMSE should be small relative to output variation

• Predicted vs. true plots cluster along $y = x$

Calibration: Uncertainty estimates are meaningful

• ~95% of true values fall within ±2σ of predictions

• Uncertainty increases where training data is sparse

Speed: Fast enough for inference

• Microseconds per prediction (vs. seconds for simulation)

• JIT compilation is essential

Robustness: Sensible behavior at boundaries

• No wild predictions just outside training range

• Ensemble uncertainty increases appropriately

Looking Ahead: The Final Project¶

You’re now ready for the final project. You’ll:

1. Generate training data: Latin Hypercube sampling in $(Q_0, a)$ space; compute summary statistics from N-body simulations using your Project 5 code

2. Build and train an emulator: MLP in Equinox; train ensemble with Optax

3. Evaluate accuracy: Test metrics, predicted vs. true plots, uncertainty behavior

4. Perform inference: NumPyro model with emulator as forward model; NUTS sampling; corner plots showing parameter recovery

This workflow—expensive simulations → fast surrogates → Bayesian inference—is exactly how frontier research operates. You’re learning it from the inside out.

🔬 Final Conceptual Checkpoint¶

Before starting the final project, ensure you can answer:

Further Reading and Resources¶

Essential (Complete Before Final Project)¶

• Equinox documentation: https://docs.kidger.site/equinox/

  • “Getting Started”

  • “Train a Neural Network” example

• Optax documentation: https://optax.readthedocs.io/

  • “Quick Start” tutorial

• NumPyro documentation: https://num.pyro.ai/

  • “Getting Started”

  • “Bayesian Regression” example

For Deeper Understanding¶

• Nielsen, Michael: Neural Networks and Deep Learning

  • Free online: http://neuralnetworksanddeeplearning.com/

  • Excellent intuitive explanations

• Goodfellow, Bengio, Courville: Deep Learning (2016)

  • Free online: https://www.deeplearningbook.org/

  • Chapters 6–8 cover feed-forward networks

• Mehta et al.: “A high-bias, low-variance introduction to Machine Learning for physicists” (Physics Reports, 2019)

  • https://arxiv.org/abs/1803.08823

  • Machine learning from a physicist’s perspective

Emulation in Astrophysics¶

• Cranmer et al. (2020): “The frontier of simulation-based inference”

  • https://arxiv.org/abs/1911.01429

  • Overview of ML for scientific simulation

• Alsing et al. (2019): “Fast likelihood-free cosmology with neural density estimators and active learning”

  • Example of emulator-based inference in cosmology

Now go build something that learns. 🚀

Next: Proceed to Final Project, or continue with Gaussian Processes (optional advanced topic)