bayesflow.experimental.rectifiers module

class bayesflow.experimental.rectifiers.DriftNetwork(*args, **kwargs)

Bases: Model

Implements a learnable velocity field for a neural ODE. Will typically be used in conjunction with a RectifyingFlow instance, as proposed by [1] in the context of unconditional image generation.

[1] Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.

__init__(target_dim, num_dense=3, dense_args=None, dropout=True, mc_dropout=False, dropout_prob=0.05, **kwargs)

Creates a learnable velocity field instance to be used in the context of rectifying flows or neural ODEs.

[1] Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.

Parameters:

target_dimint

The problem dimensionality (e.g., in parameter estimation, the number of parameters)

num_denseint, optional, default: 3

The number of hidden layers for the inner fully-connected network

dense_argsdict or None, optional, default: None

The arguments to be passed to tf.keras.layers.Dense constructor. If None, default settings will be fetched from bayesflow.default_settings.

dropoutbool, optional, default: True

Whether to use dropout in-between the hidden layers.

mc_dropoutbool, optional, default: False

Whether to use dropout Monte Carlo dropout (i.e., Bayesian approximation) during inference

dropout_probfloat in (0, 1), optional, default: 0.05

The dropout probability. Only has effecft if dropout=True or mc_dropout=True

**kwargsdict, optional, default: {}

Optional keyword arguments passed to the tf.keras.Model.__init__ method.

call(target_vars, latent_vars, time, condition, **kwargs)

Performs a linear interpolation between target and latent variables over time (i.e., a single ODE step during training).

Parameters:

target_varstf.Tensor of shape (batch_size, …, num_targets)

The variables of interest (e.g., parameters) over which we perform inference.

latent_varstf.Tensor of shape (batch_size, …, num_targets)

The sampled random variates from the base distribution.

timetf.Tensor of shape (batch_size, …, 1)

A vector of time indices in (0, 1)

conditiontf.Tensor of shape (batch_size, …, condition_dim)

The optional conditioning variables (e.g., as returned by a summary network)

**kwargsdict, optional, default: {}

Optional keyword arguments passed to the tf.keras.Model call() method

drift(target_t, time, condition, **kwargs)

Returns the drift at target_t time given optional condition(s).

Parameters:

target_ttf.Tensor of shape (batch_size, …, num_targets)

The variables of interest (e.g., parameters) over which we perform inference.

timetf.Tensor of shape (batch_size, …, 1)

A vector of time indices in (0, 1)

conditiontf.Tensor of shape (batch_size, …, condition_dim)

The optional conditioning variables (e.g., as returned by a summary network)

**kwargsdict, optional, default: {}

Optional keyword arguments passed to the drift network.

class bayesflow.experimental.rectifiers.RectifiedDistribution(*args, **kwargs)

Bases: Model

Implements a rectifying flows according to [1]. To be used as an alternative to a normalizing flow in a BayesFlow pipeline.

[1] Liu, X., Gong, C., & Liu, Q. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.

__init__(drift_net, summary_net=None, latent_dist=None, loss_fun=None, summary_loss_fun=None, **kwargs)

Initializes a composite neural network to represent an amortized approximate posterior through for a rectifying flow.

Parameters:

drift_nettf.keras.Model

A neural network for the velocity field (drift) of the learnable ODE

summary_nettf.keras.Model or None, optional, default: None

An optional summary network to compress non-vector data structures.

latent_distcallable or None, optional, default: None

The latent distribution towards which to optimize the networks. Defaults to a multivariate unit Gaussian.

loss_funcallable or None, optional, default: None

The loss function for “rectifying” the velocity field. If None, defaults to tf.keras.losses.logcosh. Sensible alternatives are MSE (as in [])

summary_loss_funcallable, str, or None, optional, default: None

The loss function which accepts the outputs of the summary network. If None, no loss is provided and the summary space will not be shaped according to a known distribution (see [2]). If summary_loss_fun='MMD', the default loss from [2] will be used.

**kwargsdict, optional, default: {}

Additional keyword arguments passed to the __init__ method of a tf.keras.Model instance.

call(input_dict, return_summary=False, num_eval_points=1, **kwargs)

Performs a forward pass through the summary and drift network given an input dictionary.

Parameters:

input_dictdict

Input dictionary containing the following mandatory keys, if DEFAULT_KEYS unchanged: targets - the latent model parameters over which a condition density is learned summary_conditions - the conditioning variables (including data) that are first passed through a summary network direct_conditions - the conditioning variables that the directly passed to the inference network

return_summarybool, optional, default: False

A flag which determines whether the learnable data summaries (representations) are returned or not.

num_eval_pointsint, optional, default: 1

The number of time points for evaluating the noisy estimator. Values larger than the default 1 may reduce the variance of the estimator, but may lead to increased memory demands, since an additional dimension is added at axis 1 of all tensors.

**kwargsdict, optional, default: {}

Additional keyword arguments passed to the networks For instance, kwargs={'training': True} is passed automatically during training.

Returns:

net_out or (net_out, summary_out)

compute_loss(input_dict, **kwargs)

Computes the loss of the posterior amortizer given an input dictionary, which will typically be the output of a Bayesian GenerativeModel instance.

Parameters:

input_dictdict

Input dictionary containing the following mandatory keys, if DEFAULT_KEYS unchanged: targets - the latent variables over which a condition density is learned summary_conditions - the conditioning variables that are first passed through a summary network direct_conditions - the conditioning variables that the directly passed to the inference network

**kwargsdict, optional, default: {}

Additional keyword arguments passed to the networks For instance, kwargs={'training': True} is passed automatically during training.

Returns:

total_losstf.Tensor of shape (1,) - the total computed loss given input variables

sample(input_dict, n_samples, to_numpy=True, step_size=0.001, **kwargs)

Generates random draws from the approximate posterior given a dictionary with conditonal variables.

Parameters:

input_dictdict

Input dictionary containing the following mandatory keys, if DEFAULT_KEYS unchanged: summary_conditions : the conditioning variables (including data) that are first passed through a summary network direct_conditions : the conditioning variables that the directly passed to the inference network

n_samplesint

The number of posterior draws (samples) to obtain from the approximate posterior

to_numpybool, optional, default: True

Flag indicating whether to return the samples as a np.ndarray or a tf.Tensor

step_sizefloat, optional, default: 0.01

The step size for the stochastic Euler solver.

**kwargsdict, optional, default: {}

Additional keyword arguments passed to the networks

Returns:

post_samplestf.Tensor or np.ndarray of shape (n_data_sets, n_samples, n_params)

The sampled parameters from the approximate posterior of each data set

log_density(input_dict, to_numpy=True, step_size=0.001, **kwargs)

Computes the log density…