Flow Estimator

Flow-based posterior estimation using normalizing flows.

Overview

Flow is the primary estimator in Falcon for learning posterior distributions. It uses dual normalizing flows (conditional and marginal) with importance sampling for adaptive proposal generation.

Key features:

• Dual flow architecture for posterior and proposal sampling
• Parameter space normalization via hypercube mapping
• Importance sampling with effective sample size monitoring
• Automatic learning rate scheduling and early stopping

Configuration

All Flow parameters are specified flat directly under estimator: in YAML. There are no nested group keys — everything is a top-level keyword argument to Flow.__init__.

estimator:
  _target_: falcon.estimators.Flow
  max_epochs: 300
  net_type: nsf
  lr: 0.01
  gamma: 0.5
  embedding:
    _target_: model.MyEmbedding
    _input_: [x]

Configuration Reference

Training Loop

Controls the training process.

Parameter	Type	Default	Description
max_epochs	int	100	Maximum training epochs
batch_size	int	128	Training batch size
early_stop_patience	int	16	Epochs without improvement before stopping
cache_sync_every	int	0	Epochs between cache syncs with the buffer (0 = every epoch)
max_cache_samples	int	0	Maximum samples to cache (0 = cache all available)
cache_on_device	bool	false	Keep cached training data on the estimator's device (e.g. GPU)
prior_epochs	int	0	Epochs to sample from prior before switching to proposal
device	str	null	Device string (e.g. "cuda:0"); auto-detected if null

Data Caching

Training data is loaded into a local cache that is periodically synced with the shared simulation buffer. This avoids repeated remote data fetches and allows fast random-access batching.

• cache_sync_every: Controls how often the cache pulls new samples from the buffer. A value of 0 (default) syncs every epoch. Higher values reduce sync overhead at the cost of slightly stale data, which can be useful when simulations are slow.
• max_cache_samples: Caps the number of samples held in the cache. Set to 0 to cache everything. A positive value randomly subsamples, which helps limit GPU memory usage for very large buffers.
• cache_on_device: When true, cached tensors are moved to the estimator's device (typically GPU) once during sync rather than per-batch. This eliminates CPU-to-GPU transfer overhead during training but increases device memory usage.

Network Architecture

Defines the neural network structure.

Parameter	Type	Default	Description
net_type	str	zuko_nice	Flow architecture (see FlowDensity for all types)
theta_norm	bool	true	Normalize parameter space
norm_momentum	float	0.01	Momentum for online normalization updates
use_log_update	bool	false	Use log-space variance updates
adaptive_momentum	bool	false	Sample-dependent momentum

Embedding

The embedding network processes observations before they enter the flow.

embedding:
  _target_: model.MyEmbedding
  _input_: [x]

See Embeddings for details on the declarative embedding system, including multi-input and nested pipeline configurations.

Optimizer

Controls learning rate and scheduling.

Parameter	Type	Default	Description
lr	float	0.01	Initial learning rate
lr_decay_factor	float	0.1	LR multiplier when plateau detected
lr_patience	int	8	Epochs without improvement before LR decay
betas	tuple	(0.9, 0.9)	AdamW beta coefficients

Inference

Controls posterior sampling and amortization.

Parameter	Type	Default	Description
gamma	float	0.5	Amortization mixing coefficient (0=focused, 1=amortized)
discard_samples	bool	true	Discard low-likelihood samples during training
log_ratio_threshold	float	-20	Log-likelihood threshold for sample discarding
sample_reference_posterior	bool	false	Sample from reference posterior
use_best_models	bool	true	Use best validation model for sampling
num_proposals	int	256	Candidate samples drawn from the flow for importance sampling
reference_samples	int	128	Samples used to evaluate the reference posterior
hypercube_bound	float	2.0	Out-of-bounds threshold in hypercube space
out_of_bounds_penalty	float	100.0	Log-weight penalty applied to out-of-bounds proposals
nan_replacement	float	-100.0	Log-weight substituted for NaN values during importance sampling

Understanding gamma

gamma controls how broadly the proposal distribution covers the parameter space:

• gamma=0: Proposal equals the posterior — tightest possible proposal, best for refining a good posterior estimate
• Higher gamma: Broader proposal — more exploration, more robust to poorly-initialized posteriors
• gamma=0.5: Default, works well in most cases

