Neural Likelihoods¶

In data fitting, a ParameterEstimationComposition computes the likelihood of its data under each set of parameter values it considers by simulating the model num_estimates times and estimating a density from the simulated outcomes. A neural likelihood is a density estimator that is trained beforehand, on data simulated from the model, and then used in place of those simulations. Each evaluation of the likelihood then requires a single evaluation of the estimator, and the likelihood varies smoothly with the parameter values.

Training an Estimator¶

train_neural_likelihood simulates the model at parameter values drawn uniformly from within the ranges specified in bounds, and trains an estimator on the results. The model can be specified directly, together with the inputs used to run it:

likelihood = pnl.train_neural_likelihood(
    bounds={"rate": (-1.5, 1.5), "threshold": (0.3, 1.5)},
    outcome_names=("decision", "response_time"),
    pec=pec,
    inputs={comp: trial_inputs},
)
likelihood.save("ddm_nle.pt")

The keys of bounds must name the fitted parameters in the order in which the model lists them, and outcome_names must be the column names of the data that the estimator will be used to fit, which are in the order of the model’s outcome_variables. Each set of parameter values is simulated on the trials in inputs, each as many times as the model’s num_estimates.

To distribute the simulations over a Dask cluster, specify a pec_factory in place of pec, together with distributed_options (see Distributed Fitting):

likelihood = pnl.train_neural_likelihood(
    bounds={"rate": (-1.5, 1.5), "threshold": (0.3, 1.5)},
    outcome_names=("decision", "response_time"),
    pec_factory=build_pec,
    n_trials_per_sample=500,
    distributed_options={"n_workers": 8},
)

build_pec(data) -> (pec, inputs) is a top-level function that builds the model, as for distributed fitting. Each worker calls it with a table of n_trials_per_sample rows to build its own copy of the model.

train_neural_likelihood generates an error if the estimator’s negative log-likelihood on the data held out from training is not finite, and records it per trial as metadata.val_nll. This shows that training worked, not that the estimator is accurate; for that, fit data simulated at known parameter values.

Fitting with an Estimator¶

To use a trained estimator, specify likelihood_estimator="neural", and the estimator, or the path to one saved with its save method, as the "artifact" of likelihood_estimator_kwargs:

pec = pnl.ParameterEstimationComposition(
    nodes=[comp],
    parameters={("rate", decision): np.linspace(-1.5, 1.5, 1000),
                ("threshold", decision): np.linspace(0.3, 1.5, 1000)},
    outcome_variables=[decision.output_ports[pnl.DECISION_OUTCOME],
                       decision.output_ports[pnl.RESPONSE_TIME]],
    data=data,
    optimization_function="differential_evolution",
    likelihood_estimator="neural",
    likelihood_estimator_kwargs={"artifact": "ddm_nle.pt"},
)
pec.run(inputs={comp: trial_inputs})

The composition being fitted is not simulated when fitting with an estimator, so it is not compiled. log_likelihood also uses the estimator.

A parameter that depends on a condition (specified in depends_on) is fit separately for each condition, and each trial is scored with the value for its condition. The estimator for such a fit is trained on the model without depends_on, over the range of each parameter.

In a hierarchical fit, likelihood_estimator is specified for the participant models built by the pec_factory, and not for the group.

Trial Features¶

Where trials differ from one another (for example, congruent and incongruent trials), an estimator represents the distribution of outcomes on each kind of trial, which it distinguishes by the values of the model’s inputs on each trial.

When fitting, the same inputs are taken from those specified for run or log_likelihood, including any that do not vary in the data being fit. These must be laid out as they were for training, though the nodes to which they are given can be listed in any order, and an input that did not vary in training must have the same value; otherwise an error is generated, since the estimator has no information about such trials.

Matching an Estimator to a Model¶

An estimator is valid only for the model, and the ranges of parameter values, for which it was trained; within those ranges, it can be reused for any fit of that model. It records the following, and generates an error if it is used to fit a model that differs in any of them:

the fitted parameters, or their order;
the range of any parameter, which must lie within the range used for training;
the outcome variables, their order, or which of them are categorical;
the values of a categorical outcome, which must be among those simulated in training.

