Gaussian Estimator

Full covariance Gaussian posterior estimation.

Overview

GaussianFullCov 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

GaussianFullCov requires a Product prior with "standard_normal" mode.

Configuration

Like Flow, all GaussianFullCov parameters are specified flat directly under estimator: in YAML — no nested group keys.

estimator:
  _target_: falcon.estimators.GaussianFullCov
  max_epochs: 1000
  batch_size: 128
  patience_epochs: 32
  hidden_dim: 128
  num_layers: 3
  momentum: 0.01
  min_var: 1.0e-20
  eig_update_freq: 1
  embedding:
    _target_: model.E_identity
    _input_: [x]
  lr: 0.01
  lr_decay_factor: 1.0
  lr_patience_epochs: 8
  gamma: 0.5
  discard_samples: false
  log_ratio_threshold: -20.0

Configuration Reference

Network Parameters

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

The training loop, optimizer, and inference parameters (max_rounds, patience_rounds, max_epochs, patience_epochs, val_every_epochs, batch_size, lr, lr_decay_factor, lr_patience_epochs, prior_rounds, gamma, discard_samples, log_ratio_threshold, etc.) are identical to those in Flow; see Training Loop for how rounds work. GaussianFullCov has a single network group: the Gaussian posterior together with its embedding.

gamma for GaussianFullCov

Unlike Flow, where gamma controls importance-sampling breadth, in GaussianFullCov it controls eigenvalue tempering of the covariance matrix. Smaller gamma (e.g. 0.1) produces a broader proposal relative to the posterior; the relationship is not the same as in 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.GaussianFullCov
      max_epochs: 8000
      batch_size: 128
      patience_epochs: 128
      hidden_dim: 128
      num_layers: 3
      momentum: 0.01
      min_var: 1.0e-20
      eig_update_freq: 1
      embedding:
        _target_: model.E_identity
        _input_: [x]
      lr: 0.01
      lr_decay_factor: 1.0
      lr_patience_epochs: 8
      gamma: 0.1
      discard_samples: false
      log_ratio_threshold: -20.0

    ray:
      num_train_gpus: 0.5
      num_sample_gpus: 0.5

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

sample:
  posterior:
    n: 1000

Class Reference

GaussianFullCov

GaussianFullCov(*, max_epochs=100, lr=0.01, gamma=0.5, embedding=None, device=None, batch_size=128, patience_epochs=16, val_every_epochs=1, max_rounds=None, patience_rounds=10, prior_rounds=0, cache_on_device=False, max_cache_samples=0, hidden_dim=128, num_layers=3, momentum=0.01, min_var=1e-20, eig_update_freq=1, betas=(0.9, 0.9), lr_decay_factor=1.0, lr_patience_epochs=8, discard_samples=False, log_ratio_threshold=-20.0)

Bases: TorchModel

Full-covariance Gaussian posterior estimator for TransformedPrior simulators.

Works in the standard-normal latent space; samples are mapped back to parameter space after generation. Training runs in rounds of epochs on fixed data; see RoundTrainer.

Parameters:

Name	Type	Description	Default
max_epochs	int	Maximum epochs per round.	100
lr	float	Learning rate; reset at the start of every round.	0.01
gamma	float	Proposal tempering coefficient.	0.5
embedding		Embedding config dict or None.	None
device		Device string; auto-detected if None.	None
batch_size	int	Mini-batch size.	128
patience_epochs	int	End the round once the best validation loss is this many epochs old (checked at validations).	16
val_every_epochs	int	Validate after every N-th epoch (and at the last epoch).	1
max_rounds	Optional[int]	Maximum number of rounds (None = unlimited).	None
patience_rounds	int	Stop training after this many rejected rounds in a row.	10
prior_rounds	int	Rounds that simulate from the prior before switching to the learned proposal. The first round always uses the prior.	0
cache_on_device	bool	Cache training data on the estimator's device.	False
max_cache_samples	int	Cap on cached training samples (0 = all).	0
hidden_dim	int	MLP hidden layer width.	128
num_layers	int	MLP depth.	3
momentum	float	EMA momentum for running statistics.	0.01
min_var	float	Minimum variance for numerical stability.	1e-20
eig_update_freq	int	Eigendecomposition update frequency.	1
betas	tuple	AdamW beta coefficients.	(0.9, 0.9)
lr_decay_factor	float	LR decay factor (1.0 = no decay).	1.0
lr_patience_epochs	int	Epochs without validation improvement before LR decay.	8
discard_samples	bool	Run a discard sweep after every accepted round.	False
log_ratio_threshold	float	Log-ratio cutoff for discarding.	-20.0

