Part II: Gaussian Processes - Foundations and Theory

Learning Outcomes¶

By the end of Part II, you will be able to:

• Articulate why surrogate modeling is essential for modern computational astrophysics and identify when emulation is appropriate

• Explain Gaussian Processes as probability distributions over functions and derive predictions from Gaussian conditioning

• Decide when to use GPs versus neural networks versus other approaches based on problem characteristics

• Interpret kernel functions as encoding physical assumptions about smoothness, periodicity, and structure

• Distinguish epistemic from aleatoric uncertainty and explain how GPs quantify both

• Connect GP regression to the Bayesian inference framework from Module 5, seeing it as infinite-dimensional parameter inference

• Apply GP concepts to emulate N-body star cluster simulations, predicting cluster evolution from initial conditions

The Big Picture: The Computational Crisis in Modern Astrophysics¶

The Problem We’re Solving¶

You’ve spent the semester building physics from first principles:

Module 1-2: Statistical Foundations

• Central Limit Theorem, Maximum Entropy, moments as information compression

• Statistical mechanics: from microstates to macrostates

• The Boltzmann equation: evolution of distribution functions

Module 3: Dynamics in Phase Space

• Collisionless Boltzmann equation for stellar systems

• Virial theorem: when systems reach equilibrium

• Two-body relaxation: how clusters dissolve

Module 4: Radiative Transfer

• Photon transport through stellar atmospheres

• Monte Carlo methods for solving integral equations

• Importance sampling: making rare events tractable

Module 5: Bayesian Inference

• Beliefs as probability distributions over parameters

• MCMC (Markov Chain Monte Carlo): sampling complex posteriors when no closed form exists

• Model comparison via marginal likelihoods

Projects 2, 3, 5: Building the Machinery

• N-body simulations of star cluster evolution (gravity from scratch)

• Monte Carlo radiative transfer (photon-by-photon tracking)

• JAX implementations (fast, differentiable, GPU-accelerated)

  • JAX: Just-In-Time compiled array computing for Python with automatic differentiation

  • GPU: Graphics Processing Unit (massively parallel hardware accelerator)

  • JIT: Just-In-Time compilation (jit) for speed

Each simulator encodes centuries of physics compressed into a few hundred lines of code. They work beautifully. They give correct answers.

They’re also prohibitively slow.

Output: Cluster properties at time $t = 10 \, t_{\rm dyn}$, where the dynamical time is defined as $t_{\rm dyn} \equiv \sqrt{R_h^3 / (G M_{\rm tot})}$ (the characteristic timescale for orbital motion at the half-mass radius)

Note on Plummer profiles: For the Plummer density profile you’re using, the half-mass radius relates to the scale radius as $R_h \approx 1.305 a$ (the exact factor depends on definitions and conventions; verify in your simulator’s documentation). So clusters with larger $a$ have longer dynamical times (slower evolution): $t_{\rm dyn} \propto a^{3/2} / \sqrt{G M_{\rm tot}}$. This means $a$ directly controls the evolutionary timescale—important for interpreting your emulator’s learned lengthscale $\ell_a$!

• Bound fraction: what percentage of stars remain gravitationally bound?

• Core radius $R_{\rm core}$: how compact is the cluster center?

• Half-mass radius $R_h$: characteristic size

• Velocity dispersion $\sigma_v$: kinetic energy indicator

Computational Cost:

• Single simulation: ~10 minutes (1000 particles, 100 timesteps, direct N-body)

• With your JAX implementation: ~2 minutes (GPU acceleration, jit compilation)

Now suppose you want to answer scientific questions:

1. Parameter Space Exploration: “Which initial conditions lead to cluster survival vs dissolution?”

   • Need: ~5,000 simulations to sample 3D space densely (roughly $\sim 17$ points per dimension for $17^3 \approx 5{,}000$ total)

   • Cost: 5,000 × 2 min = ~167 hours = 7 days of compute

2. Bayesian Parameter Inference: “What initial conditions produced the Pleiades cluster?”

   • Need: ~100,000 MCMC samples × likelihood evaluations

   • Cost: 100,000 × 2 min = ~3,333 hours = 138 days = 4.5 months

3. Optimization: “What initial setup maximizes cluster lifetime?”

   • Need: ~50,000 gradient evaluations for optimization

   • Cost: 50,000 × 2 min = ~1,666 hours = 69 days = 2.3 months

This is not a toy problem. This is real research.

### The Traditional Approach: Suffer or Approximate

Historically, astronomers faced a choice:

**Option 1: Just Wait**

- Run the expensive simulations
- Use supercomputers, parallelize, optimize
- Wait weeks/months for results
- **Problem**: Kills iteration speed, makes exploration impossible

**Option 2: Use Simple Models**