Note: gamma is not the same as "amortization" in the NPE literature. It sets the de-weighting of the conditional flow relative to the marginal during importance sampling.

Embedding Networks

Flow requires an embedding network to process observations. The embedding maps high-dimensional observations to a lower-dimensional summary statistic.

Basic Embedding

embedding:
  _target_: model.MyEmbedding
  _input_: [x]

Multi-Input Embedding

embedding:
  _target_: model.MyEmbedding
  _input_: [x, y]  # Multiple observation nodes

Nested Embedding Pipeline

embedding:
  _target_: model.Concatenate
  _input_:
    - _target_: timm.create_model
      model_name: resnet18
      pretrained: true
      num_classes: 0
      _input_:
        _target_: model.Unsqueeze
        _input_: [image]
    - _target_: torch.nn.Linear
      in_features: 64
      out_features: 32
      _input_: [metadata]

Complete Example

graph:
  z:
    evidence: [x]

    simulator:
      _target_: falcon.priors.Product
      priors:
        - ['uniform', -100.0, 100.0]
        - ['uniform', -100.0, 100.0]
        - ['uniform', -100.0, 100.0]

    estimator:
      _target_: falcon.estimators.Flow
      max_epochs: 100
      batch_size: 128
      early_stop_patience: 16
      cache_sync_every: 0
      max_cache_samples: 0
      net_type: zuko_nice
      theta_norm: true
      norm_momentum: 0.01
      embedding:
        _target_: model.E
        _input_: [x]
      lr: 0.01
      lr_decay_factor: 0.1
      lr_patience: 8
      gamma: 0.5
      discard_samples: true
      log_ratio_threshold: -20

    ray:
      num_gpus: 0

  x:
    parents: [z]
    simulator:
      _target_: model.Simulate
    observed: "./data/obs.npz['x']"

Training Strategies

Standard Training

Default configuration with continuous simulation:

buffer:
  min_samples: 4096
  max_samples: 32768
  simulate_count: 128
  simulate_when_full: true
  simulate_interval: 10

Amortized Training

Fixed dataset without simulation (for learning across many observations):

buffer:
  min_samples: 32000
  max_samples: 32000
  simulate_count: 0       # No simulation
  simulate_when_full: false

graph:
  z:
    estimator:
      _target_: falcon.estimators.Flow
      gamma: 0.8          # Higher gamma for amortization

Round-Based Training

Large batch renewal for sequential refinement:

buffer:
  min_samples: 8000
  max_samples: 8000
  simulate_count: 8000    # Full renewal
  simulate_when_full: true
  simulate_interval: 30

graph:
  z:
    estimator:
      _target_: falcon.estimators.Flow
      discard_samples: true   # Remove poor samples

Logged Metrics

Flow logs the following metrics during training:

Metric	Description
train:loss	Training loss (negative log-likelihood)
val:loss	Validation loss
lr	Current learning rate
epoch	Training epoch
checkpoint:conditional	Epoch when conditional flow was checkpointed
checkpoint:marginal	Epoch when marginal flow was checkpointed

Tips

1. Start with defaults: The default configuration works well for most problems
2. Increase max_epochs for complex posteriors
3. Enable discard_samples if training becomes unstable with outliers
4. Use GPU (ray.num_gpus: 1) for faster training with large embeddings
5. Lower gamma for single-observation inference, higher for amortization
6. Adjust early_stop_patience based on expected convergence time
7. Set cache_on_device: true when GPU memory permits, to eliminate per-batch CPU-to-GPU transfers
8. Increase cache_sync_every (e.g. 5-10) when simulations are slow and training data changes infrequently

Class Reference

Flow

Flow(*, max_epochs=100, net_type='zuko_nice', 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, theta_norm=True, norm_momentum=0.01, adaptive_momentum=False, use_log_update=False, betas=(0.9, 0.9), lr_decay_factor=0.1, lr_patience=8, discard_samples=True, log_ratio_threshold=-20.0, sample_reference_posterior=False, use_best_models=True, num_proposals=256, proposal_mixture_beta=0.5, reference_samples=128, hypercube_bound=2.0, out_of_bounds_penalty=100.0, nan_replacement=-100.0)

