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
• Round-based training with automatic learning rate scheduling and early stopping (see Training Loop)

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. Training runs in rounds of epochs on fixed data; see Training Loop for the terminology and the full description.

Parameter	Type	Default	Description
max_rounds	int	null	Maximum number of rounds (null = unlimited)
patience_rounds	int	10	Stop training after this many rejected rounds in a row
max_epochs	int	100	Maximum epochs per round
patience_epochs	int	16	End the round once the best validation loss is this many epochs old
val_every_epochs	int	1	Validate after every N-th epoch (and at the last epoch of a round)
batch_size	int	128	Training batch size
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_rounds	int	0	Rounds that simulate from the prior before switching to the learned proposal (the first round always does)
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 refreshed from the shared simulation buffer at the start of every round and then stays fixed for the round. This avoids repeated remote data fetches and allows fast batching.

• 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	Learning rate at the start of every round
lr_decay_factor	float	1.0	LR multiplier when plateau detected (1.0 = no decay)
lr_patience_epochs	int	8	Epochs without validation 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	After each accepted round, discard low-likelihood samples
log_ratio_threshold	float	-20	Log-likelihood threshold for sample discarding
sample_reference_posterior	bool	false	Sample from reference posterior
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
      patience_epochs: 16
      patience_rounds: 10
      batch_size: 128
      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: 1.0
      lr_patience_epochs: 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 of the conditional flow (negative log-likelihood)
train:loss_aux	Training loss of the marginal flow
val:loss, val:loss_aux	Validation losses (only for validated epochs)
lr	Current learning rate
round, epoch	Current round, and epoch within the round
checkpoint:conditional	Epoch at which the conditional flow reached its best validation loss in the round
checkpoint:marginal	Epoch at which the marginal flow reached its best validation loss in the round
round:accepted	1 if the round was accepted
round:conditional:promoted, round:marginal:promoted	1 if the flow replaced its best network

See Training Loop for all round metrics.

Tips

1. Start with defaults: The default configuration works well for most problems
2. Increase max_epochs (per round) 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 patience_epochs based on expected convergence time within a round
7. Set cache_on_device: true when GPU memory permits, to eliminate per-batch CPU-to-GPU transfers
8. Increase val_every_epochs (e.g. 3) when validation is expensive; patience_epochs keeps its meaning

Class Reference

Flow

Flow(*, max_epochs=100, net_type='zuko_nice', 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, theta_norm=True, norm_momentum=0.01, adaptive_momentum=False, use_log_update=False, betas=(0.9, 0.9), lr_decay_factor=1.0, lr_patience_epochs=8, discard_samples=True, log_ratio_threshold=-20.0, sample_reference_posterior=False, 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.

Training runs in rounds of epochs on fixed data; see StepwiseEstimator. The conditional flow (with the embedding) and the marginal flow are promoted to best networks independently; a round is accepted when the conditional flow improves.

Parameters:

Name	Type	Description	Default
max_epochs	int	Maximum epochs per round.	100
net_type	str	Flow architecture (zuko_nice, nsf, maf, zuko_gf, ...).	'zuko_nice'
lr	float	Learning rate; reset at the start of every round.	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
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 device.	False
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 (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.	True
log_ratio_threshold	float	Log-ratio cutoff for discarding.	-20.0
sample_reference_posterior	bool	Sample reference posterior for proposals.	False
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,
    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
    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 = 1.0,
    lr_patience_epochs: int = 8,
    # Inference
    discard_samples: bool = True,
    log_ratio_threshold: float = -20.0,
    sample_reference_posterior: bool = False,
    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.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.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_epochs = lr_patience_epochs
    self.discard_samples = discard_samples
    self.log_ratio_threshold = log_ratio_threshold
    self.sample_reference_posterior = sample_reference_posterior
    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._summary(batch, "conditional", 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()

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

val_step

val_step(batch, use_best=False)

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

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

    s = self._summary(batch, "conditional", conditions_device, use_best=use_best)
    loss_cond, loss_marg = self._compute_flow_losses(u_device, s, train=False, use_best=use_best)

    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._has_best:
        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._use_prior_proposal():
        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")
    self._save_round_state(node_dir)

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"))

    # A resumed run starts from the best networks
    for group in self._network_groups().values():
        self._copy_modules(group.best_modules(), group.current_modules())

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

Training Loop

Flow inherits its round loop from StepwiseEstimator. The loop parameters (max_rounds, patience_rounds, max_epochs, patience_epochs, val_every_epochs, batch_size, max_cache_samples, cache_on_device, prior_rounds) are Flow.__init__ arguments, documented above.

StepwiseEstimator

Bases: BaseEstimator

Estimator with a round-based training loop.

Training runs in rounds (see docs/training.md for the terminology):

1. Refresh the training and validation caches from the buffer; the data then stays fixed for the whole round.
2. Start from the best networks and train for up to max_epochs epochs (passes over the training set), validating every val_every_epochs epochs and at the last epoch. The round ends early once the best validation loss is patience_epochs epochs old.
3. Restore each network group to its best validated epoch of the round and compare it with the best network on the round's validation set. Groups that improve are promoted; the round is accepted if the first (primary) group was promoted.
4. After an accepted round, run one discard sweep over the training and validation sets with the new best networks.

Training stops after patience_rounds rejected rounds in a row, after max_rounds rounds, or when interrupted. Proposal sampling can run between any two training steps; it uses the best networks, which only change when a round is accepted.

Subclasses must implement: - train_step() / val_step() / on_validation_end() / discard_mask() - _network_groups() - sample_prior/posterior/proposal - save/load

The loop reads its parameters off self -- max_rounds, patience_rounds, max_epochs, patience_epochs, val_every_epochs, batch_size, max_cache_samples, cache_on_device, discard_samples -- which each subclass sets from its own init kwargs.

train async

train(buffer)

Main training loop: rounds of epochs on fixed data.