- Fit polynomials, power laws, linear regression
- Fast to evaluate, easy to interpret
- **Problem**: Physics is nonlinear, complex, high-dimensional. Simple models fail.

**Option 3: Dimensional Reduction**

- Derive simplified analytic approximations (e.g., virial theorem $Q = 1$ for equilibrium)
- Use scaling relations (e.g., $t_{\rm relax} \propto N / \ln N$)
- **Problem**: Valid only in limited regimes, miss interesting dynamics (unless validated with new simulations)

**Option 4: Give Up on Some Questions**

- Don't do parameter inference, don't optimize, don't explore broadly
- **Problem**: Science suffers

### The Modern Solution: Emulation

**Key Insight**: Your simulator $f_{\rm sim}$ is deterministic (or nearly so). Once you fix initial conditions $\mathbf{x}$, the output $y$ is determined by physics:

$$
y = f_{\rm sim}(\mathbf{x}) + \epsilon
$$

where $\epsilon$ represents numerical noise (finite timesteps, roundoff errors, stochastic elements).

**The Emulation Strategy**:

1. **One-Time Investment**: Run $N = 200\text{-}500$ simulations with carefully chosen initial conditions
   - Cost: $500 \times 2 \text{ min} \approx 16 \text{ hours}$ (one night of compute)
   - Generate training dataset: $\mathcal{D} = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$

2. **Learn the Function**: Train a *surrogate model* (emulator) $f_{\rm emu}$ that approximates $f_{\rm sim}$
   - Cost: Minutes to hours (one-time training)
   - Result: Fast function $f_{\rm emu}(\mathbf{x})$ that mimics $f_{\rm sim}(\mathbf{x})$

3. **Deploy Everywhere**: Use $f_{\rm emu}$ instead of $f_{\rm sim}$ for:
   - Parameter space exploration: evaluate $f_{\rm emu}$ at millions of points
   - MCMC sampling: cheap likelihood evaluations
   - Optimization: fast gradient computation (if emulator is differentiable)
   - **Speedup**: $10^4$ to $10^6$ times faster than simulator

4. **Quantify Trust**: Emulator provides **uncertainty estimates**
   - "I predict bound fraction = $0.65 \pm 0.08$"
   - Know when to trust predictions, when to run more simulations
   - Principled decision-making about computational budget

:::{admonition} Connection to Module 5: Bayesian Inference
:class: note

In Project 4, you used MCMC to sample posterior distributions:

$$
p(\theta | \mathcal{D}) \propto p(\mathcal{D} | \theta) \, p(\theta)
$$

Each MCMC step required evaluating the likelihood $p(\mathcal{D} | \theta)$. When your model was simple (e.g., Gaussian with unknown mean/variance), this was cheap—just compute a probability.

But what if the likelihood requires *running a simulation*?

$$
p(\mathcal{D} | \theta) = p(\text{observed cluster properties} | \text{initial conditions } \theta)
$$

Then each MCMC step costs 2 minutes, and 100,000 samples = 138 days. **Intractable.**

**Emulation makes Bayesian inference feasible**: Replace $f_{\rm sim}(\theta)$ with $f_{\rm emu}(\theta)$ in the likelihood. Now each MCMC step costs milliseconds, and 100,000 samples = 1 hour.

This workflow is called **simulation-based inference** or **likelihood-free inference**, and it's how modern experiments (LIGO, DESI, Planck) constrain cosmological and astrophysical parameters.

[FIGURE 1.1: Emulation Workflow]

First Principles: Why Gaussian Processes Work¶

The Core Idea: Functions as Infinite-Dimensional Vectors¶

You’re familiar with probability distributions over scalars and vectors:

Scalar random variable: $x \sim \mathcal{N}(\mu, \sigma^2)$

• Describes uncertainty about a single number

• Specified by mean $\mu$ and variance $\sigma^2$

Vector random variable: $\mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$

• Describes uncertainty about $n$ numbers jointly

• Specified by mean vector $\boldsymbol{\mu} \in \mathbb{R}^n$ and covariance matrix $\boldsymbol{\Sigma} \in \mathbb{R}^{n \times n}$

A Gaussian Process extends this to functions:

• Describes uncertainty about an entire function $f: \mathbb{R}^D \to \mathbb{R}$

• Specified by:

  • Mean function $m: \mathbb{R}^D \to \mathbb{R}$ (often $m(\mathbf{x}) = 0$ after centering data)

  • Kernel function $k: \mathbb{R}^D \times \mathbb{R}^D \to \mathbb{R}$ (covariance between function values)

Definition (Gaussian Process):

A stochastic process $f(\mathbf{x})$ is a Gaussian Process if, for any finite collection of input points $\{\mathbf{x}_1, \ldots, \mathbf{x}_n\}$, the joint distribution of function values is a multivariate Gaussian:

Intuition:

• A function is an infinite-dimensional object: it assigns a value to every point in input space