Bases: StepwiseEstimator

Flow-based posterior estimation using a conditional + marginal flow pair.

Parameters:

Name	Type	Description	Default
max_epochs	int	Maximum training epochs.	100
net_type	str	Flow architecture (zuko_nice, nsf, maf, zuko_gf, ...).	'zuko_nice'
lr	float	Learning rate.	0.01
gamma	float	Proposal tempering coefficient.	0.5
embedding		Embedding config dict (with _target_ etc.) or None.	None
device	Optional[str]	Device string (e.g. "cuda:0"); 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 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
theta_norm	bool	Normalise parameter space online.	True
norm_momentum	float	EMA momentum for online normalisation.	0.01
adaptive_momentum	bool	Adaptive momentum for normalisation.	False
use_log_update	bool	Use log-space normalisation update.	False
betas	tuple	AdamW beta coefficients.	(0.9, 0.9)
lr_decay_factor	float	LR decay factor for plateau scheduler.	0.1
lr_patience	int	Plateau patience before LR decay.	8
discard_samples	bool	Discard low log-ratio training samples.	True
log_ratio_threshold	float	Log-ratio cutoff for discarding.	-20.0
sample_reference_posterior	bool	Sample reference posterior for proposals.	False
use_best_models	bool	Use best-checkpoint networks for sampling.	True
num_proposals	int	Importance sampling proposal count.	256
proposal_mixture_beta	float	Fraction of proposals drawn from the conditional flow in the multiple-importance-sampling mixture (balance heuristic); the rest are drawn from the marginal flow. 1.0 = conditional only (reproduces the single-proposal behaviour), 0.0 = marginal only, 0.5 = even defensive mixture.	0.5
reference_samples	int	Reference posterior sample count.	128
hypercube_bound	float	Hypercube clipping bound for proposals.	2.0
out_of_bounds_penalty	float	Log-weight penalty for out-of-bounds samples.	100.0
nan_replacement	float	Replacement for NaN/−∞ log-weights.	-100.0

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

def __init__(
    self,
    *,
    # Most commonly changed
    max_epochs: int = 100,
    net_type: str = "zuko_nice",
    lr: float = 1e-2,
    gamma: float = 0.5,
    embedding=None,
    device: Optional[str] = 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
    theta_norm: bool = True,
    norm_momentum: float = 1e-2,
    adaptive_momentum: bool = False,
    use_log_update: bool = False,
    # Optimizer
    betas: tuple = (0.9, 0.9),
    lr_decay_factor: float = 0.1,
    lr_patience: int = 8,
    # Inference
    discard_samples: bool = True,
    log_ratio_threshold: float = -20.0,
    sample_reference_posterior: bool = False,
    use_best_models: bool = True,
    num_proposals: int = 256,
    proposal_mixture_beta: float = 0.5,
    reference_samples: int = 128,
    hypercube_bound: float = 2.0,
    out_of_bounds_penalty: float = 100.0,
    nan_replacement: float = -100.0,
):
    self.max_epochs = max_epochs
    self.net_type = net_type
    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.theta_norm = theta_norm
    self.norm_momentum = norm_momentum
    self.adaptive_momentum = adaptive_momentum
    self.use_log_update = use_log_update
    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
    self.sample_reference_posterior = sample_reference_posterior
    self.use_best_models = use_best_models
    self.num_proposals = num_proposals
    self.proposal_mixture_beta = proposal_mixture_beta
    self.reference_samples = reference_samples
    self.hypercube_bound = hypercube_bound
    self.out_of_bounds_penalty = out_of_bounds_penalty
    self.nan_replacement = nan_replacement

train_step

train_step(batch)

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

def train_step(self, batch) -> Dict[str, float]:
    ids, theta, theta_logprob, conditions, u, u_device, conditions_device = \
        self._unpack_batch(batch, "train")

    if not self.networks_initialized:
        self._initialize_networks(u, conditions)

    s = self._embed(conditions_device, train=True)

    with torch.no_grad():
        self.history["theta_mins"].append(theta.min(dim=0).values.cpu().numpy())
        self.history["theta_maxs"].append(theta.max(dim=0).values.cpu().numpy())

    self._optimizer.zero_grad()
    loss_cond, loss_marg = self._compute_flow_losses(u_device, s, train=True)
    (loss_cond + loss_marg).backward()
    self._optimizer.step()

    if self.discard_samples:
        discard_mask = self._compute_discard_mask(theta, theta_logprob, conditions_device)
        batch.discard(discard_mask)

    return {"loss": loss_cond.item(), "loss_aux": loss_marg.item()}

