#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""EM-Based Empirical Gaussian Process.
This module implements an EM (Expectation-Maximization) based Empirical Gaussian
Process that learns its prior mean μ and covariance Σ from K independent datasets.
See [lin2026empirical]_ for details.
**Key Concept**: Given K independent datasets assumed to be generated by the
same underlying stochastic process, we estimate the GP prior (mean and covariance)
using the EM algorithm with closed-form updates.
**Advantages over Sampling**:
- Deterministic (no Monte Carlo variance)
- Exact gradients for hyperparameter optimization
- Simpler implementation
**Pre-training Workflow**: The recommended workflow for production use is:
1. Call ``pretrain_em_prior()`` once on historical data.
2. Optionally save the ``EMPriorContainer`` for later use.
3. Create multiple ``EMEmpiricalGaussianProcess`` instances with the same
container for different (train_X, train_Y) test sets.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
from botorch.models import SingleTaskGP
from botorch.models.empirical_gps.utils import (
build_unique_inputs,
ExperimentDataset,
project_psd,
trace_matched_shrinkage,
UniqueInputs,
)
from botorch.models.gpytorch import GPyTorchModel
from botorch.models.model_list_gp_regression import ModelListGP
from gpytorch.distributions import MultivariateNormal
from gpytorch.kernels import Kernel
from gpytorch.likelihoods import GaussianLikelihood, Likelihood
from gpytorch.means import Mean
from gpytorch.mlls import (
MarginalLogLikelihood,
SumMarginalLogLikelihood as GPyTorchSumMarginalLogLikelihood,
)
from gpytorch.models import ExactGP
from linear_operator import to_linear_operator
from linear_operator.operators import LinearOperator
from linear_operator.utils.cholesky import psd_safe_cholesky
from torch import Tensor
# =============================================================================
# Pre-trained Prior Container
# =============================================================================
[docs]
@dataclass(frozen=True)
class EMPriorContainer:
"""Container for pre-trained EM prior quantities.
This container holds all quantities computed during EM pre-training
that are needed for GP conditioning in EMEmpiricalGaussianProcess.
The container is designed to be:
- Immutable after creation (frozen=True)
- Serializable (can be saved/loaded for caching)
- Complete (contains everything needed for forward pass)
All tensor attributes are DETACHED to prevent gradient flow through
pre-training when used in a model's forward pass.
Args:
datasets: The K experiment datasets used for EM pre-training.
mean_module: Parametric mean module for shift interpolation.
covar_module: Parametric covariance module for shift interpolation.
likelihood_noise: Noise variance used during pre-training.
num_em_iterations: Number of EM iterations used.
use_mean_prior: Whether kernel prior on μ was used.
use_covar_prior: Whether Inverse-Wishart prior on Σ was used.
iw_nu: Degrees of freedom for IW prior (or None).
em_convergence_tol: Convergence tolerance used.
enable_interpolation: Whether interpolation is enabled.
X_inducing: (M, d) inducing point locations.
mu_inducing: (M,) EM-converged mean at inducing points.
Sigma_inducing: (M, M) EM-converged covariance at inducing points.
L_kernel_inducing: (M, M) Cholesky factor of K(Z, Z) for interpolation.
delta_mu: (M,) mean shift = mu_inducing - m(Z).
unique_inputs: UniqueInputs structure for index lookup.
"""
# Configuration (stored for reference/serialization)
datasets: list[ExperimentDataset]
mean_module: Mean
covar_module: Kernel
likelihood_noise: Tensor
num_em_iterations: int
use_mean_prior: bool
use_covar_prior: bool
iw_nu: float | None
em_convergence_tol: float | None
enable_interpolation: bool
use_inducing_points: bool # True if inducing points were explicitly provided
# Computed State (ALL DETACHED)
X_inducing: Tensor
mu_inducing: Tensor
Sigma_inducing: Tensor
L_kernel_inducing: Tensor
delta_mu: Tensor
unique_inputs: UniqueInputs
[docs]
def save(self, path: str) -> None:
"""Save the container to disk.
Args:
path: File path to save the container.
"""
torch.save(self, path)
[docs]
@classmethod
def load(cls, path: str) -> "EMPriorContainer":
"""Load a container from disk.
Args:
path: File path to load the container from.
Returns:
Loaded EMPriorContainer.
"""
return torch.load(path, weights_only=False)
# =============================================================================
# Core EM Primitives
# =============================================================================
def _evaluate_mean(mean_module: Mean, X: Tensor) -> Tensor:
"""Evaluate ``mean_module`` at ``X``, returning a 1-dim mean vector.
Mean modules may return either a ``(..., n)`` or ``(..., n, 1)`` tensor; this
normalizes the result by dropping a trailing singleton output dimension.
"""
mean = mean_module(X)
return mean.squeeze(-1) if mean.dim() > 1 else mean
def _compute_prior_scale_matrix(
Sigma_init: Tensor,
nu: float,
N_obs: int,
) -> Tensor:
"""Compute IW scale matrix Ψ such that prior mode = Sigma_init.
For Σ ~ IW(Ψ, ν), the mode is Ψ / (ν + N + 1).
To achieve mode = Sigma_init, we set Ψ = (ν + N + 1) × Sigma_init.
Args:
Sigma_init: Initial covariance matrix from parametric kernel.
nu: Degrees of freedom for Inverse-Wishart prior.
N_obs: Number of observed points (used for consistent scaling).
Returns:
Psi: Scaled prior matrix such that mode(IW(Psi, nu)) = Sigma_init.
"""
return (nu + N_obs + 1) * Sigma_init
def _compute_observation_factors(
Sigma_SS: Tensor,
y: Tensor,
mu_S: Tensor,
likelihood_noise: Tensor | None = None,
) -> tuple[Tensor, Tensor]:
"""Compute Cholesky factor and solve for observation conditioning.
Args:
Sigma_SS: (n, n) covariance at observation locations.
y: (n,) observation values.
mu_S: (n,) mean at observation locations.
likelihood_noise: Optional scalar noise variance to add.
Returns:
L: (n, n) Cholesky factor of Σ_{S,S} + noise.
alpha: (n,) = (Σ_{S,S} + noise)^{-1} @ (y - μ_S).
"""
n = Sigma_SS.shape[0]
if likelihood_noise is not None:
# Move noise to the covariance's device/dtype to avoid mismatches
# (e.g. a CPU likelihood noise with CUDA data). Use diag_embed so a
# single-observation dataset (n == 1) does not collapse to a 0-d scalar.
noise_diag = likelihood_noise.to(Sigma_SS).reshape(-1).expand(n)
Sigma_SS = Sigma_SS + torch.diag_embed(noise_diag)
L = psd_safe_cholesky(Sigma_SS)
residual = y - mu_S
alpha = torch.cholesky_solve(residual.unsqueeze(-1), L).squeeze(-1)
return L, alpha
def _compute_conditional_moments(
mu_T: Tensor,
Sigma_TT: Tensor,
Sigma_T_S: Tensor,
L: Tensor,
alpha: Tensor,
) -> tuple[Tensor, Tensor]:
"""Compute conditional mean and covariance.
The formulas are:
E[θ_T | y] = μ_T + Σ_{T,S} @ α
Cov[θ_T | y] = Σ_{T,T} - V_T^T @ V_T where V_T = L^{-1} @ Σ_{S,T}
Args:
mu_T: (|T|,) mean at target locations.
Sigma_TT: (|T|, |T|) covariance at target locations.
Sigma_T_S: (|T|, n) cross-covariance between target and observations.
L: (n, n) Cholesky factor from _compute_observation_factors.
alpha: (n,) solve result from _compute_observation_factors.
Returns:
cond_mean: (|T|,) E[θ_T | y].
cond_cov: (|T|, |T|) Cov[θ_T | y].
"""
# Conditional mean: μ_T + Σ_{T,S} @ α
cond_mean = mu_T + Sigma_T_S @ alpha
# Conditional covariance: Σ_{T,T} - V_T^T @ V_T
V_T = torch.linalg.solve_triangular(L, Sigma_T_S.T, upper=False)
cond_cov = torch.addmm(Sigma_TT, V_T.T, V_T, beta=1, alpha=-1)
return cond_mean, cond_cov
def _e_step(
datasets: list[ExperimentDataset],
mu: Tensor,
Sigma: Tensor,
likelihood_noise: Tensor | None = None,
experiment_indices: list[Tensor] | None = None,
X_inducing: Tensor | None = None,
mean_module: Mean | None = None,
covar_module: Kernel | None = None,
) -> tuple[list[Tensor], list[Tensor]]:
"""E-step: compute conditional distributions for each dataset.
For each dataset i, computes E[θ | y_i] and Cov[θ | y_i].
When experiment_indices is provided (standard case), uses direct indexing.
When X_inducing is provided (inducing point case), uses shift interpolation
to compute the joint distribution and conditions on observations.
Args:
datasets: K experiment datasets.
mu: (N,) current mean estimate (at all locations or inducing points).
Sigma: (N, N) current covariance estimate (at all locations or inducing points).
likelihood_noise: Optional scalar noise variance.
experiment_indices: List of K index tensors for each dataset (standard case).
X_inducing: (M, d) inducing point locations (inducing point case).
mean_module: Parametric mean module for interpolation (required if X_inducing).
covar_module: Parametric covariance module for interpolation (required for both
X_inducing, and if test points could be different from the set of historical
inputs).
Returns:
cond_means: List of K tensors - E[θ | y_i].
cond_covs: List of K tensors - Cov[θ | y_i].
"""
cond_means: list[Tensor] = []
cond_covs: list[Tensor] = []
# Precompute inducing point quantities if using shift interpolation
if X_inducing is not None:
K_ZZ = covar_module(X_inducing, X_inducing).to_dense()
L_ZZ = psd_safe_cholesky(K_ZZ)
m_Z = _evaluate_mean(mean_module, X_inducing)
delta_mu = mu - m_Z
for i, dataset in enumerate(datasets):
# Note (potential optimization): when using inducing points, the
# interpolated covariance at observation points X_k is low-rank-plus-
# dense:
# Σ(X_k) + σ²I = (Λ_k + σ²I) + W_k @ Σ_em @ W_k^T
# where Λ_k = K(X_k,X_k) - K(X_k,Z) K(Z,Z)^{-1} K(Z,X_k) is the (dense)
# Nyström residual and W_k = K(X_k,Z) K(Z,Z)^{-1}. A Woodbury / matrix-
# determinant-lemma factorization in O(n_k M² + M³) is only available
# if the base term (Λ_k + σ²I) is cheaply invertible, i.e. if Λ_k is
# approximated by its diagonal (an FITC-style approximation). With the
# exact dense residual used here for extrapolation, the O(n_k³)
# Cholesky is unavoidable. (When multiple datasets share X_k, W_k and
# Λ_k are identical and need only be computed once.)
# The target locations T are the full set of estimation points, so the
# target moments are always the current (mu, Sigma); only the observed
# block (mu_S, Sigma_SS) and cross-covariance Sigma_TS differ by case.
mu_T, Sigma_TT = mu, Sigma
if X_inducing is not None:
# Inducing point case: use shift interpolation
mu_S, Sigma_SS, Sigma_TS = _interpolate_prior(
X=dataset.X,
mean_module=mean_module,
covar_module=covar_module,
X_inducing=X_inducing,
L_ZZ=L_ZZ,
delta_mu=delta_mu,
Sigma_inducing=Sigma,
include_cross_covariance=True,
)
else:
# Standard case: use indexing
mu_S, Sigma_SS, Sigma_TS = _index_prior(
mu_full=mu,
Sigma_full=Sigma,
indices=experiment_indices[i],
include_cross_covariance=True,
)
L, alpha = _compute_observation_factors(
Sigma_SS, dataset.Y.squeeze(-1), mu_S, likelihood_noise
)
cond_mean, cond_cov = _compute_conditional_moments(
mu_T, Sigma_TT, Sigma_TS, L, alpha
)
cond_means.append(cond_mean)
cond_covs.append(cond_cov)
return cond_means, cond_covs
def _m_step(
cond_means: list[Tensor],
cond_covs: list[Tensor],
Psi: Tensor | None = None,
nu: float | None = None,
N_obs: int | None = None,
K_mu: Tensor | None = None,
m_0: Tensor | None = None,
Sigma_current: Tensor | None = None,
psd_stabilization: bool = False,
shrinkage: float = 0.0,
shrinkage_target: Tensor | None = None,
) -> tuple[Tensor, Tensor]:
"""M-step: update μ and Σ from conditional distributions.
**ML Updates** (when Psi is None):
μ_new = (1/K) Σ_i E[θ_i | Y_i]
Σ_new = (1/K) Σ_i [Cov + (E - μ_new)(E - μ_new)^T]
**MAP Updates** (when Psi provided):
Σ_new = (scatter + Ψ) / (K + ν + N_obs + 1)
Args:
cond_means: List of K conditional means.
cond_covs: List of K conditional covariances.
Psi: Pre-scaled IW scale matrix (ν + N_obs + 1) × Sigma_init.
nu: IW degrees of freedom (required if Psi provided).
N_obs: Number of observed points (for consistent denominator).
K_mu: Kernel matrix for GP prior on μ.
m_0: Prior mean for μ (default: zeros).
Sigma_current: Current Σ estimate (required if K_mu provided).
psd_stabilization: If True, ensure Sigma_new is positive definite by
clipping negative eigenvalues (default: False). This prevents
numerical instability when working with partial observations.
shrinkage: M-step linear-shrinkage intensity ``alpha in [0, 1]`` (default
0). When ``> 0`` and ``shrinkage_target`` is provided, blends Sigma_new
toward the trace-matched target via ``trace_matched_shrinkage``.
shrinkage_target: ``(N_obs, N_obs)`` structured shrinkage target (e.g. the
base-kernel gram ``K(Z, Z)``); required for shrinkage to take effect.
Returns:
mu_new: Updated mean.
Sigma_new: Updated covariance.
"""
K = len(cond_means)
N = cond_means[0].shape[0]
device = cond_means[0].device
dtype = cond_means[0].dtype
# Sum of conditional means
sum_cond_means = sum(cond_means)
# === Update μ ===
if K_mu is None:
# ML: simple average
mu_new = sum_cond_means / K
else:
# MAP with kernel prior: μ ~ N(m_0, K_μ)
if m_0 is None:
m_0 = torch.zeros(N, device=device, dtype=dtype)
L_Sigma = psd_safe_cholesky(Sigma_current)
L_K_mu = psd_safe_cholesky(K_mu)
Sigma_inv = torch.cholesky_inverse(L_Sigma)
K_mu_inv = torch.cholesky_inverse(L_K_mu)
A = K * Sigma_inv + K_mu_inv
rhs = Sigma_inv @ sum_cond_means + K_mu_inv @ m_0
L_A = psd_safe_cholesky(A)
mu_new = torch.cholesky_solve(rhs.unsqueeze(-1), L_A).squeeze(-1)
# === Update Σ ===
scatter = torch.zeros(N, N, device=device, dtype=dtype)
for i in range(K):
diff = cond_means[i] - mu_new
scatter = scatter + cond_covs[i]
scatter = scatter + torch.outer(diff, diff)
# Apply prior regularization
if Psi is None:
Sigma_new = scatter / K
else:
n_for_denom = N_obs if N_obs is not None else N
denom = K + nu + n_for_denom + 1
Sigma_new = (scatter + Psi) / denom
# Linear shrinkage toward a structured target (e.g. the base kernel K(Z, Z)).
# A flexible generalization of the Inverse-Wishart prior: it blends the empirical
# covariance with the target using a free intensity in [0, 1], decoupled from the
# number of inducing points. The target is trace-matched to Sigma_new so only its
# correlation structure is imposed, not its absolute scale.
if shrinkage > 0.0 and shrinkage_target is not None:
Sigma_new = trace_matched_shrinkage(Sigma_new, shrinkage_target, shrinkage)
# === PSD Stabilization ===
# With partial observations and floating-point arithmetic, the covariance
# matrix can develop small negative eigenvalues due to accumulated
# numerical errors. We fix this by eigendecomposition and clipping.
if psd_stabilization:
Sigma_new = project_psd(Sigma_new, min_eigval=1e-10)
return mu_new, Sigma_new
# =============================================================================
# Shared EM Algorithm
# =============================================================================
def _run_em_algorithm(
datasets: list[ExperimentDataset],
mu_init: Tensor,
Sigma_init: Tensor,
X_inducing: Tensor,
mean_module: Mean,
covar_module: Kernel,
likelihood_noise: Tensor | None,
experiment_indices: list[Tensor] | None,
use_inducing_points: bool,
num_em_iterations: int,
Psi: Tensor | None,
K_mu: Tensor | LinearOperator | None,
iw_nu: float | None,
em_convergence_tol: float | None,
N_inducing: int,
iteration_history: list[tuple[Tensor, Tensor]] | None = None,
covariance_shrinkage: float = 0.0,
shrinkage_target: Tensor | None = None,
) -> tuple[Tensor, Tensor]:
"""Core EM algorithm - single source of truth for EM computation.
This function encapsulates the EM loop and is used by both `pretrain_em_prior()`
and the internal `EMEmpiricalGaussianProcess._run_em()` method to ensure
consistent behavior.
Args:
datasets: K experiment datasets.
mu_init: (N_inducing,) initial mean estimate.
Sigma_init: (N_inducing, N_inducing) initial covariance estimate.
X_inducing: (N_inducing, d) inducing point locations.
mean_module: Parametric mean module.
covar_module: Parametric covariance module.
likelihood_noise: Scalar noise variance, or None.
experiment_indices: List of index tensors for each dataset (direct indexing).
use_inducing_points: If True, use shift interpolation in E-step.
num_em_iterations: Maximum number of EM iterations.
Psi: IW prior scale matrix, or None for ML estimation.
K_mu: Kernel prior matrix for mean, or None for ML estimation.
iw_nu: IW degrees of freedom, or None.
em_convergence_tol: Convergence tolerance for early stopping, or None.
N_inducing: Number of inducing points (for M-step normalization).
iteration_history: Optional list to record convergence diagnostics. When
provided, the ``(mu, Sigma)`` snapshot after each EM iteration is
appended (as clones); pass ``None`` (default) to skip recording, in
which case no snapshots are computed.
covariance_shrinkage: M-step linear-shrinkage intensity ``alpha in [0, 1]``
(default 0), passed through to each ``_m_step``.
shrinkage_target: ``(N_inducing, N_inducing)`` shrinkage target (e.g.
``K(Z, Z)``); required for ``covariance_shrinkage`` to take effect.
Returns:
``(mu_em, Sigma_em)``, the EM-estimated mean and covariance at the
inducing points. If ``iteration_history`` was provided, it is populated
in place with one ``(mu, Sigma)`` per iteration (up to early stopping).
"""
mu, Sigma = mu_init.clone(), Sigma_init.clone()
for _ in range(num_em_iterations):
mu_prev, Sigma_prev = mu, Sigma
cond_means, cond_covs = _e_step(
datasets,
mu,
Sigma,
likelihood_noise=likelihood_noise,
experiment_indices=None if use_inducing_points else experiment_indices,
X_inducing=X_inducing if use_inducing_points else None,
mean_module=mean_module if use_inducing_points else None,
covar_module=covar_module if use_inducing_points else None,
)
mu, Sigma = _m_step(
cond_means,
cond_covs,
Psi=Psi,
nu=iw_nu,
N_obs=N_inducing,
K_mu=K_mu,
Sigma_current=Sigma,
psd_stabilization=True,
shrinkage=covariance_shrinkage,
shrinkage_target=shrinkage_target,
)
# Record state after this iteration
if iteration_history is not None:
iteration_history.append((mu.clone(), Sigma.clone()))
# Early stopping if converged
if em_convergence_tol is not None:
mu_change = (mu - mu_prev).abs().max()
Sigma_change = (Sigma - Sigma_prev).abs().max()
if mu_change < em_convergence_tol and Sigma_change < em_convergence_tol:
break
return mu, Sigma
# =============================================================================
# Pre-training Function
# =============================================================================
[docs]
def pretrain_em_prior(
datasets: list[ExperimentDataset],
mean_module: Mean,
covar_module: Kernel,
likelihood_noise: Tensor | float | None = None,
inducing_points: Tensor | None = None,
num_em_iterations: int = 16,
use_mean_prior: bool = False,
use_covar_prior: bool = False,
iw_nu: float | None = None,
em_convergence_tol: float | None = 1e-6,
enable_interpolation: bool = True,
init_mode: str = "kernel",
iteration_history: list[tuple[Tensor, Tensor]] | None = None,
covariance_shrinkage: float = 0.0,
) -> EMPriorContainer:
"""Pre-train an EM-based empirical prior from historical datasets.
This function runs the EM algorithm on historical datasets to estimate
the GP prior mean and covariance at inducing points. The returned
container can be used with EMEmpiricalGaussianProcess to skip EM
computation during model instantiation.
This is the recommended workflow for production use:
1. Call pretrain_em_prior() once on historical data.
2. Optionally save the container for later use.
3. Create multiple EMEmpiricalGaussianProcess instances with the
same container for different (train_X, train_Y) test sets.
All returned tensors in the container are DETACHED to prevent gradient
flow through pre-training when used in a model's forward pass.
Args:
datasets: List of ExperimentDataset objects containing historical data.
mean_module: Parametric mean module for shift interpolation.
covar_module: Parametric covariance module for shift interpolation.
likelihood_noise: Noise variance for observations. If None, defaults to 1e-10.
inducing_points: Optional (M, d) tensor of inducing point locations.
If None, uses unique inputs from historical data.
num_em_iterations: Maximum number of EM iterations (default: 16).
use_mean_prior: If True, use covar_module as kernel prior on μ.
use_covar_prior: If True, use Inverse-Wishart prior on Σ.
iw_nu: Degrees of freedom for IW prior. If None and use_covar_prior=True,
automatically set to M + 2 where M is the number of inducing points.
em_convergence_tol: Convergence tolerance for early stopping (default: 1e-6).
enable_interpolation: If True (default), enable shift interpolation.
init_mode: Initialization mode for the EM algorithm. ``"kernel"``
(default) initializes mu and Sigma from the parametric GP prior
(mean_module and covar_module); ``"naive"`` initializes mu to zeros
and Sigma to the identity matrix (useful for testing the effect of
initialization on results).
iteration_history: Optional list to record convergence diagnostics. When
provided, the ``(mu, Sigma)`` snapshot after each EM iteration is
appended (as clones); pass ``None`` (default) to skip recording.
covariance_shrinkage: Optional M-step linear-shrinkage intensity
``alpha in [0, 1]`` (default 0). When ``> 0``, each M-step blends the
estimated covariance toward the trace-matched base-kernel gram
``K(Z, Z)`` -- a trace-matched (empirical-Bayes) Inverse-Wishart-style
regularizer for the rank-limited empirical covariance (e.g. the
data-starved tail of early-stopped curves).
Returns:
EMPriorContainer with all pre-computed quantities (all tensors detached).
If ``iteration_history`` was provided, it is populated in place with one
``(mu, Sigma)`` per EM iteration.
Example:
>>> em_prior = pretrain_em_prior(
... datasets=historical_datasets,
... mean_module=mean_module,
... covar_module=covar_module,
... )
>>> em_prior.save("path/to/prior.pt")
>>>
>>> # Later, reuse for multiple test sets
>>> for train_X, train_Y in test_sets:
... model = EMEmpiricalGaussianProcess(
... train_X=train_X, train_Y=train_Y, em_prior=em_prior
... )
"""
# Build unique inputs from datasets
unique_inputs = build_unique_inputs(datasets, X_forward=None)
# Determine inducing points
if inducing_points is not None:
if inducing_points.dim() != 2:
raise ValueError(
f"inducing_points must be 2D, got {inducing_points.dim()}D"
)
# Check feature dimension matches data
sample_X = datasets[0].X
if inducing_points.shape[-1] != sample_X.shape[-1]:
raise ValueError(
f"inducing_points feature dimension ({inducing_points.shape[-1]}) "
f"must match data feature dimension ({sample_X.shape[-1]})"
)
X_inducing = inducing_points
use_inducing_points = True
else:
X_inducing = unique_inputs.X_all
use_inducing_points = False
N_inducing = X_inducing.shape[0]
# Handle likelihood noise
if likelihood_noise is None:
likelihood_noise = torch.tensor(
1e-10, dtype=X_inducing.dtype, device=X_inducing.device
)
elif not isinstance(likelihood_noise, Tensor):
likelihood_noise = torch.tensor(
likelihood_noise, dtype=X_inducing.dtype, device=X_inducing.device
)
else:
# Normalize a provided tensor (e.g. likelihood.noise) to the inducing
# point device/dtype so downstream EM stays device-consistent.
likelihood_noise = likelihood_noise.to(
dtype=X_inducing.dtype, device=X_inducing.device
)
# Handle iw_nu for Inverse-Wishart prior
if use_covar_prior:
if iw_nu is None:
iw_nu = float(N_inducing + 2)
elif iw_nu <= N_inducing - 1:
raise ValueError(
f"iw_nu must be > M - 1 = {N_inducing - 1} for "
f"Inverse-Wishart on {N_inducing}×{N_inducing} matrix, "
f"got iw_nu={iw_nu}"
)
# Parametric prior at inducing points. m(Z) and K(Z, Z) are needed for the
# interpolation cache below, and also serve as the "kernel" initialization,
# so compute them once here and reuse.
m_inducing = _evaluate_mean(mean_module, X_inducing)
K_kernel_inducing = covar_module(X_inducing, X_inducing).to_dense()
# Initialize mu and Sigma based on init_mode
if init_mode == "kernel":
# Initialize from parametric prior at inducing points
mu_init = m_inducing
Sigma_init = K_kernel_inducing
elif init_mode == "naive":
# Naive initialization: zero mean, identity covariance
mu_init = torch.zeros(
N_inducing, dtype=X_inducing.dtype, device=X_inducing.device
)
Sigma_init = torch.eye(
N_inducing, dtype=X_inducing.dtype, device=X_inducing.device
)
else:
raise ValueError(
f"Unknown init_mode: {init_mode}. Must be 'kernel' or 'naive'."
)
# Compute prior matrices
Psi = (
_compute_prior_scale_matrix(Sigma_init, iw_nu, N_inducing)
if use_covar_prior
else None
)
K_mu = Sigma_init if use_mean_prior else None
# Get experiment indices for direct indexing case
experiment_indices = unique_inputs.experiment_indices
# Run EM algorithm
mu_em, Sigma_em = _run_em_algorithm(
datasets=datasets,
mu_init=mu_init,
Sigma_init=Sigma_init,
X_inducing=X_inducing,
mean_module=mean_module,
covar_module=covar_module,
likelihood_noise=likelihood_noise,
experiment_indices=experiment_indices,
use_inducing_points=use_inducing_points,
num_em_iterations=num_em_iterations,
Psi=Psi,
K_mu=K_mu,
iw_nu=iw_nu,
em_convergence_tol=em_convergence_tol,
N_inducing=N_inducing,
iteration_history=iteration_history,
covariance_shrinkage=covariance_shrinkage,
shrinkage_target=(K_kernel_inducing if covariance_shrinkage > 0.0 else None),
)
# Compute interpolation cache quantities (reuse m_inducing and
# K_kernel_inducing computed once above).
L_kernel_inducing = psd_safe_cholesky(K_kernel_inducing)
delta_mu = mu_em - m_inducing
# Return container with ALL TENSORS DETACHED
container = EMPriorContainer(
# Configuration
datasets=datasets,
mean_module=mean_module,
covar_module=covar_module,
likelihood_noise=likelihood_noise.detach(),
num_em_iterations=num_em_iterations,
use_mean_prior=use_mean_prior,
use_covar_prior=use_covar_prior,
iw_nu=iw_nu,
em_convergence_tol=em_convergence_tol,
enable_interpolation=enable_interpolation,
use_inducing_points=use_inducing_points,
# Computed state (ALL DETACHED)
X_inducing=X_inducing.detach(),
mu_inducing=mu_em.detach(),
Sigma_inducing=Sigma_em.detach(),
L_kernel_inducing=L_kernel_inducing.detach(),
delta_mu=delta_mu.detach(),
unique_inputs=unique_inputs,
)
return container
# =============================================================================
# Model Class
# =============================================================================
[docs]
class EMEmpiricalGaussianProcess(ExactGP, GPyTorchModel):
"""Empirical GP with EM-based prior estimation.
Uses closed-form E-step and M-step updates instead of Monte Carlo sampling.
This provides deterministic, exact gradient flow for hyperparameter optimization.
**Two Initialization Modes**:
1. With em_prior (recommended for production):
- Skip EM entirely
- Use pre-computed prior for conditioning
- forward() never re-runs EM
- ~100x faster when reusing prior for many test sets
2. Without em_prior (backward compatible):
- Run EM in constructor
- Re-run EM in training mode forward()
- Same behavior as before
**Recommended Workflow**:
.. code-block:: python
# Pre-train once
em_prior = pretrain_em_prior(datasets, mean_module, covar_module)
# Reuse for many test sets (no EM computation!)
for train_X, train_Y in test_sets:
model = EMEmpiricalGaussianProcess(
train_X=train_X, train_Y=train_Y, em_prior=em_prior
)
posterior = model(test_X)
**Transforms are not yet supported** in this implementation. Do not use
input_transform or outcome_transform with this model.
Args:
train_X: (n, d) tensor of training inputs for GP conditioning.
train_Y: (n, m) tensor of training targets for GP conditioning.
em_prior: Pre-trained EMPriorContainer. If provided, skips EM entirely
and uses the pre-computed prior. This is the recommended mode for
production use when applying the same prior to multiple test sets.
datasets: List of ExperimentDataset objects. Required if em_prior is None.
These K independent datasets are used to estimate the GP prior via EM.
mean_module: Initial prior mean module. Required if em_prior is None.
covar_module: Initial prior covariance module. Required if em_prior is None.
likelihood: A likelihood. If omitted, uses a GaussianLikelihood.
num_em_iterations: Number of EM iterations (default: 16). Only used if
em_prior is None.
use_mean_prior: If True, use covar_module as kernel prior on μ.
use_covar_prior: If True, use Inverse-Wishart prior on Σ.
iw_nu: Degrees of freedom for Inverse-Wishart prior. Must be > M - 1
where M is the number of inducing points. If None (default),
automatically set to M + 2 when use_covar_prior=True.
inducing_points: Optional (M, d) tensor of inducing point locations.
enable_interpolation: If True (default), enable shift interpolation.
warm_start_em: If True, warm-start EM iterations from previously
converged values during hyperparameter optimization.
em_convergence_tol: Convergence tolerance for early stopping of EM.
init_mode: Initialization mode for the EM algorithm. ``"kernel"``
(default) initializes mu and Sigma from the parametric GP prior
(mean_module and covar_module); ``"naive"`` initializes mu to zeros
and Sigma to the identity matrix (useful for testing the effect of
initialization on results).
covariance_shrinkage: Optional M-step linear-shrinkage intensity ``alpha
in [0, 1]`` (default 0). When ``> 0`` (and ``em_prior`` is None), each
EM M-step blends the estimated covariance toward the trace-matched
base-kernel gram ``K(Z, Z)`` via ``trace_matched_shrinkage`` -- a
full-rank, structured regularizer for the rank-limited empirical
covariance. It is a trace-matched (empirical-Bayes) Inverse-Wishart-
style regularizer -- equivalent to MAP-EM under an IW prior on ``Sigma``
whose scale (proportional to ``K(Z, Z)``) is re-matched to the data each
step, with ``alpha`` setting the prior strength ``nu`` via
``alpha = (nu + M + 1) / (K + nu + M + 1)`` -- and pins otherwise
data-starved directions to the base structure. This is the *in-EM*
regularizer; for a *post-EM* conditioning-time **additive** base kernel,
use ``from_pretrained(..., base_covar_module=...)`` instead.
"""
# Explicitly mark that transforms are not supported
_supports_input_transform: bool = False
_supports_outcome_transform: bool = False
_num_outputs: int = 1
def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
# Option 1: Pre-trained prior (RECOMMENDED)
em_prior: EMPriorContainer | None = None,
# Option 2: Train from scratch (backward compatible)
datasets: list[ExperimentDataset] | None = None,
mean_module: Mean | None = None,
covar_module: Kernel | None = None,
likelihood: Likelihood | None = None,
# EM parameters (only used if em_prior is None)
num_em_iterations: int = 16,
use_mean_prior: bool = False,
use_covar_prior: bool = False,
iw_nu: float | None = None,
inducing_points: Tensor | None = None,
enable_interpolation: bool = True,
warm_start_em: bool = False,
em_convergence_tol: float | None = 1e-6,
init_mode: str = "kernel",
learnable_inducing_points: bool = False,
covariance_shrinkage: float = 0.0,
) -> None:
"""Initialize the model. See the class docstring for argument details."""
# Validate likelihood
if likelihood is None:
likelihood = GaussianLikelihood()
elif not isinstance(likelihood, GaussianLikelihood):
raise ValueError(
"EMEmpiricalGaussianProcess only supports GaussianLikelihood, "
f"got {type(likelihood).__name__}. "
"FixedNoiseGaussianLikelihood and other heteroscedastic likelihoods "
"are not supported because the noise tensor is fixed to a specific "
"number of data points, but each historical dataset may have a "
"different number of observations. You can use a GaussianLikelihood "
"with fixed noise instead: "
"likelihood = GaussianLikelihood(); "
"likelihood.noise = your_noise_value; "
"likelihood.raw_noise.requires_grad_(False)"
)
# Unified path: create container if not provided
using_pretrained = em_prior is not None
if using_pretrained and covariance_shrinkage > 0.0:
raise ValueError(
"covariance_shrinkage only applies when fitting the EM prior from "
"`datasets`; with a pretrained `em_prior`, use "
"from_pretrained(..., base_covar_module=...) instead."
)
if em_prior is None:
# Validate required arguments for from-scratch initialization
if datasets is None:
raise ValueError(
"Either em_prior or datasets must be provided. "
"Use em_prior for pre-trained prior, or datasets for "
"from-scratch EM computation."
)
if mean_module is None or covar_module is None:
raise ValueError(
"mean_module and covar_module are required when "
"em_prior is not provided."
)
# Create container via pretrain_em_prior
em_prior = pretrain_em_prior(
datasets=datasets,
mean_module=mean_module,
covar_module=covar_module,
likelihood_noise=getattr(likelihood, "noise", None),
inducing_points=inducing_points,
num_em_iterations=num_em_iterations,
use_mean_prior=use_mean_prior,
use_covar_prior=use_covar_prior,
iw_nu=iw_nu,
em_convergence_tol=em_convergence_tol,
enable_interpolation=enable_interpolation,
init_mode=init_mode,
covariance_shrinkage=covariance_shrinkage,
)
# Single initialization path via container
self._init_from_container(
container=em_prior,
train_X=train_X,
train_Y=train_Y,
likelihood=likelihood,
using_pretrained=using_pretrained,
warm_start_em=warm_start_em,
learnable_inducing_points=learnable_inducing_points,
covariance_shrinkage=covariance_shrinkage,
)
def _init_from_container(
self,
container: EMPriorContainer,
train_X: Tensor,
train_Y: Tensor,
likelihood: Likelihood,
using_pretrained: bool,
warm_start_em: bool,
learnable_inducing_points: bool = False,
covariance_shrinkage: float = 0.0,
) -> None:
"""Initialize model state from an EMPriorContainer.
This method provides a single initialization path regardless of whether
the container was pre-computed or created during __init__.
Args:
container: EMPriorContainer with all pre-computed quantities.
train_X: Training inputs for GP conditioning.
train_Y: Training targets for GP conditioning.
likelihood: Gaussian likelihood.
using_pretrained: Whether the container was pre-trained (vs created
in __init__). Controls whether forward() re-runs EM.
warm_start_em: Whether to warm-start EM in training mode.
learnable_inducing_points: If True, register inducing point locations
as a learnable nn.Parameter for joint optimization with kernel
hyperparameters via the observed-data MLL.
"""
# Call parent __init__
super().__init__(train_X, train_Y.squeeze(-1), likelihood)
# Ensure the likelihood matches the data device/dtype so the EM re-run
# and per-dataset MLL stay device-consistent (the default likelihood is
# created on CPU).
self.likelihood.to(device=train_X.device, dtype=train_X.dtype)
# Store configuration from container
self.datasets = container.datasets
self.initial_mean_module = container.mean_module
self.initial_covar_module = container.covar_module
self.num_em_iterations = container.num_em_iterations
self.use_mean_prior = container.use_mean_prior
self.use_covar_prior = container.use_covar_prior
self.enable_interpolation = container.enable_interpolation
self.iw_nu = container.iw_nu
self.em_convergence_tol = container.em_convergence_tol
self.warm_start_em = warm_start_em
# M-step shrinkage intensity, re-applied whenever forward() re-runs EM
# during from-scratch fitting (the pretrained/frozen path never re-runs EM).
self._covariance_shrinkage = covariance_shrinkage
# Optional post-EM additive base kernel (set by from_pretrained); default
# None so the attribute is always present.
self._additive_base = None
# Register modules for parameter tracking
self.mean_module = container.mean_module
self.covar_module = container.covar_module
# Store unique inputs
self._unique_inputs_obs = container.unique_inputs
self._X_obs = container.unique_inputs.X_all
self._N_obs = self._X_obs.shape[0]
# Inducing point locations are always stored as a Parameter; gradient
# tracking is toggled to make them either learnable or fixed.
self._X_inducing = torch.nn.Parameter(
container.X_inducing.clone().detach(),
requires_grad=learnable_inducing_points,
)
self._N_inducing = container.X_inducing.shape[0]
# Use the flag from the container that tracks whether inducing points
# were explicitly provided to pretrain_em_prior
self._use_inducing_points = container.use_inducing_points
# Store EM results (these may be detached if from pretrained)
self._mu_inducing = container.mu_inducing
self._Sigma_inducing = container.Sigma_inducing
# Store cached interpolation quantities
self._cached_L_kernel_inducing = container.L_kernel_inducing
self._cached_delta_mu = container.delta_mu
# Store flag to control EM re-running in forward()
self._using_pretrained_prior = using_pretrained
# Store initial values for warm-starting (needed for _run_em)
self._mu_init_inducing = container.mu_inducing
self._Sigma_init_inducing = container.Sigma_inducing
[docs]
@classmethod
def from_pretrained(
cls,
em_prior: EMPriorContainer,
train_X: Tensor,
train_Y: Tensor,
likelihood: Likelihood | None = None,
freeze_pretrained: bool = True,
learnable_inducing_points: bool = False,
base_covar_module: Kernel | None = None,
) -> "EMEmpiricalGaussianProcess":
"""Factory method for creating a model from a pre-trained prior.
This is equivalent to calling the constructor with em_prior, but
provides a clearer API for the pre-training use case.
Args:
em_prior: Pre-trained EMPriorContainer.
train_X: Training inputs for GP conditioning.
train_Y: Training targets for GP conditioning.
likelihood: Optional likelihood. Defaults to GaussianLikelihood.
freeze_pretrained: If True (default), freeze all pre-trained
parameters. Set to False for fine-tuning.
learnable_inducing_points: If True, register inducing point
locations as a learnable nn.Parameter.
base_covar_module: Optional base kernel added additively at conditioning
time. It contributes its full covariance ``K_base(X, X)`` at the query
points, on top of the EM covariance (which is Nyström-interpolated
from the inducing set):
``Sigma(X, X) = Sigma_EM(X, X) + K_base(X, X)``. The base variance is
therefore preserved away from the inducing set, so the model degrades
gracefully toward a standard GP with this kernel as target data grow --
including in higher dimensions where the inducing set is sparse. Which
of its parameters are fit by the marginal likelihood -- jointly with the
observation noise and *without* backprop through EM (``Sigma_inducing``
is fixed) -- is controlled by the ``requires_grad`` flags the caller
sets on it. Pass a fresh ``ScaleKernel(MaternKernel(...))`` with all
parameters trainable for a fully-adaptive base (recommended for automl).
None (default) recovers the pure EM-EGP.
Returns:
EMEmpiricalGaussianProcess instance using the pre-trained prior.
Example:
>>> em_prior = pretrain_em_prior(datasets, mean_module, covar_module)
>>> model = EMEmpiricalGaussianProcess.from_pretrained(
... em_prior, train_X, train_Y
... )
"""
model = cls(
train_X=train_X,
train_Y=train_Y,
em_prior=em_prior,
likelihood=likelihood,
learnable_inducing_points=learnable_inducing_points,
)
# Optionally freeze pre-trained parameters
if freeze_pretrained:
model.freeze_pretrained_parameters()
# Optional additive base kernel at conditioning time: Sigma + K_base(Z, Z).
# Set after freezing so the caller-controlled base params stay trainable.
model._additive_base = base_covar_module
return model
[docs]
def freeze_pretrained_parameters(self) -> None:
"""Freeze all pre-trained parameters to prevent gradient updates.
This method sets requires_grad=False on:
- initial_mean_module parameters (kernel hyperparameters for prior mean)
- initial_covar_module parameters (kernel hyperparameters for prior covariance)
- Pre-computed inducing point quantities (mu_inducing, Sigma_inducing, etc.)
Call this after creating a model with `from_pretrained()` if you want
to ensure the pre-trained prior is completely frozen during any
subsequent training operations.
Example:
>>> model = EMEmpiricalGaussianProcess.from_pretrained(
... prior, train_X, train_Y, likelihood
... )
>>> model.freeze_pretrained_parameters() # Explicit freezing
"""
# Freeze mean module parameters
for param in self.initial_mean_module.parameters():
param.requires_grad_(False)
# Freeze covariance module parameters
for param in self.initial_covar_module.parameters():
param.requires_grad_(False)
# Freeze pre-computed quantities (already detached, but be explicit).
# _init_from_container always sets these attributes, so plain None checks
# suffice (no hasattr guard needed).
if self._mu_inducing is not None:
self._mu_inducing.requires_grad_(False)
if self._Sigma_inducing is not None:
self._Sigma_inducing.requires_grad_(False)
if self._cached_L_kernel_inducing is not None:
self._cached_L_kernel_inducing.requires_grad_(False)
if self._cached_delta_mu is not None:
self._cached_delta_mu.requires_grad_(False)
def _effective_Sigma_inducing(self) -> Tensor:
"""EM covariance at the inducing points, plus an optional additive base."""
base = getattr(self, "_additive_base", None)
if base is None:
return self._Sigma_inducing
k_base = base(self._X_inducing, self._X_inducing)
if hasattr(k_base, "to_dense"):
k_base = k_base.to_dense()
return self._Sigma_inducing + k_base
def _update_cache(self) -> None:
"""Update cached quantities for shift interpolation.
This method recomputes the cached Cholesky factor and mean shift
based on the current _mu_inducing and _Sigma_inducing values.
Should be called after EM is re-run with new hyperparameters.
The EM-learned prior is extended to new points via shift interpolation:
mu(X) = m(X) + W @ Delta_mu
Sigma(X) = Lambda(X) + W @ Sigma_inducing @ W^T
where Delta_mu = mu_inducing - m(X_inducing),
Lambda(X) = K(X,X) - K(X,Z) K(Z,Z)^{-1} K(Z,X) is the Nystrom residual,
and W = K(X, X_inducing) @ K_inducing^{-1}.
"""
K_kernel_inducing = self.initial_covar_module(
self._X_inducing, self._X_inducing
).to_dense()
m_inducing = _evaluate_mean(self.initial_mean_module, self._X_inducing)
self._cached_L_kernel_inducing = psd_safe_cholesky(K_kernel_inducing)
# Mean shift: Delta_mu = mu_inducing - m(X_inducing)
self._cached_delta_mu = self._mu_inducing - m_inducing
def _interpolate_prior_to_X(self, X: Tensor) -> tuple[Tensor, Tensor]:
"""Interpolate the EM-learned prior to arbitrary locations X.
Uses shift interpolation from inducing points Z to query locations X:
μ(X) = m(X) + W @ Δμ
Σ(X) = Λ(X) + W @ Σ_inducing @ W^T [+ K_base(X, X)]
where W = K(X, Z) @ K(Z, Z)^{-1} and
Λ(X) = K(X,X) - K(X,Z) K(Z,Z)^{-1} K(Z,X) is the Nyström residual.
Only the rank-limited, inducing-defined EM covariance is Nyström-
interpolated through W. An optional additive base kernel contributes its
*full* covariance ``K_base(X, X)`` directly at the query points -- not the
Nyström projection ``W K_base(Z, Z) W^T`` -- so its variance is preserved
away from the inducing set (important in higher dimensions where Z is
sparse). At X == Z the two coincide (W == I), so this is consistent with the
direct-indexing path.
Args:
X: (n, d) query locations.
Returns:
mu: (n,) interpolated mean.
Sigma: (n, n) interpolated covariance.
"""
mu, Sigma, _ = _interpolate_prior(
X=X,
mean_module=self.initial_mean_module,
covar_module=self.initial_covar_module,
X_inducing=self._X_inducing,
L_ZZ=self._cached_L_kernel_inducing,
delta_mu=self._cached_delta_mu,
Sigma_inducing=self._Sigma_inducing,
include_cross_covariance=False,
)
base = getattr(self, "_additive_base", None)
if base is not None:
k_base = base(X, X)
if hasattr(k_base, "to_dense"):
k_base = k_base.to_dense()
Sigma = Sigma + k_base
return mu, Sigma
def _get_prior_at_indices(self, indices: Tensor) -> tuple[Tensor, Tensor]:
"""Extract prior mean and covariance at the given indices.
This is a fast O(n²) operation for when X is a known subset of
the inducing points.
Args:
indices: (n,) index tensor into _mu_inducing and _Sigma_inducing.
Returns:
mu: (n,) mean at the indexed locations.
Sigma: (n, n) covariance at the indexed locations.
"""
mu, Sigma, _ = _index_prior(
mu_full=self._mu_inducing,
Sigma_full=self._effective_Sigma_inducing(),
indices=indices,
include_cross_covariance=False,
)
return mu, Sigma
def _get_em_initialization(
self,
) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
"""Get initial values for EM algorithm, with optional warm-starting.
This helper encapsulates the initialization logic for the EM algorithm.
When warm-starting without priors, it skips the expensive kernel evaluation
entirely. Otherwise, it computes Sigma_init for prior matrices and/or
cold-start initialization.
Warm-starting accelerates EM convergence during hyperparameter optimization
by initializing from the previous optimization step's converged values.
The warm-start values are detached to avoid gradient graph issues when
backward() was called between forward() calls.
Returns:
mu: (N_inducing,) initial mean for EM iterations.
Sigma: (N_inducing, N_inducing) initial covariance for EM iterations.
Psi: IW prior scale matrix, or None if not using covariance prior.
K_mu: Kernel prior matrix for mean, or None if not using mean prior.
"""
# Determine if we're warm-starting
is_warm_starting = self.warm_start_em and hasattr(self, "_mu_inducing")
# Determine if we need to compute Sigma_init (expensive kernel evaluation):
# - Always needed for cold-start (to initialize Sigma)
# - Needed when using priors (Psi and K_mu depend on Sigma_init)
needs_sigma_init = (
(not is_warm_starting) or self.use_mean_prior or self.use_covar_prior
)
Sigma_init: Tensor | None = None
if needs_sigma_init:
# Compute fresh Sigma_init for priors and/or cold-start initialization
Sigma_init = self.initial_covar_module(
self._X_inducing, self._X_inducing
).to_dense()
K_mu = Sigma_init if self.use_mean_prior else None
Psi = (
_compute_prior_scale_matrix(Sigma_init, self.iw_nu, self._N_inducing)
if self.use_covar_prior
else None
)
else:
# Warm-starting without priors: skip expensive kernel evaluation
K_mu = None
Psi = None
# Determine initial mu and Sigma for EM iterations
if is_warm_starting:
# Warm-start: detach to avoid gradient graph issues
# The detach is critical because backward() may have been called
# on the previous computation graph, which would have freed it.
mu = self._mu_inducing.detach()
Sigma = self._Sigma_inducing.detach()
else:
# Cold start: use parametric initialization. Sigma_init is guaranteed
# to be set here because cold start (not is_warm_starting) implies
# needs_sigma_init above.
mu = _evaluate_mean(self.initial_mean_module, self._X_inducing)
Sigma = Sigma_init
return mu, Sigma, Psi, K_mu
def _run_em(self) -> tuple[Tensor, Tensor]:
"""Run EM iterations at inducing points (or historical X if no inducing).
This is the core EM algorithm that learns the prior at the fixed set of
inducing points. The results can then be interpolated to any query location
via shift interpolation in forward().
When self.warm_start_em is True, the EM algorithm starts from the
previously converged values (detached) instead of the parametric
initialization. This can significantly accelerate convergence during
hyperparameter optimization, but may produce slightly different gradients
since the warm-start values are detached.
Returns:
mu: (N_inducing,) final mean estimate at inducing points.
Sigma: (N_inducing, N_inducing) final covariance estimate.
"""
# Get initialization (with optional warm-starting) and prior matrices
mu_init, Sigma_init, Psi, K_mu = self._get_em_initialization()
likelihood_noise = getattr(self.likelihood, "noise", None)
# Re-apply M-step covariance shrinkage on EM re-runs (mirrors the initial
# pretrain_em_prior call); the trace-matched target is the parametric gram
# K(Z, Z). None when shrinkage is disabled so the M-step blend is a no-op.
shrinkage_target = None
if self._covariance_shrinkage > 0.0:
shrinkage_target = self.initial_covar_module(
self._X_inducing, self._X_inducing
).to_dense()
# Delegate to shared EM algorithm
return _run_em_algorithm(
datasets=self.datasets,
mu_init=mu_init,
Sigma_init=Sigma_init,
X_inducing=self._X_inducing,
mean_module=self.initial_mean_module,
covar_module=self.initial_covar_module,
likelihood_noise=likelihood_noise,
experiment_indices=self._unique_inputs_obs.experiment_indices,
use_inducing_points=self._use_inducing_points,
num_em_iterations=self.num_em_iterations,
Psi=Psi,
K_mu=K_mu,
iw_nu=self.iw_nu,
em_convergence_tol=self.em_convergence_tol,
N_inducing=self._N_inducing,
covariance_shrinkage=self._covariance_shrinkage,
shrinkage_target=shrinkage_target,
)
[docs]
def forward(self, X: Tensor) -> MultivariateNormal:
"""Compute GP prior at X after running EM iterations.
**Training mode**: Re-runs EM with warm-starting for gradient flow,
then interpolates to X.
**Eval mode**: Uses cached EM results with shift interpolation to X.
When interpolation is enabled (default), shift interpolation extends the
EM-learned prior from inducing points Z to query locations X::
μ(X) = m(X) + W @ Δμ
Σ(X) = Λ(X) + W @ Σ_inducing @ W^T
where W = K(X, Z) @ K(Z, Z)^{-1} and Λ(X) = K(X, X) - W @ K(Z, X) is
the Nyström residual. This numerically stable decomposition is used
instead of the algebraically equivalent K(X, X) + W @ ΔΣ @ W^T; see
``_interpolate_prior``.
When interpolation is disabled, X must be a subset of the historical
observations, and we use direct indexing (faster but restrictive).
Args:
X: (n, d) tensor of query locations.
Returns:
MultivariateNormal distribution at X.
Raises:
ValueError: If interpolation is disabled and X is not a subset of
historical observations.
"""
# Only re-run EM if training AND not using pretrained prior
if self.training and not self._using_pretrained_prior:
# Re-run EM for gradient flow through kernel hyperparameters
mu_em, Sigma_em = self._run_em()
self._mu_inducing = mu_em
self._Sigma_inducing = Sigma_em
self._update_cache()
if self.enable_interpolation or self._use_inducing_points:
# Shift interpolation from inducing points to X
mean, covar = self._interpolate_prior_to_X(X)
else:
# Direct indexing: X must be subset of historical observations
indices = self._find_indices_in_historical(X)
mean, covar = self._get_prior_at_indices(indices)
return MultivariateNormal(mean, to_linear_operator(covar))
def _find_indices_in_historical(self, X: Tensor) -> Tensor:
"""Find indices of X in historical observations.
Args:
X: (n, d) query locations that must be a subset of _X_inducing.
Returns:
indices: (n,) index tensor such that _X_inducing[indices] == X.
Raises:
ValueError: If any point in X is not found in historical observations.
"""
# Compute pairwise distances to find matches
# X: (n, d), _X_inducing: (N, d)
dists = torch.cdist(X, self._X_inducing) # (n, N)
min_dists, indices = dists.min(dim=1)
# Check that all points were found (within numerical tolerance)
tol = 1e-6
not_found = min_dists > tol
if not_found.any():
num_missing = not_found.sum().item()
raise ValueError(
f"{num_missing} query point(s) not found in historical observations. "
f"When not using inducing points, forward(X) requires X to be a "
f"subset of the historical input locations. Consider using "
f"inducing_points to enable interpolation to arbitrary locations."
)
return indices
# =============================================================================
# Shift Interpolation Helpers
# =============================================================================
def _interpolate_prior(
X: Tensor,
mean_module: Mean,
covar_module: Kernel,
X_inducing: Tensor,
L_ZZ: Tensor,
delta_mu: Tensor,
Sigma_inducing: Tensor,
include_cross_covariance: bool = False,
) -> tuple[Tensor, Tensor, Tensor | None]:
"""Interpolate the EM-learned prior from inducing points to query locations.
Uses shift interpolation to extend the prior from inducing points Z to X:
μ(X) = m(X) + K(X,Z) @ K(Z,Z)^{-1} @ Δμ
Σ(X,X) = Λ(X) + W @ Σ_inducing @ W^T
Σ(Z,X) = Σ_inducing @ K(Z,Z)^{-1} @ K(Z,X) (if requested)
where Λ(X) = K(X,X) - K(X,Z) K(Z,Z)^{-1} K(Z,X) is the Nyström residual,
and W = K(X,Z) K(Z,Z)^{-1}.
**Numerical stability**: We use the decomposition Σ(X,X) = Λ(X) + W Σ_inducing W^T
instead of the algebraically equivalent K(X,X) + W ΔΣ W^T (where ΔΣ = Σ_inducing
- K(Z,Z)). The former is guaranteed PSD (sum of two PSD matrices: Λ is PSD by
Schur complement theory, W Σ_inducing W^T is PSD since Σ_inducing is PSD). The
latter suffers from catastrophic cancellation when K(X,X) ≈ W K(Z,Z) W^T.
Args:
X: (n, d) query locations.
mean_module: Parametric mean module.
covar_module: Parametric covariance module.
X_inducing: (M, d) inducing point locations.
L_ZZ: (M, M) Cholesky factor of K(Z, Z).
delta_mu: (M,) mean shift = μ_Z - m(Z).
Sigma_inducing: (M, M) EM-estimated covariance at inducing points.
include_cross_covariance: If True, also compute and return Σ(Z, X).
Returns:
mu: (n,) interpolated mean at X.
Sigma: (n, n) interpolated covariance at X.
cross_covariance: (M, n) cross-covariance Σ(Z, X), or None if not requested.
"""
# Compute kernel matrices. Batch-aware: X may be ``(*batch, q, d)`` (e.g. the
# ``q=1`` t-batches that analytic acquisitions such as LogExpectedImprovement
# feed in), so use ``.mT`` and batched matmuls throughout. On 2D ``(n, d)``
# input these reduce exactly to the original operations.
K_XZ = covar_module(X, X_inducing).to_dense() # (*b, q, M)
K_XX = covar_module(X, X).to_dense() # (*b, q, q)
K_ZX = K_XZ.mT # (*b, M, q)
# Compute V = L_ZZ^{-1} K_ZX and alpha_ZX = K_ZZ^{-1} K_ZX = L_ZZ^{-T} V
# via two triangular solves (L_ZZ broadcasts over any leading batch dims).
V = torch.linalg.solve_triangular(L_ZZ, K_ZX, upper=False)
alpha_ZX = torch.linalg.solve_triangular(L_ZZ.mT, V, upper=True)
# Interpolate mean: μ(X) = m(X) + K_XZ @ K_ZZ^{-1} @ Δμ
alpha_mu = torch.cholesky_solve(delta_mu.unsqueeze(-1), L_ZZ) # (M, 1)
interp = (K_XZ @ alpha_mu).squeeze(-1) # (*b, q)
m_X = mean_module(X)
# Drop a trailing singleton *output* dim if the mean module emits ``(..., n, 1)``,
# while preserving any t-batch/q structure so batched inputs stay ``(*b, q)``.
if m_X.dim() == interp.dim() + 1 and m_X.shape[-1] == 1:
m_X = m_X.squeeze(-1)
mu = m_X + interp
# Interpolate covariance using the numerically stable decomposition:
# Σ(X,X) = Λ(X) + W @ Σ_inducing @ W^T
# where Λ(X) = K(X,X) - Vᵀ V is the Nyström residual (Schur complement of
# K(Z,Z) in the joint kernel — PSD in exact arithmetic). Small floating-point
# errors are handled downstream by psd_safe_cholesky jitter (forward path) and
# the σ²I noise buffer (E-step path).
W = alpha_ZX.mT # (*b, q, M)
Lambda = K_XX - V.mT @ V
Sigma = Lambda + W @ Sigma_inducing @ W.mT
# Optionally compute cross-covariance: Σ(Z,X) = Σ_inducing @ K(Z,Z)^{-1} @ K(Z,X)
cross_covariance = None
if include_cross_covariance:
cross_covariance = Sigma_inducing @ alpha_ZX
return mu, Sigma, cross_covariance
def _index_prior(
mu_full: Tensor,
Sigma_full: Tensor,
indices: Tensor,
include_cross_covariance: bool = False,
) -> tuple[Tensor, Tensor, Tensor | None]:
"""Extract prior mean and covariance at the given indices.
This is a fast O(n²) operation for when X is a known subset of
the inducing points.
Args:
mu_full: (N,) full mean vector at all inducing points.
Sigma_full: (N, N) full covariance matrix at all inducing points.
indices: (n,) index tensor into mu_full and Sigma_full.
include_cross_covariance: If True, also return cross-covariance Σ(all, indices).
Returns:
mu: (n,) mean at the indexed locations.
Sigma: (n, n) covariance at the indexed locations.
cross_covariance: (N, n) cross-covariance, or None if not requested.
"""
mu = mu_full[indices]
Sigma = Sigma_full[indices][:, indices]
cross_covariance = None
if include_cross_covariance:
cross_covariance = Sigma_full[:, indices]
return mu, Sigma, cross_covariance
# =============================================================================
# Marginal Log-Likelihood Classes
# =============================================================================
[docs]
class EMEmpiricalMarginalLogLikelihood(MarginalLogLikelihood):
"""Marginal log-likelihood for EMEmpiricalGaussianProcess.
Computes the sum of marginal log-likelihoods across ALL K datasets used
for empirical prior estimation.
Compatible with BoTorch's fit_gpytorch_mll.
Example:
>>> model = EMEmpiricalGaussianProcess(...)
>>> mll = EMEmpiricalMarginalLogLikelihood(model.likelihood, model)
>>> output = model(train_X)
>>> loss = -mll(output, train_Y)
>>> loss.backward()
"""
def _dataset_mll(self, mu_S: Tensor, Sigma_SS: Tensor, Y: Tensor) -> Tensor:
"""Compute the GP marginal log-likelihood log N(Y | μ_S, Σ_SS + σ²I).
Args:
mu_S: (n,) prior mean at observation locations.
Sigma_SS: (n, n) prior covariance at observation locations.
Y: (n, 1) or (n,) observed targets.
Returns:
mll: Scalar log marginal likelihood for this dataset.
"""
# to_linear_operator improves numerical stability (adds jitter if needed).
prior_dist = MultivariateNormal(mu_S, to_linear_operator(Sigma_SS))
output_dist = self.likelihood(prior_dist)
return output_dist.log_prob(Y.squeeze(-1))
[docs]
def forward(
self,
function_dist: MultivariateNormal,
target: Tensor,
) -> Tensor:
"""Compute the sum of marginal log-likelihoods over all K datasets.
Note: function_dist and target arguments are NOT used. Instead, this
method runs EM with current hyperparameters to get the refined prior.
When warm_start_em is enabled on the model, the EM algorithm is
initialized from the previously converged values, which can accelerate
convergence during hyperparameter optimization.
Args:
function_dist: Ignored (kept for API compatibility).
target: Ignored (kept for API compatibility).
Returns:
mll: Scalar log marginal likelihood, summed across datasets (plus
any hyperparameter prior terms) and normalized per observation
(divided by the total number of observations across datasets).
"""
# When using a pretrained prior, skip the expensive EM re-run.
# The EM-estimated mu and Sigma are already cached (and detached).
# We only re-evaluate the kernel at inducing points so that gradients
# flow through the kernel hyperparameters (for coordinate ascent).
if not self.model._using_pretrained_prior:
mu_em, Sigma_em = self.model._run_em()
self.model._mu_inducing = mu_em
self.model._Sigma_inducing = Sigma_em
# Always refresh the interpolation cache with current kernel params.
# This re-evaluates m_phi(Z) and K_phi(Z,Z), enabling gradient flow
# through the kernel to the observed-data MLL below, while the
# detached mu_em and Sigma_em block gradients through EM iterations.
self.model._update_cache()
# Compute sum of MLLs over all K datasets
total_mll = torch.tensor(
0.0,
device=self.model._mu_inducing.device,
dtype=self.model._mu_inducing.dtype,
)
total_data_points = 0
# Get experiment indices for direct indexing case (if applicable)
experiment_indices = (
None
if self.model._use_inducing_points
else self.model._unique_inputs_obs.experiment_indices
)
for i, dataset in enumerate(self.model.datasets):
# Get prior at observation locations via interpolation or direct indexing
if self.model._use_inducing_points:
mu_S, Sigma_SS = self.model._interpolate_prior_to_X(dataset.X)
else:
mu_S, Sigma_SS = self.model._get_prior_at_indices(experiment_indices[i])
total_mll = total_mll + self._dataset_mll(mu_S, Sigma_SS, dataset.Y)
total_data_points += dataset.X.shape[0]
# Add log probs of priors on hyperparameters
for _, module, prior, closure, _ in self.model.named_priors():
prior_term = prior.log_prob(closure(module))
total_mll = total_mll + prior_term.sum()
# Scale by total amount of data
return total_mll / total_data_points
# =============================================================================
# Helper Functions for fit_gpytorch_mll Compatibility
# =============================================================================
[docs]
def build_shared_gp_model_list(
datasets: list[ExperimentDataset],
mean_module: Mean,
covar_module: Kernel,
observation_noise: float | None = None,
) -> tuple[ModelListGP, GPyTorchSumMarginalLogLikelihood]:
"""Build a ModelListGP with shared mean/kernel across all GPs.
All GPs in the returned ModelList share the SAME mean_module and covar_module
instances, so optimizing the ModelList's MLL optimizes a single set of
hyperparameters using gradients from all K datasets.
This uses BoTorch's ModelListGP which provides full compatibility with
fit_gpytorch_mll, including transform_inputs and other BoTorch model methods.
Args:
datasets: K ExperimentDataset objects with potentially different sizes.
mean_module: Shared mean module (same instance used by all GPs).
covar_module: Shared kernel module (same instance used by all GPs).
observation_noise: Optional observation noise variance. If provided,
the noise is fixed (not optimized). If None, noise is learned.
Returns:
Tuple of:
- model_list: ModelListGP of K SingleTaskGP instances.
- mll: GPyTorch SumMarginalLogLikelihood ready for fit_gpytorch_mll.
Example:
>>> from botorch.fit import fit_gpytorch_mll
>>> from gpytorch.kernels import ScaleKernel, RBFKernel
>>> from gpytorch.means import ZeroMean
>>>
>>> # Create shared modules
>>> mean = ZeroMean()
>>> kernel = ScaleKernel(RBFKernel())
>>>
>>> # Build model list with shared modules
>>> model_list, mll = build_shared_gp_model_list(
... datasets, mean, kernel, observation_noise=1e-2
... )
>>>
>>> # Optimize hyperparameters using fit_gpytorch_mll
>>> fit_gpytorch_mll(mll)
>>>
>>> # The shared modules now have optimized parameters
>>> print(kernel.base_kernel.lengthscale)
"""
models = []
for dataset in datasets:
# Create a SingleTaskGP for this dataset
# Key: pass the SAME mean_module and covar_module instances to all GPs
gp = SingleTaskGP(
train_X=dataset.X,
train_Y=dataset.Y,
mean_module=mean_module, # Shared across all GPs
covar_module=covar_module, # Shared across all GPs
)
# Set observation noise if provided
if observation_noise is not None:
gp.likelihood.noise = observation_noise
gp.likelihood.noise_covar.raw_noise.requires_grad_(False)
models.append(gp)
# Create ModelListGP (BoTorch wrapper with full fit_gpytorch_mll compatibility)
model_list = ModelListGP(*models)
# Create SumMarginalLogLikelihood (GPyTorch's version for fit_gpytorch_mll)
mll = GPyTorchSumMarginalLogLikelihood(model_list.likelihood, model_list)
return model_list, mll