• A GP says: “I don’t know the exact function, but I have beliefs about what it looks like”

• Those beliefs are encoded in the kernel $k(\mathbf{x}, \mathbf{x}')$: “how similar should $f(\mathbf{x})$ and $f(\mathbf{x}')$ be?”

From Finite to Infinite: The Consistency Requirement¶

Why use Gaussians for functions? Two reasons:

1. Computational Tractability

Gaussian distributions have magical properties (recall Module 1):

• Marginals are Gaussian (if you ignore some dimensions, what’s left is still Gaussian)

• Conditionals are Gaussian (if you observe some values, predictions for others are Gaussian)

• Linear transformations are Gaussian

These properties mean we can compute predictions in closed form — no MCMC, no numerical integration, just linear algebra.

2. Consistency Under Marginalization

Suppose I specify a GP prior $f \sim \mathcal{GP}(0, k)$. Then by definition:

But also:

Consistency: The distribution over $[f(\mathbf{x}_1), f(\mathbf{x}_2)]$ obtained by marginalizing out $f(\mathbf{x}_3)$ from the 3D Gaussian must match the 2D Gaussian. Gaussians satisfy this automatically.

This is called the Kolmogorov consistency theorem, and it’s why GPs are mathematically well-defined even though functions are infinite-dimensional.

🔴 The Mathematical Machinery: Gaussian Conditioning¶

Reviewing Multivariate Gaussians¶

Before we predict with GPs, we need the key formula from Module 1: Gaussian conditioning.

Suppose we have a joint Gaussian distribution over two sets of variables:

where:

• $\mathbf{f}_1 \in \mathbb{R}^n$ are observed variables (training data)

• $\mathbf{f}_2 \in \mathbb{R}^m$ are unobserved variables (test points)

• $\boldsymbol{\Sigma}_{12} = \boldsymbol{\Sigma}_{21}^T$ (covariance is symmetric)

Question: If we observe $\mathbf{f}_1 = \mathbf{y}$ (actual data), what is the distribution of $\mathbf{f}_2$?

Theorem (Gaussian Conditioning):

The conditional distribution is Gaussian:

where:

This is the entire machinery of Gaussian Process regression.

Interpreting the Formulas¶

Let’s unpack these equations physically:

Posterior Mean $\boldsymbol{\mu}_{2|1} = \boldsymbol{\mu}_2 + \boldsymbol{\Sigma}_{21} \boldsymbol{\Sigma}_{11}^{-1} (\mathbf{y} - \boldsymbol{\mu}_1)$

• Prior guess: $\boldsymbol{\mu}_2$ (what we’d predict before seeing data)

• Residual: $(\mathbf{y} - \boldsymbol{\mu}_1)$ (how much the observations differ from prior expectations)

• Regression weights: $\boldsymbol{\Sigma}_{21} \boldsymbol{\Sigma}_{11}^{-1}$ (how much each observation influences the prediction)

• Update: Add weighted combination of residuals to prior

Physical interpretation: If test point $\mathbf{x}_*$ is highly correlated with training point $\mathbf{x}_i$ (large $\Sigma_{21}$), then observing $y_i$ strongly influences prediction at $\mathbf{x}_*$.

Posterior Covariance $\boldsymbol{\Sigma}_{2|1} = \boldsymbol{\Sigma}_{22} - \boldsymbol{\Sigma}_{21} \boldsymbol{\Sigma}_{11}^{-1} \boldsymbol{\Sigma}_{12}$

• Prior uncertainty: $\boldsymbol{\Sigma}_{22}$ (uncertainty before seeing data)

• Reduction: $\boldsymbol{\Sigma}_{21} \boldsymbol{\Sigma}_{11}^{-1} \boldsymbol{\Sigma}_{12}$ (how much observations reduce uncertainty)

• Key insight: Posterior covariance does not depend on observed values $\mathbf{y}$!

Physical interpretation: Uncertainty depends only on where you observed data (the $\mathbf{x}_i$), not on what you observed (the $y_i$). This makes sense: uncertainty is about information geometry, not values.

From Conditioning to GP Prediction¶

Now apply this to Gaussian Processes. We have:

Training data: $\mathcal{D} = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$ where $y_i = f(\mathbf{x}_i) + \epsilon_i$

• Collected input-output pairs from expensive simulator

• Noise $\epsilon_i \sim \mathcal{N}(0, \sigma_n^2)$ (numerical errors, stochasticity)

GP prior: $f \sim \mathcal{GP}(m, k)$

• Before seeing data, we believe $f$ has mean function $m$ and kernel $k$

• For simplicity, often take $m(\mathbf{x}) = 0$ (data can be centered)

Goal: Predict $f(\mathbf{x}_*)$ at new test point $\mathbf{x}_*$

Setup: Form joint distribution over training outputs and test output:

