Population Level Parameter Recovery Problem

General
1

Dear BayesFlow Team,

thanks for developing such an interesting tool, and I really enjoy working with it!

I was trying to train a network for a hierarchical two-choice drift diffusion model. It has 5 parameters in the individual level (two drift rates (b0 & b1), boundary separation (\alpha), response bias (\beta), and non-decision time (\tau)), and 10 parameters in the population level (the mean \mu and standard deviation \sigma of the five distributions where the individual parameters were drawn from).

The network settings are:

summary_net = bf.summary_networks.HierarchicalNetwork([
    #from the bottom level to the top level
    bf.networks.DeepSet(summary_dim=50), 
    bf.networks.DeepSet(summary_dim=200)
])

local_inference_net = bf.networks.InvertibleNetwork(
    num_params=5,
    coupling_settings={"dense_args": dict(kernel_regularizer=None), "dropout": False},
    name="local_ddm_inference"
)

hyper_inference_net = bf.networks.InvertibleNetwork(
    num_params=10,
    coupling_settings={"dense_args": dict(kernel_regularizer=None), "dropout": False},
    name="hyper_ddm_inference"
)

The training was fine, and the individual level parameters are nicely recovered after this much of training:

history = trainer.train_online(epochs=300, iterations_per_epoch=300, batch_size=10)

individual level parameters recovery
individual level parameters recovery961×727 94.4 KB

However, the population level parameter recovery is problematic for one parameter (in this case \sigma_{b0}):

population level parameters recovery
population level parameters recovery945×1141 166 KB

From my experience of retraining the model several times, the problematic recovery only appears in the estimation of population-level standard deviation, but not on a certain standard deviation. For example, It can happen on \sigma_{b0}, \sigma_{b1} or \sigma_\tau, and the plot here is only the output of one neural network model.

I am quite confused about this problem and I am seeking for help on this. Please let me know if anything is unclear here.

Thanks a lot!

Best,
Yufei

2

Some additional information is:

I trained these networks on GPU, and when I tried to investigate the latent space z with:

f = bf.diagnostics.plot_latent_space_2d(z_samples[1])

the z samples (of the population level I guess) has 32 dimensions.
But from the tutorial I read that the dimensionality of the latent space is equal to the num_params in the inference network. So that was a bit confusing as well.

Thanks!

3

I would start by dropping the DeepSet in favor of a SetTransformer summary network.