Source code in src/falcon/estimators/gaussian_fullcov.py

def __init__(
    self,
    *,
    # Most commonly changed
    max_epochs: int = 100,
    lr: float = 1e-2,
    gamma: float = 0.5,
    embedding=None,
    device=None,
    # Training loop
    batch_size: int = 128,
    patience_epochs: int = 16,
    val_every_epochs: int = 1,
    max_rounds: Optional[int] = None,
    patience_rounds: int = 10,
    prior_rounds: int = 0,
    cache_on_device: bool = False,
    max_cache_samples: int = 0,
    # Network architecture
    hidden_dim: int = 128,
    num_layers: int = 3,
    momentum: float = 0.01,
    min_var: float = 1e-20,
    eig_update_freq: int = 1,
    # Optimizer
    betas: tuple = (0.9, 0.9),
    lr_decay_factor: float = 1.0,
    lr_patience_epochs: int = 8,
    # Inference / sampling
    discard_samples: bool = False,
    log_ratio_threshold: float = -20.0,
):
    self.max_epochs = max_epochs
    self.lr = lr
    self.gamma = gamma
    self.embedding = embedding
    self.device = device
    self.batch_size = batch_size
    self.patience_epochs = patience_epochs
    self.val_every_epochs = val_every_epochs
    self.max_rounds = max_rounds
    self.patience_rounds = patience_rounds
    self.prior_rounds = prior_rounds
    self.cache_on_device = cache_on_device
    self.max_cache_samples = max_cache_samples
    self.hidden_dim = hidden_dim
    self.num_layers = num_layers
    self.momentum = momentum
    self.min_var = min_var
    self.eig_update_freq = eig_update_freq
    self.betas = betas
    self.lr_decay_factor = lr_decay_factor
    self.lr_patience_epochs = lr_patience_epochs
    self.discard_samples = discard_samples
    self.log_ratio_threshold = log_ratio_threshold

build

build(init_tree)

Source code in src/falcon/estimators/gaussian_fullcov.py

def build(self, init_tree) -> None:
    from falcon.estimators.embedded_posterior import EmbeddedPosterior
    from falcon.embeddings import instantiate_embedding

    theta = self._to_tensor(init_tree["theta"])
    conditions = {k: self._to_tensor(v, self.device) for k, v in init_tree["conditions"].items()}
    theta_latent = self.simulator_instance.inverse(theta, mode="standard_normal")

    embedding = instantiate_embedding(self.embedding).to(self.device)
    embedding.eval()
    with torch.no_grad():
        embedded = embedding(conditions)

    posterior = _GaussianPosterior(
        param_dim=theta_latent.shape[1],
        condition_dim=embedded.shape[1],
        hidden_dim=self.hidden_dim,
        num_layers=self.num_layers,
        momentum=self.momentum,
        min_var=self.min_var,
        eig_update_freq=self.eig_update_freq,
    ).to(self.device)

    self._model = EmbeddedPosterior(embedding, posterior)
    self._set_groups({
        "model": NetworkGroup({"model": self._model}, "loss", embedding=self._model.embedding),
    })
    debug(f"GaussianFullCov model built: param_dim={theta_latent.shape[1]}")

train_step

train_step(batch)

Source code in src/falcon/estimators/gaussian_fullcov.py

def train_step(self, batch) -> Dict[str, float]:
    if self._optimizer is None:
        self._build_optimizer()
    self._optimizer.zero_grad()
    loss, metrics = self._compute_loss(batch, train=True)
    loss.backward()
    self._optimizer.step()
    return metrics

val_step

val_step(batch)

Source code in src/falcon/estimators/gaussian_fullcov.py

def val_step(self, batch) -> Dict[str, float]:
    _, metrics = self._compute_loss(batch, train=False)
    return metrics

discard_test

discard_test(batch)

Source code in src/falcon/estimators/gaussian_fullcov.py

def discard_test(self, batch):
    theta_latent, theta_logprob, conditions = self._unpack_batch(batch)
    self._model.eval()
    s = self._summary("model", conditions)
    log_prob = self._model.posterior.log_prob(theta_latent, s).cpu()
    return (log_prob - theta_logprob.cpu()) < self.log_ratio_threshold