Other properties of the model, such as the values of parameters that are not fit, are not recorded; an estimator should be retrained if any of these are changed.

Limitations¶

At least one outcome must be continuous, such as a response time: an estimator cannot be trained for a model whose outcomes are all categorical.

Requirements¶

Neural likelihoods require the nle extra, installed with pip install "psyneulink[nle]", which includes sbi and PyTorch.

train_neural_likelihood.py trains an estimator for a drift-diffusion model, in one process or, with --distributed, over a Dask cluster, and uses it to fit simulated data.

Class Reference¶

psyneulink.core.components.functions.nonstateful.neurallikelihoodfunctions.train_neural_likelihood(bounds, outcome_names, *, pec=None, inputs=None, pec_factory=None, n_parameter_samples=20000, n_trials_per_sample=None, categorical=None, epochs=30, batch_size=512, learning_rate=0.0005, validation_fraction=0.1, seed=0, distributed_options=None)¶

Train a NeuralLikelihood on data simulated from a composition.

See Neural Likelihoods.

Parameters:

bounds (Mapping) – specifies the range (lower, upper) of each fitted parameter, by name, in the order the model lists them. The estimator is valid only within these ranges.
outcome_names (Sequence[str]) – specifies the names of the outcome variables, which must be the column names of the data the estimator will be used to fit, in the order of the model’s outcome_variables.
pec (ParameterEstimationComposition : default None) – specifies a model to simulate in this process. Requires inputs.
inputs (Mapping : default None) – specifies the inputs with which pec is run; each parameter draw simulates the trials in inputs, each as many times as the model’s num_estimates.
pec_factory (callable : default None) – specifies a function pec_factory(data) -> (pec, inputs) that builds the model, scored by simulating it, as used for distributed fitting; it is called with a table of n_trials_per_sample rows. Required by distributed_options. Exactly one of pec and pec_factory must be specified.
n_parameter_samples (int : default 20000) – specifies the number of parameter values drawn uniformly from within bounds and simulated.
n_trials_per_sample (int : default None) – specifies the number of trials in the table pec_factory is called with, which sets the trials each parameter draw simulates as inputs does for pec; 100 if not specified.
categorical (Sequence[bool] : default None) – specifies which outcome variables are categorical; if not specified, this is inferred from the simulated data.
epochs (int : default 30) – specifies the number of passes over the simulated data made in training.
batch_size (int : default 512) – specifies the number of rows in each training step.
learning_rate (float : default 5e-4) – specifies the learning rate of the Adam optimizer used for training.
validation_fraction (float : default 0.1) – specifies the fraction of the parameter draws whose simulated data are held out from training, to evaluate the estimator.
seed (int : default 0) – specifies the seed for the parameter draws, the noise simulated at each of them, and training.
distributed_options (Mapping : default None) – specifies a Dask cluster over which to distribute the simulations, as for distributed fitting. Requires pec_factory. Each worker builds its own model before simulating, so this is worthwhile only for more than a few hundred parameter draws.

Return type:

A trained NeuralLikelihood.

class psyneulink.core.components.functions.nonstateful.neurallikelihoodfunctions.NeuralLikelihood(estimator, metadata, shape_probe)¶

A trained conditional density p(outcomes | parameters, trial features).

Built by train_neural_likelihood, saved with save, and reloaded with load. trial_log_prob is differentiable with respect to the parameters.

metadata¶

what the estimator was trained for, which is checked against a model before it is used to fit one (see Matching an Estimator to a Model), and its negative log-likelihood per trial on the data held out from training (val_nll).

Type:: NeuralLikelihoodMetadata

trial_log_prob(theta, outcomes, trial_features=None)¶

Per-trial log densities, differentiable with respect to theta.

theta is one vector of parameters for every trial, or one row of them per trial.

Return type:: Tensor

log_likelihood(theta, outcomes, trial_features=None)¶

Total log-likelihood of outcomes under theta.

Return type:: float

save(path)¶: Write weights and metadata to path.

classmethod load(path)¶

Read back an estimator written by save.

Return type:: NeuralLikelihood

exception psyneulink.core.components.functions.nonstateful.neurallikelihoodfunctions.NeuralLikelihoodError¶: Raised when a neural likelihood is misconfigured or does not match its model.