Gaussian Estimator

Full covariance Gaussian posterior estimation.

Overview

Gaussian provides a simpler alternative to Flow for posterior estimation. Instead of normalizing flows, it models the posterior as a multivariate Gaussian with full covariance, using eigenvalue-based operations.

Key features:

• Full covariance matrix showing parameter correlations directly
• Eigenvalue-based tempered sampling for exploration
• Simpler and more interpretable than flow-based methods

Note

Gaussian requires a Product prior (with "standard_normal" mode) as the simulator, not Hypercube.

Configuration

Gaussian is configured through the same four groups as Flow: loop, network, optimizer, and inference.

estimator:
  _target_: falcon.estimators.Gaussian
  loop:
    num_epochs: 1000
    batch_size: 128
    early_stop_patience: 32
  network:
    hidden_dim: 128
    num_layers: 3
    momentum: 0.10
    min_var: 1.0e-20
    eig_update_freq: 1
  embedding:
    _target_: model.E_identity
    _input_: [x]
  optimizer:
    lr: 0.01
    lr_decay_factor: 0.5
    scheduler_patience: 16
  inference:
    gamma: 1.0
    discard_samples: false
    log_ratio_threshold: -20.0

Configuration Reference

Network (network)

Parameter	Type	Default	Description
hidden_dim	int	128	MLP hidden layer dimension
num_layers	int	3	Number of hidden layers
momentum	float	0.10	EMA momentum for running statistics
min_var	float	1e-20	Minimum variance for numerical stability
eig_update_freq	int	1	Eigendecomposition update frequency

The loop, optimizer, and inference groups share the same parameters as Flow.

Complete Example

graph:
  z:
    evidence: [x]

    simulator:
      _target_: falcon.priors.Product
      priors:
        - ['normal', 0.0, 1.0]
        - ['normal', 0.0, 1.0]
        - ['normal', 0.0, 1.0]

    estimator:
      _target_: falcon.estimators.Gaussian
      loop:
        num_epochs: 1000
        batch_size: 128
        early_stop_patience: 32
      network:
        hidden_dim: 128
        num_layers: 3
        momentum: 0.10
        min_var: 1.0e-20
        eig_update_freq: 1
      embedding:
        _target_: model.E_identity
        _input_: [x]
      optimizer:
        lr: 0.01
        lr_decay_factor: 0.5
        scheduler_patience: 16
      inference:
        gamma: 1.0
        discard_samples: false
        log_ratio_threshold: -20.0

    ray:
      num_gpus: 1

  x:
    parents: [z]
    simulator:
      _target_: model.ExpPlusNoise
      sigma: 1.0e-6
    observed: "./data/mock_data.npz['x']"

sample:
  posterior:
    n: 1000

Class Reference

Gaussian

Gaussian(simulator_instance, theta_key=None, condition_keys=None, config=None)

Create a LossBasedEstimator with GaussianPosterior.

This is the main entry point for using Gaussian posterior estimation. It provides sensible defaults while allowing full customization.

Parameters:

Name	Type	Description	Default
simulator_instance		Prior/simulator instance	required
theta_key	Optional[str]	Key for theta in batch data	None
condition_keys	Optional[List[str]]	Keys for condition data in batch	None
config	Optional[dict]	Configuration dict with sections: - loop: TrainingLoopConfig options - network: GaussianPosteriorConfig options - optimizer: OptimizerConfig options - inference: InferenceConfig options - embedding: Embedding configuration with target - device: Device string (optional)	None

Returns:

Type	Description
LossBasedEstimator	Configured LossBasedEstimator ready for training

Example YAML

estimator: target: falcon.estimators.Gaussian network: hidden_dim: 128 num_layers: 3 embedding: target: model.E input: [x]

Source code in falcon/estimators/gaussian.py

def Gaussian(
    simulator_instance,
    theta_key: Optional[str] = None,
    condition_keys: Optional[List[str]] = None,
    config: Optional[dict] = None,
) -> LossBasedEstimator:
    """Create a LossBasedEstimator with GaussianPosterior.

    This is the main entry point for using Gaussian posterior estimation.
    It provides sensible defaults while allowing full customization.

    Args:
        simulator_instance: Prior/simulator instance
        theta_key: Key for theta in batch data
        condition_keys: Keys for condition data in batch
        config: Configuration dict with sections:
            - loop: TrainingLoopConfig options
            - network: GaussianPosteriorConfig options
            - optimizer: OptimizerConfig options
            - inference: InferenceConfig options
            - embedding: Embedding configuration with _target_
            - device: Device string (optional)

    Returns:
        Configured LossBasedEstimator ready for training

    Example YAML:
        estimator:
          _target_: falcon.estimators.Gaussian
          network:
            hidden_dim: 128
            num_layers: 3
          embedding:
            _target_: model.E
            _input_: [x]
    """
    # Check simulator supports transformation interface
    if not isinstance(simulator_instance, TransformedPrior):
        raise TypeError(
            f"Gaussian requires a TransformedPrior (e.g., Product), "
            f"got {type(simulator_instance).__name__}. "
            f"The simulator must support forward/inverse with mode='standard_normal'."
        )

    # Merge with defaults
    schema = OmegaConf.structured(GaussianConfig)
    cfg = OmegaConf.merge(schema, config or {})

    # Extract configs as plain dicts
    embedding_config = OmegaConf.to_container(cfg.embedding, resolve=True)
    posterior_config = OmegaConf.to_container(cfg.network, resolve=True)

    return LossBasedEstimator(
        simulator_instance=simulator_instance,
        posterior_cls=GaussianPosterior,
        embedding_config=embedding_config,
        loop_config=cfg.loop,
        optimizer_config=cfg.optimizer,
        inference_config=cfg.inference,
        posterior_config=posterior_config,
        theta_key=theta_key,
        condition_keys=condition_keys,
        device=cfg.device,
        latent_mode="standard_normal",  # GaussianPosterior assumes N(0,I) prior
    )

