Difficulty Estimating a Multivariate Normal Outcome

General
5

Thanks for sharing. It seems that the summary networks (both SetTransformer and DeepSet) fail to learn the relevant summary statistics required to estimate the variances. I further reduced the code to a minimal example that I would have expected to work, but it doesn’t.

import keras
import bayesflow as bf
import numpy as np

D = 12
N = 10
rng = np.random.default_rng(2025)

def prior():
    variance = rng.uniform(size=D)
    return {"variance": variance}


def likelihood(variance):
    y = rng.normal(loc=np.zeros_like(variance), scale=np.sqrt(variance), size=(N, D))
    return {"y": y}

simulator = bf.make_simulator([prior, likelihood])

validation_data = simulator.sample(100)
print("shapes", keras.tree.map_structure(keras.ops.shape, validation_data))

summary_network = bf.networks.SetTransformer(summary_dim=2*D)

workflow = bf.BasicWorkflow(
    simulator=simulator,
    summary_network=summary_network,
    summary_variables="y",
    inference_variables=["variance"],
    initial_learning_rate=1e-3,
    standardize="all",
)

workflow.fit_online(num_batches_per_epoch=1000, epochs=10, validation_data=validation_data, batch_size=32);

test_data = simulator.sample(200)

workflow.plot_default_diagnostics(test_data);

This produces the following result. While it learns the relevant summary statistics for some dimensions, it fails to do so for the others. Note that they are all interchangeable, so which one will be learned is random.

image
image897×296 98.1 KB

Using a SetTransformer instead of the DeepSet shows the same problem, but to a somewhat lesser degree:

image
image1495×494 245 KB

I’m not sure if I missed anything, tagging @KLDivergence, @paul.buerkner, @elseml if they have any ideas or insights to share on how one would approach this.

1 Like