where:

• $\mathbf{y} = [y_1, \ldots, y_N]^T$ (training outputs, observed)

• $f_* = f(\mathbf{x}_*)$ (test output, unobserved)

• $\mathbf{K} \in \mathbb{R}^{N \times N}$ with $K_{ij} = k(\mathbf{x}_i, \mathbf{x}_j)$ (training-training covariance)

• $\mathbf{k}_* \in \mathbb{R}^N$ with $[\mathbf{k}_*]_i = k(\mathbf{x}_i, \mathbf{x}_*)$ (training-test covariance)

• $k_{**} = k(\mathbf{x}_*, \mathbf{x}_*)$ (test-test covariance, usually 1 for normalized kernels)

• $\sigma_n^2 \mathbf{I}$ accounts for observation noise

Apply Gaussian conditioning: Condition on observed $\mathbf{y}$ to get posterior over $f_*$:

These are the fundamental equations of GP regression.

What These Equations Mean¶

Let’s interpret each term:

Predictive Mean: $\mu(\mathbf{x}_*) = m(\mathbf{x}_*) + \mathbf{k}_*^T (\mathbf{K} + \sigma_n^2 \mathbf{I})^{-1} (\mathbf{y} - \mathbf{m})$

1. Prior mean $m(\mathbf{x}_*)$: Default guess (often zero)

2. Data term $(\mathbf{y} - \mathbf{m})$: How training outputs differ from prior expectations

3. Weights $(\mathbf{K} + \sigma_n^2 \mathbf{I})^{-1} (\mathbf{y} - \mathbf{m})$: How much influence each training point has (depends on noise and correlations)

4. Similarity $\mathbf{k}_*$: How similar test point is to each training point

5. Weighted average: Prediction is weighted sum of training residuals, weighted by similarity

Predictive Variance: $\sigma^2(\mathbf{x}_*) = k_{**} - \mathbf{k}_*^T (\mathbf{K} + \sigma_n^2 \mathbf{I})^{-1} \mathbf{k}_*$

1. Prior variance $k_{**}$: Uncertainty before seeing data (usually normalized to 1)

2. Reduction $\mathbf{k}_*^T (\mathbf{K} + \sigma_n^2 \mathbf{I})^{-1} \mathbf{k}_*$: How much training data reduces uncertainty

3. Key property: Variance is small (confident) when test point is similar to training points (large $\mathbf{k}_*$)

4. Another key property: Variance is large (uncertain) far from training data

🔴 Kernels: Encoding Physical Intuition¶

What Kernels Do¶

The kernel function $k: \mathbb{R}^D \times \mathbb{R}^D \to \mathbb{R}$ is the heart of the GP. It encodes all your prior beliefs about the function you’re learning.

