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
  early_stop_patience: 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: 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_epochs, batch_size, early_stop_patience, lr, lr_decay_factor, lr_patience, prior_epochs, gamma, discard_samples, log_ratio_threshold, etc.) are identical to those in Flow.

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
      early_stop_patience: 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: 8
      gamma: 0.1
      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

GaussianFullCov

GaussianFullCov(*, max_epochs=100, lr=0.01, gamma=0.5, embedding=None, device=None, batch_size=128, early_stop_patience=16, prior_epochs=0, cache_on_device=False, cache_sync_every=0, 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=8, discard_samples=False, log_ratio_threshold=-20.0)

Bases: StepwiseEstimator

Full-covariance Gaussian posterior estimator for TransformedPrior simulators.

Works in the standard-normal latent space; samples are mapped back to parameter space after generation.

Parameters:

Name	Type	Description	Default
max_epochs	int	Maximum training epochs.	100
lr	float	Learning rate.	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
early_stop_patience	int	Epochs without improvement before stopping.	16
prior_epochs	int	Epochs to sample from prior before switching to proposal.	0
cache_on_device	bool	Cache training data on the estimator's device.	False
cache_sync_every	int	Resync buffer cache every N epochs (0 = every epoch).	0
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	int	Plateau patience before LR decay.	8
discard_samples	bool	Discard low log-ratio training samples.	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,
    early_stop_patience: int = 16,
    prior_epochs: int = 0,
    cache_on_device: bool = False,
    cache_sync_every: int = 0,
    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: 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.early_stop_patience = early_stop_patience
    self.prior_epochs = prior_epochs
    self.cache_on_device = cache_on_device
    self.cache_sync_every = cache_sync_every
    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 = lr_patience
    self.discard_samples = discard_samples
    self.log_ratio_threshold = log_ratio_threshold

train_step

train_step(batch)

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

def train_step(self, batch) -> Dict[str, float]:
    if not self.networks_initialized:
        self._initialize_model(batch)

    self._optimizer.zero_grad()
    self._model.train()
    loss, metrics = self._compute_loss(batch)
    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]:
    with torch.no_grad():
        self._model.eval()
        _, metrics = self._compute_loss(batch)
    return metrics

sample_prior

sample_prior(num_samples, conditions=None)

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

def sample_prior(self, num_samples: int, conditions=None) -> dict:
    if conditions:
        raise ValueError("Conditions are not supported for sample_prior.")
    samples = self.simulator_instance.simulate_batch(num_samples)
    return {"value": samples, "log_prob": np.zeros(num_samples)}

sample_posterior

sample_posterior(num_samples, conditions=None)

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

def sample_posterior(self, num_samples: int, conditions=None) -> dict:
    return self._sample(num_samples, conditions, gamma=self._posterior_gamma)

sample_proposal

sample_proposal(num_samples, conditions=None)

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

def sample_proposal(self, num_samples: int, conditions=None) -> dict:
    if self._total_epochs_trained < self.prior_epochs:
        return self.sample_prior(num_samples)
    result = self._sample(num_samples, conditions, gamma=self._proposal_gamma)
    log({
        "sample_proposal:mean": result["value"].mean(),
        "sample_proposal:std": result["value"].std(),
        "sample_proposal:logprob": result["log_prob"].mean(),
    })
    return result

save

save(node_dir)

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

def save(self, node_dir) -> None:
    node_dir = Path(node_dir)
    if not self.networks_initialized:
        raise RuntimeError("Cannot save: model not initialised.")

    torch.save(self._best_model.state_dict(), node_dir / "model.pth")
    torch.save(
        {"theta": self._init_theta, "conditions": self._init_conditions},
        node_dir / "init_tensors.pth",
    )
    torch.save(self._total_epochs_trained, node_dir / "total_epochs_trained.pth")

    torch.save(self.history["train_ids"], node_dir / "train_id_history.pth")
    torch.save(self.history["val_ids"], node_dir / "validation_id_history.pth")
    torch.save(self.history["epochs"], node_dir / "epochs.pth")
    torch.save(self.history["train_loss"], node_dir / "loss_train_posterior.pth")
    torch.save(self.history["val_loss"], node_dir / "loss_val_posterior.pth")
    torch.save(self.history["n_samples"], node_dir / "n_samples_total.pth")
    torch.save(self.history["elapsed_min"], node_dir / "elapsed_minutes.pth")

load

load(node_dir)

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

def load(self, node_dir) -> None:
    node_dir = Path(node_dir)

    data = torch.load(node_dir / "init_tensors.pth")
    self._init_theta = data["theta"]
    self._init_conditions = data["conditions"]
    self._model = self._create_model(self._init_theta, self._init_conditions)
    self._best_model = self._create_model(self._init_theta, self._init_conditions)

    saved_state = torch.load(node_dir / "model.pth")
    self._best_model.load_state_dict(saved_state)
    self._model.load_state_dict(saved_state)

    self._build_optimizer()
    self.networks_initialized = True

    tep = node_dir / "total_epochs_trained.pth"
    self._total_epochs_trained = torch.load(tep) if tep.exists() else 0