GaussianPosterior

GaussianPosterior(param_dim, condition_dim, hidden_dim=128, num_layers=3, momentum=0.01, min_var=1e-06, eig_update_freq=1)

Bases: Module

Full covariance Gaussian posterior with eigenvalue-based operations.

Implements the Posterior contract

• loss(theta, conditions) -> Tensor
• sample(conditions, gamma=None) -> Tensor
• log_prob(theta, conditions) -> Tensor

The posterior is parameterized as

p(theta | conditions) = N(mu(conditions), Sigma)

where mu(conditions) is predicted by an MLP with whitening, and Sigma is the residual covariance matrix estimated from training data.

Source code in falcon/estimators/gaussian.py

def __init__(
    self,
    param_dim: int,
    condition_dim: int,
    hidden_dim: int = 128,
    num_layers: int = 3,
    momentum: float = 0.01,
    min_var: float = 1e-6,
    eig_update_freq: int = 1,
):
    super().__init__()
    self.param_dim = param_dim
    self.condition_dim = condition_dim
    self.momentum = momentum
    self.min_var = min_var
    self.eig_update_freq = eig_update_freq
    self._step_counter = 0

    # MLP for mean prediction
    self.net = build_mlp(condition_dim, hidden_dim, param_dim, num_layers)

    # Input statistics (conditions) - diagonal whitening
    self.register_buffer("_input_mean", torch.zeros(condition_dim))
    self.register_buffer("_input_std", torch.ones(condition_dim))

    # Output statistics (theta) - diagonal whitening
    self.register_buffer("_output_mean", torch.zeros(param_dim))
    self.register_buffer("_output_std", torch.ones(param_dim))

    # Residual covariance (prediction error) - full covariance
    self.register_buffer("_residual_cov", torch.eye(param_dim))
    self.register_buffer("_residual_eigvals", torch.ones(param_dim))
    self.register_buffer("_residual_eigvecs", torch.eye(param_dim))

loss

loss(theta, conditions)

Compute negative log likelihood loss, updating statistics only during training.

Source code in falcon/estimators/gaussian.py

def loss(self, theta: torch.Tensor, conditions: torch.Tensor) -> torch.Tensor:
    """Compute negative log likelihood loss, updating statistics only during training."""
    if self.training:
        self._update_stats(theta, conditions)
        self._update_residual_cov(theta, conditions)
    log_prob = self.log_prob(theta, conditions)
    return -log_prob.mean()

log_prob

log_prob(theta, conditions)

Compute Gaussian log probability using eigendecomposition.

Source code in falcon/estimators/gaussian.py

def log_prob(self, theta: torch.Tensor, conditions: torch.Tensor) -> torch.Tensor:
    """Compute Gaussian log probability using eigendecomposition."""
    mean = self._forward_mean(conditions)
    residuals = theta - mean

    V = self._residual_eigvecs.detach()
    d = self._residual_eigvals.detach()

    log_det = torch.log(d).sum()
    r_proj = V.T @ residuals.T
    mahal = (r_proj**2 / d.unsqueeze(1)).sum(dim=0)

    return -0.5 * (self.param_dim * np.log(2 * np.pi) + log_det + mahal)

sample

sample(conditions, gamma=None)

Sample from posterior, optionally tempered.

Parameters:

Name	Type	Description	Default
conditions	Tensor	Condition tensor of shape (batch, condition_dim)	required
gamma	Optional[float]	Tempering parameter. None = untempered, <1 = widened.	None

Source code in falcon/estimators/gaussian.py

def sample(self, conditions: torch.Tensor, gamma: Optional[float] = None) -> torch.Tensor:
    """Sample from posterior, optionally tempered.

    Args:
        conditions: Condition tensor of shape (batch, condition_dim)
        gamma: Tempering parameter. None = untempered, <1 = widened.
    """
    mean = self._forward_mean(conditions)
    V = self._residual_eigvecs
    d = self._residual_eigvals

    if gamma is None:
        eps = torch.randn_like(mean)
        return mean + (V @ (torch.sqrt(d).unsqueeze(1) * eps.T)).T

    # Tempered sampling
    a = gamma / (1 + gamma)
    lambda_like = (1.0 / d - 1.0).clamp(min=0)
    lambda_prop = a * lambda_like + 1.0
    var_prop = 1.0 / lambda_prop

    mean_proj = V.T @ mean.T
    alpha = a / (d * lambda_prop)
    mean_prop = (V @ (alpha.unsqueeze(1) * mean_proj)).T

    eps = torch.randn_like(mean)
    return mean_prop + (V @ (torch.sqrt(var_prop).unsqueeze(1) * eps.T)).T

Configuration Classes

GaussianConfig dataclass

GaussianConfig(loop=TrainingLoopConfig(), network=GaussianPosteriorConfig(), optimizer=_default_optimizer_config(), inference=InferenceConfig(), embedding=None, device=None)

Top-level Gaussian estimator configuration.

GaussianPosteriorConfig dataclass

GaussianPosteriorConfig(hidden_dim=128, num_layers=3, momentum=0.01, min_var=1e-20, eig_update_freq=1)

Configuration for GaussianPosterior network.