Definition: The kernel $k(\mathbf{x}, \mathbf{x}')$ measures the covariance between function values at inputs $\mathbf{x}$ and $\mathbf{x}'$:

Interpretation:

• Large $k(\mathbf{x}, \mathbf{x}')$ → knowing $f(\mathbf{x})$ tells you a lot about $f(\mathbf{x}')$ (strongly correlated)

• Small $k(\mathbf{x}, \mathbf{x}')$ → knowing $f(\mathbf{x})$ tells you little about $f(\mathbf{x}')$ (weakly correlated)

• $k(\mathbf{x}, \mathbf{x}) = \sigma_f^2$ → prior variance of $f$ at any point (signal variance)

Key Properties (for a valid kernel):

1. Symmetry: $k(\mathbf{x}, \mathbf{x}') = k(\mathbf{x}', \mathbf{x})$ (covariance is symmetric)

2. Positive Semi-Definite: For any set of points $\{\mathbf{x}_i\}$, the matrix $K_{ij} = k(\mathbf{x}_i, \mathbf{x}_j)$ must have non-negative eigenvalues

3. Stationarity (common but not required): Many kernels are stationary, $k(\mathbf{x}, \mathbf{x}') = k(\mathbf{x} - \mathbf{x}')$ (depends only on distance, not absolute location). Non-stationary kernels can model varying lengthscales (e.g., input warping, Gibbs kernel).

The Squared Exponential (RBF) Kernel¶

The most common kernel is the Squared Exponential (SE), also called Radial Basis Function (RBF):

Hyperparameters:

• $\sigma_f^2$ : Signal variance (overall amplitude of function, units: output$^2$)

• $\ell$ : Lengthscale (how far you must move in input space for function to change significantly, units: input)

Anisotropic version (Automatic Relevance Determination - ARD):

When dimensions have different importance, use per-dimension lengthscales:

where each $\ell_d$ controls sensitivity to dimension $d$. Small $\ell_d$ → high sensitivity, large $\ell_d$ → low sensitivity.

Physical interpretation:

1. Exponential decay: Correlation decreases exponentially with distance $\|\mathbf{x} - \mathbf{x}'\|$

2. Lengthscale: If $\|\mathbf{x} - \mathbf{x}'\| = \ell$, then $k = \sigma_f^2 e^{-1/2} \approx 0.6 \sigma_f^2$ (60% correlated)

3. Short lengthscale: Function varies rapidly, needs data points close together

4. Long lengthscale: Function varies slowly, data points far apart still informative

5. Infinite smoothness: Functions drawn from SE GP are infinitely differentiable (unrealistic for many physical systems!)

Visualization: Imagine bound fraction vs virial ratio $Q$:

• Small $\ell = 0.01$: Bound fraction changes dramatically between $Q = 0.50$ and $Q = 0.51$ (highly sensitive)

• Large $\ell = 0.5$: Bound fraction changes smoothly across entire range $Q \in [0.3, 0.7]$ (insensitive)

The Matérn Family: More Realistic Smoothness¶

The SE kernel assumes infinitely smooth functions. Real physics is often rougher. The Matérn family allows tunable smoothness:

where:

• $\nu > 0$ : Smoothness parameter

• $K_\nu$ : Modified Bessel function of the second kind

• $\Gamma(\nu)$ : Gamma function

Key cases:

1. $\nu = 1/2$ (Matérn-1/2):

   • Equivalent to Ornstein-Uhlenbeck process

   • Functions are continuous but not differentiable (rough, can have kinks)

   • With $r = \|\mathbf{x} - \mathbf{x}'\|$:

     $$k(\mathbf{x}, \mathbf{x}') = \sigma_f^2 \exp(-r / \ell)$$

     (30)

   • Note: Still continuous! Doesn’t model true discontinuities, just permits roughness.

2. $\nu = 3/2$ (Matérn-3/2):

   • Functions are once differentiable (smooth but can have kinks in second derivative)

   • With $r = \|\mathbf{x} - \mathbf{x}'\|$:

     $$k(\mathbf{x}, \mathbf{x}') = \sigma_f^2 \left(1 + \sqrt{3} \frac{r}{\ell}\right) \exp\left(-\sqrt{3} \frac{r}{\ell}\right)$$

     (31)

3. $\nu = 5/2$ (Matérn-5/2):

   • Functions are twice differentiable (smoother, often good default)

   • With $r = \|\mathbf{x} - \mathbf{x}'\|$:

     $$k(\mathbf{x}, \mathbf{x}') = \sigma_f^2 \left(1 + \sqrt{5} \frac{r}{\ell} + \frac{5r^2}{3\ell^2}\right) \exp\left(-\sqrt{5} \frac{r}{\ell}\right)$$

     (32)

4. $\nu \to \infty$: Recovers SE kernel (infinitely smooth)

When to use which:

• Matérn-5/2 (default): Start here for most physical systems (smooth but not unrealistically so)

• Matérn-3/2: If Matérn-5/2 is over-smooth (under-coverage in calibration, residuals show structure)

• SE: When you know function is truly smooth (interpolating well-behaved data)

• Matérn-1/2: When function is rough/non-smooth, but still continuous (not for true discontinuities - GPs struggle with those)

Practical tuning guide:

• If predictions are under-covered (empirical coverage < 68% at 1σ) → try lower $\nu$ (Matérn-3/2 instead of 5/2) or regularize $\sigma_n^2$

• If predictions are numerically wiggly → increase $\nu$ (higher smoothness)

• If lengthscales blow up → kernel may be too flexible, try Matérn-5/2 or SE

Periodic Kernels: Exploiting Symmetries¶

If your function has periodic structure, you can encode this:

Note: This is the 1D periodic kernel. For vector inputs, periodicity is typically applied per dimension: $k_{\text{Per}}(\mathbf{x}, \mathbf{x}') = \prod_d \sigma_f^2 \exp\left(-\frac{2\sin^2(\pi(x_d - x_d')/p_d)}{\ell_d^2}\right)$. Use only if you have strong physical reason to expect periodicity in a specific dimension (e.g., orbital angular position, stellar rotation phase).

where:

• $p$ : Period (units: input)

• $\ell$ : Lengthscale within period (controls smoothness)

When to use:

• Time series with known periodicity (stellar rotation, orbital period)

• Spatial data with translational symmetry

N-body example: If you’re emulating cluster properties in a rotating frame, periodic kernel might capture rotational symmetry. (Rare, but possible!)

Compositional Kernels: Building Complexity¶

You can combine kernels to build more complex priors:

1. Addition $k(\mathbf{x}, \mathbf{x}') = k_1(\mathbf{x}, \mathbf{x}') + k_2(\mathbf{x}, \mathbf{x}')$

• Functions are sums of two independent processes

• Example: Smooth trend + noise, Long-range + short-range structure

2. Multiplication $k(\mathbf{x}, \mathbf{x}') = k_1(\mathbf{x}, \mathbf{x}') \cdot k_2(\mathbf{x}, \mathbf{x}')$

• Functions have coupled structure

• Example: Periodic function with varying amplitude

3. Per-dimension kernels:

• Different structure in each dimension

• Encodes independence assumptions

🔴 Decision Framework: When to Use What¶

You now understand what GPs are and how they work. But when should you actually use them?

The Method Landscape¶

Modern surrogate modeling has three main approaches:

When to Use Gaussian Processes¶

GPs Excel When:

✅ Low-dimensional inputs (single-digit to low-tens, roughly $D \lesssim 10$–20)

• Kernel methods scale poorly with dimension (curse of dimensionality)

• Need exponentially more data to cover space as $D$ increases

• Effective dimensionality often lower than nominal $D$ when using ARD (automatic relevance determination learns which dimensions matter)

• Your N-body emulation: $D = 3$ → perfect for GPs!

✅ Limited training data ($N = 100$-5000)

• GPs are data-efficient (Bayesian priors help)

• Can make reasonable predictions from tens of samples

• NNs need thousands to millions of samples

✅ Uncertainty quantification is critical

• GPs give Bayesian uncertainty automatically

• Know when predictions are trustworthy

• Essential for scientific applications (error bars matter!)

✅ Smooth or structured functions

• Physics is usually continuous, differentiable

• Kernels encode this structure explicitly

• GPs exploit smoothness for efficient learning

✅ Interpretability matters

• Hyperparameters have clear physical meaning

  • $\ell_Q = 0.05$ → “Bound fraction changes significantly over $\Delta Q = 0.05$”

  • $\sigma_f^2 = 0.1$ → “Typical variation in bound fraction is $\pm 0.3$”

• Can diagnose what GP learned about physics

✅ Need gradients

• GP predictions are differentiable (can compute $\nabla_{\mathbf{x}} \mu(\mathbf{x})$)

• Useful for optimization, sensitivity analysis

• JAX makes this automatic!

When NOT to Use Gaussian Processes¶

❌ High-dimensional inputs ($D > 20$)

• Kernel evaluations become expensive

• Need impractical amounts of data to cover space

• Consider dimensionality reduction first, or use NNs

❌ Large datasets ($N > 5000$-10,000)

• $O(N^3)$ training cost becomes prohibitive

• $O(N^2)$ memory for covariance matrix

• Solution: Sparse/inducing-point GPs (see below)

❌ Discontinuous or highly irregular functions

• GPs assume some smoothness

• Can’t represent sharp discontinuities well

• Matérn-1/2 helps, but NNs might be better

❌ Complex structured outputs

• GPs naturally handle scalar outputs

• Multi-output GPs possible for moderate cases:

  • Independent GPs: Train separate GP per output (simple, loses correlations between outputs)

  • Multi-output GP (coregionalization): Model output correlations via cross-covariance (ICM/LMC methods)

    • Example: If $R_{\rm core}$ and $R_h$ are correlated, joint GP can leverage this

    • Benefit: More data-efficient than training separate GPs if outputs are correlated

    • Caveat: More hyperparameters to optimize (cross-covariances + per-output lengthscales)

    • Trade-off: More complex training, but can improve predictions when training data is scarce

  • Your project: For $(R_{\rm core}, R_h, \sigma_v)$, independent GPs are recommended (3× training cost, straightforward implementation, maintains modularity). Use multi-output only if data-scarce and outputs strongly correlated.

• For high-dimensional outputs (images, sequences, graphs) → use NNs

❌ Real-time constraints

• Training is one-time cost (acceptable if slow)

• But prediction scales as $O(N M)$ for $M$ test points

• If need millisecond inference for millions of queries, NNs may be faster

Decision Tree for Your N-Body Emulation¶

Let’s apply this to your project:

Your Problem:

• Input: $(Q, N, a) \in \mathbb{R}^3$ as primary parameters

  • $M_{\rm tot}$ is derived: $M_{\rm tot} = N \langle m \rangle_{\rm Kroupa}$ (fixed IMF, $\langle m \rangle \approx 0.5$ M$_\odot$)

• Output: Bound fraction, $R_{\rm core}$, etc. (scalar properties)

• Training data: ~250 simulations (200 provided + 50 yours)

• Goal: Predict at new initial conditions, quantify uncertainty

Analysis:

• ✅ $D = 3$ → Low-dimensional (GPs shine here)

• ✅ $N \approx 250$ → Limited data (GPs data-efficient)

• ✅ Smooth physics (cluster properties vary continuously with ICs)

• ✅ Need uncertainty (must know when emulator is trustworthy)

• ✅ Interpretability (want to understand which parameters matter)

Verdict: Gaussian Processes are ideal for this problem!

But also try Neural Networks because:

• Good learning exercise (compare approaches)

• NNs scale better if you want to add more simulations later

• Can handle multi-output prediction more naturally

• Shows you the tradeoffs empirically

[FIGURE 5.1: Emulation Method Decision Tree]

🎓 What We’ve Learned: Theory Summary¶

Let’s step back and see what we’ve covered in Part I.

First Principles¶

1. The Problem: Scientific simulations encode physics but are too slow for exploration, optimization, inference

2. The Solution: Emulation—learn a fast surrogate from modest training data

3. Why GPs Work:

   • Bayesian inference in function space

   • Gaussian conditioning gives predictions in closed form

   • Kernels encode physical structure (smoothness, lengthscales)

   • Uncertainty quantification is built-in (know when to trust predictions)

4. When to Use GPs:

   • Low-dimensional problems ($D < 20$)

   • Limited data ($N < 5000$)

   • Smooth physics

   • Uncertainty matters

   • Interpretability desired

Connection to Everything You’ve Learned¶

Module 1 (Statistics):

• GPs extend multivariate Gaussians to infinite dimensions

• Central Limit Theorem: Why Gaussian assumptions are reasonable

• Moments: GP mean and covariance are first two moments

Module 2-3 (Stellar Systems):

• N-body simulations define functions $f(\text{ICs}) = \text{cluster properties}$

• Emulation lets you explore phase space without full dynamics

• Lengthscales relate to characteristic scales (dynamical time, relaxation time)

Module 4 (Radiative Transfer):

• Monte Carlo RT is expensive (like N-body)

• Could emulate emergent spectrum vs stellar parameters

• Same workflow applies!

Module 5 (Bayesian Inference):

• GPs are Bayesian: Prior (kernel) + Likelihood (data) = Posterior (predictions)

• MCMC becomes tractable with emulator replacing expensive likelihood

• Marginal likelihood for hyperparameters = evidence for model selection

Project 5 (JAX):

• Automatic differentiation gives gradients for optimization

• JIT compilation makes GP fast

• Vectorization (vmap) enables efficient prediction at many test points

• GPU acceleration makes $N=1000$ feasible

The Philosophical Point¶

Machine learning is not magic. It’s principled probabilistic inference with mathematical foundations.

GPs exemplify this:

• Every component has clear interpretation (kernel, hyperparameters, predictions, uncertainty)

• Connections to statistics, physics, computation

• Uncertainty quantification is first-class citizen (not afterthought)

• Hyperparameters have physical meaning (not just numbers to tune)

When you implement a GP, you’re not using a black box—you’re applying a century of mathematics (Kolmogorov, Wiener, Rasmussen) to solve a 21st-century astrophysics problem.

This is computational astrophysics in 2025!

🔮 Preview: Part II¶

In Part II: GP Implementation, you’ll learn:

• The complete emulation workflow (training data → predictions)

• Hyperparameter optimization via marginal likelihood

• Numerical implementation in JAX

• Validation and diagnostics

• How to use emulators for scientific inference

See you in Part II! Continue to Part II: Gaussian Processes - Implementation and Practice.

Self-Check Rubric: Do You Understand GPs?¶

Before Part II, honestly assess yourself using this rubric. This will help you identify gaps and ensure you’re ready for implementation.

1. Conceptual Understanding¶

• Can you explain why Gaussians are special? (closure under conditioning and marginalization)

• Can you explain what a kernel does in one sentence without using the word “covariance”?

• Can you sketch how $\sigma^2(x)$ varies with position when you have training data at $x \in \{1, 3, 5\}$?

• Can you distinguish epistemic from aleatoric uncertainty and give an astrophysics example of each?

If you checked <3 boxes: Reread “First Principles” and “Uncertainty” sections before Part II.

2. Practical Knowledge¶

• Can you list 5 hyperparameters of an SE kernel with ARD and explain what each controls?

• Can you write pseudocode for computing $\mathbf{K}$ from scratch? (triple loop is fine)

• Can you explain why you should never compute $\mathbf{K}^{-1}$ explicitly?

• Can you design a training dataset for emulating your N-body simulator? (what range of $Q, N, a$? how many points?)

If you checked <3 boxes: Work through the “Kernels” and “Computational” sections with paper and pen.

3. Physical Intuition¶

• Can you guess approximate lengthscales for each parameter in your star cluster problem?

• Can you explain why extrapolating a GP beyond the training range is risky?

• Can you sketch $\mu(x) \pm 2\sigma(x)$ for a 1D problem with 3 training points?

• Can you explain the connection between GP lengthscale and physical smoothness of your simulator?

If you checked <3 boxes: Reread the “Why This Matters” sections and think about your N-body simulations.

4. Implementation Readiness¶

• Can you write a JAX function that computes the SE kernel for two input vectors?

• Can you explain how to standardize training data and denormalize predictions?

• Can you list 3 potential numerical stability issues and how to fix each?

• Do you know what “jitter” is and approximately how large to make it?

If you checked <3 boxes: Start Part II prepared to learn; ask questions early.

Scoring:

• 14+ / 16 boxes checked: Ready for Part II! You have a solid conceptual foundation.

• 12-14: You’re close; review weak areas before starting Part II.

• < 12: You’ll learn in Part II, but prioritize catching up on foundational concepts.

Note: Honest self-assessment helps you learn faster. If you’re uncertain about a box, leave it unchecked and revisit that topic.

📖 Glossary of Key Terms¶

ARD (Automatic Relevance Determination): Kernel with per-dimension lengthscales $\ell_d$, allowing the GP to learn which input dimensions matter most. Small $\ell_d$ → high sensitivity to dimension $d$, large $\ell_d$ → low sensitivity (GP effectively ignores that dimension).

LHS (Latin Hypercube Sampling): Stratified sampling method ensuring uniform coverage in each dimension separately. Better than random sampling for space-filling experimental designs. Widely used for training surrogate models.

OOD (Out-of-Distribution): Test points far from training data, where the emulator extrapolates and uncertainty is high. Detect via kernel similarity: $\max_i k(\mathbf{x}_*, \mathbf{x}_i) / \sigma_f^2 < \text{threshold}$ (e.g., 0.05; validate empirically for your problem as threshold depends on learned hyperparameters).

ICM/LMC (Intrinsic Coregionalization Model / Linear Model of Coregionalization): Multi-output GP methods that model correlations between outputs via cross-covariance functions. Beyond scope of this course; see Álvarez et al. (2012) for details.

Epistemic Uncertainty: Uncertainty about the latent function $f$ (reducible with more training data). From GP: $\sigma_*^2(\mathbf{x})$. Also called “model uncertainty” or “knowledge uncertainty.”

Aleatoric Uncertainty: Irreducible noise in observations (measurement error, simulation stochasticity). From likelihood: $\sigma_n^2$. Also called “data uncertainty.”

Marginal Likelihood (Evidence): $p(\mathbf{y} | X, \boldsymbol{\theta})$, the probability of observing training data given hyperparameters, integrating over all possible functions. Used for Type-II maximum likelihood hyperparameter optimization. Implements Occam’s razor automatically.

Kernel (Covariance Function): $k(\mathbf{x}, \mathbf{x}')$, measures correlation between function values at inputs $\mathbf{x}$ and $\mathbf{x}'$. Encodes smoothness assumptions and prior beliefs about function structure.

Lengthscale ($\ell$): Characteristic distance in input space over which function values remain correlated. Small $\ell$ → rapid variation (wiggly functions), large $\ell$ → smooth variation (slowly varying functions). Units: same as input dimensions.

Signal Variance ($\sigma_f^2$): Prior variance of function values (overall amplitude of variation). Units: output$^2$.

Noise Variance ($\sigma_n^2$): Observation noise level (aleatoric uncertainty). Units: output$^2$.

Cholesky Factorization: Decomposition $\mathbf{K}_y = \mathbf{L} \mathbf{L}^T$ where $\mathbf{L}$ is lower-triangular. Numerically stable way to solve linear systems and compute determinants. Essential for GP implementation.

Jitter: Small constant added to kernel matrix diagonal to ensure numerical stability: $\mathbf{K}_y = \mathbf{K} + (\sigma_n^2 + \epsilon_{\rm jitter}) \mathbf{I}$. Prevents Cholesky failures from near-singular matrices. Typical values: $\epsilon_{\rm jitter} \sim 10^{-6}$ (scale-aware).

SE (Squared Exponential) Kernel: Also called RBF (Radial Basis Function). Assumes infinitely smooth functions. Often too smooth for real physics.

Matérn Kernel: Family of kernels with tunable smoothness $\nu$. Matérn-5/2 is recommended default (twice differentiable, realistic for physics).

PSD (Positive Semi-Definite): Property required for valid kernels. Ensures covariance matrices have non-negative eigenvalues (variances can’t be negative).

MCMC (Markov Chain Monte Carlo): Sampling method for complex posterior distributions (from Module 5). GPs enable MCMC-based inference by providing fast, differentiable emulators of expensive simulators.

JAX: Just-In-Time compiled array computing library for Python with automatic differentiation. Makes GP implementation fast and enables gradient-based hyperparameter optimization.

JIT (Just-In-Time Compilation): Compilation technique that speeds up Python code to near-C performance. Use @jit decorator in JAX.

vmap (Vectorized Map): JAX function for automatic vectorization. Maps a function over batch dimension efficiently (avoids Python loops).

Continue to Part II for implementation! 🚀

“The best way to understand a Gaussian Process is to implement one.” — Every computational scientist ever