val_step

val_step(batch)

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

def val_step(self, batch) -> Dict[str, float]:
    _, theta, theta_logprob, conditions, u, u_device, conditions_device = \
        self._unpack_batch(batch, "val")

    s = self._embed(conditions_device, train=False)

    with torch.no_grad():
        loss_cond, loss_marg = self._compute_flow_losses(u_device, s, train=False)

    return {"loss": loss_cond.item(), "loss_aux": loss_marg.item()}

sample_prior

sample_prior(num_samples, conditions=None)

Source code in src/falcon/estimators/flow.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)
    log_prob = np.ones(num_samples) * (-np.log(2 * self.hypercube_bound) ** self.param_dim)
    return {'value': samples, 'log_prob': log_prob}

sample_posterior

sample_posterior(num_samples, conditions=None)

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

def sample_posterior(self, num_samples: int, conditions=None) -> dict:
    if not self.networks_initialized:
        return self.sample_prior(num_samples)
    samples, logprob = self._importance_sample(num_samples, mode="posterior", conditions=conditions or {})
    return {'value': samples.numpy(), 'log_prob': logprob.numpy()}

sample_proposal

sample_proposal(num_samples, conditions=None)

Source code in src/falcon/estimators/flow.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)
    if not self.networks_initialized:
        return self.sample_prior(num_samples)

    conditions = conditions or {}
    if self.sample_reference_posterior:
        post_samples, _ = self._importance_sample(
            self.reference_samples, mode="posterior", conditions=conditions
        )
        mean, std = post_samples.mean(dim=0).cpu(), post_samples.std(dim=0).cpu()
        log({f"sample_proposal:posterior_mean_{i}": mean[i].item() for i in range(len(mean))})
        log({f"sample_proposal:posterior_std_{i}": std[i].item() for i in range(len(std))})

    samples, logprob = self._importance_sample(num_samples, mode="proposal", conditions=conditions)
    log({
        "sample_proposal:mean": samples.mean().item(),
        "sample_proposal:std": samples.std().item(),
        "sample_proposal:logprob": logprob.mean().item(),
    })
    return {'value': samples.numpy(), 'log_prob': logprob.numpy()}

save

save(node_dir)

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

def save(self, node_dir: Path) -> None:
    debug(f"Saving: {node_dir}")
    if not self.networks_initialized:
        raise RuntimeError("Networks not initialized.")

    torch.save(self._best_conditional_flow.state_dict(), node_dir / "conditional_flow.pth")
    torch.save(self._best_marginal_flow.state_dict(), node_dir / "marginal_flow.pth")
    torch.save(self._init_parameters, node_dir / "init_parameters.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["theta_mins"], node_dir / "theta_mins_batches.pth")
    torch.save(self.history["theta_maxs"], node_dir / "theta_maxs_batches.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")

    if self._best_embedding is not None:
        torch.save(self._best_embedding.state_dict(), node_dir / "embedding.pth")

load

load(node_dir)

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

def load(self, node_dir: Path) -> None:
    debug(f"Loading: {node_dir}")
    init_parameters = torch.load(node_dir / "init_parameters.pth")
    self._initialize_networks(init_parameters[0], init_parameters[1])

    self._best_conditional_flow.load_state_dict(torch.load(node_dir / "conditional_flow.pth"))
    self._best_marginal_flow.load_state_dict(torch.load(node_dir / "marginal_flow.pth"))

    if (node_dir / "embedding.pth").exists() and self._best_embedding is not None:
        self._best_embedding.load_state_dict(torch.load(node_dir / "embedding.pth"))

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

Configuration Classes

TrainingLoopConfig dataclass

TrainingLoopConfig(max_epochs=100, batch_size=128, early_stop_patience=16, cache_sync_every=0, max_cache_samples=0, cache_on_device=False, prior_epochs=0)

Generic training loop parameters.