Linear regression example does not produce accurate posteriors

[ErikRingen]

I've been following the development of 2.0 as I find the changes to the interface very appealing. But I have been having a hard time getting accurate posterior inference for even very simple models, like linear regression example notebook. I have tried comparing the NPE with HMC, and doing SBC, and I always find that the resulting posteriors are not calibrated. The closest I have gotten to achieving calibration is with a CouplingFlow inference network after 30+ epochs, but it is still not that accurate.

I am surprised, because when using the main branch version of BayesFlow I can get valid NPE for such models quite rapidly using bf.networks.InvertibleNetwork().

[paul-buerkner]

I believe this has to do with our current hyperparamter settings in the dev branch which are not yet well set by default. before merging into main we will make sure that this is fixed.

Stefan, what are your thoughts on this?

Erik Ringen @.***> schrieb am So., 27. Okt. 2024, 11:03:

I've been following the development of 2.0 as I find the changes to the interface very appealing. But I have been having a hard time getting accurate posterior inference for even very simple models, like linear regression example notebook. I have tried comparing the NPE with HMC, and doing SBC, and I always find that the resulting posteriors are not calibrated. The closest I have gotten to achieving calibration is with a CouplingFlow inference network after 30+ epochs, but it is still not that accurate.

I am surprised, because when using the main branch version of BayesFlow I can get valid NPE for such models quite rapidly using bf.networks.InvertibleNetwork().

— Reply to this email directly, view it on GitHub https://github.com/bayesflow-org/bayesflow/issues/227, or unsubscribe https://github.com/notifications/unsubscribe-auth/ADCW2AESYSNPIFLRCDEXZV3Z5SM5RAVCNFSM6AAAAABQVT3LAWVHI2DSMVQWIX3LMV43ASLTON2WKOZSGYYTMNBTGYZTINA . You are receiving this because you are subscribed to this thread.Message ID: @.***>

[stefanradev93]

Hey Erik, indeed, the default hyperparameter settings are not the same between branches, but those for the CouplingFlow are pretty close. For instance, flow matching will be pretty shallow when using the defaults, so you may find it needs a deeper internal network (e.g., as in the Two Moons notebook). Our plan is to have a large-scale hyperparameter optimization study on multiple models before merging. Would it be possible to share the model and settings you are currently working with (if they go beyond the linear regression case)?

[ErikRingen]

Hey, thanks for your responses. I don't know if it is useful, but here is an example with a simple poisson polynomial regression model and settings that work pretty well--but these same settings do not work well for the linear regression notebook, for some reason.

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import os
if "KERAS_BACKEND" not in os.environ:
    os.environ["KERAS_BACKEND"] = "jax"
import keras
import bayesflow as bf
from scipy.stats import rankdata
from statsmodels.stats.proportion import proportion_confint
from scipy.stats import halfnorm

def meta(batch_shape):
    N = int(10)
    return dict(N=N)

def prior():
    beta = np.random.normal([-1, 0, 0], [1, 0.25, 0.25])
    return dict(beta=beta)

def likelihood(beta, N):
    # x: predictor variable
    x = np.random.normal(0, 1, size=N)
    # y: response variable
    mu = np.exp(beta[0] + beta[1] * x + beta[2] * (x**2))
    y = np.random.poisson(mu, size=N)
    return dict(y=y, x=x)

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

adapter = bf.adapters.Adapter()
adapter.to_array()
adapter.broadcast("N")
adapter.convert_dtype(from_dtype="float64", to_dtype="float32")
adapter.as_set(["x", "y"])
adapter.standardize(exclude=["N"])
adapter.concatenate(["beta"], into="inference_variables")
adapter.concatenate(["x", "y"], into="summary_variables")
adapter.rename("N", "inference_conditions")

epochs = 30
num_batches = 512
total_steps = num_batches * epochs

inference_network = bf.networks.CouplingFlow()

summary_network = bf.networks.DeepSet(summary_dim = 6, subnet_kwargs={"base_distribution": "normal"})

approximator = bf.ContinuousApproximator(
   inference_network=inference_network,
   summary_network=summary_network,
   adapter=adapter,
)

initial_learning_rate = 5e-4
scheduled_lr = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=initial_learning_rate,
    decay_steps=total_steps,
    alpha=1e-8
)

optimizer = keras.optimizers.AdamW(learning_rate=scheduled_lr)
approximator.compile(optimizer=optimizer)

history = approximator.fit(
    epochs=epochs,
    num_batches=num_batches,
    batch_size=64,
    simulator=simulator,
)

# Check calibration of NPE
def rank_SBC(nsims, num_samples):
    rank = np.empty(nsims)

    for i in range(nsims):
        prior_sims = prior()
        N = 10
        data_sims = likelihood(beta = prior_sims["beta"], N = N)
        conditions = {
                "y": data_sims["y"],
                "x": data_sims["x"],
                "N": N
            }
        pdraws = approximator.sample(conditions=conditions, batch_size=1, num_samples = num_samples)
        ranks = rankdata(np.append(pdraws["beta"][:,:,1], prior_sims["beta"][1]))
        rank[i] = ranks[-1]
        
    return rank

def ecdf(data, num_samples):
    sorted_data = np.sort(data)
    return sorted_data/(num_samples+1)

def compute_ecdf_diff_confidence_uniform_sampling(n, alpha=0.05):
    x_values = np.linspace(0, 1, n)
    uniform_cdf = x_values  # For uniform data, CDF(x) = x
    lower_bounds = []
    upper_bounds = []

    for i in range(n):
        successes = i + 1
        # Compute Wilson confidence interval for the ECDF at this point
        ci_low, ci_upp = proportion_confint(successes, n, alpha=alpha, method='wilson')
        
        lower_bounds.append(ci_low - uniform_cdf[i])
        upper_bounds.append(ci_upp - uniform_cdf[i])

    return x_values, lower_bounds, upper_bounds

num_samples = 999
nsims = 500
rank_intercept = rank_SBC(nsims = nsims, num_samples=num_samples)
ecdf_diff = np.linspace(0, 1, nsims) - ecdf(rank_intercept, num_samples)
plt.hist(rank_intercept)

plt.figure()
plt.step(np.sort(rank_intercept)/(num_samples + 1), ecdf_diff, where="post")

x_vals, lower_ci, upper_ci = compute_ecdf_diff_confidence_uniform_sampling(n = nsims)
plt.fill_between(x_vals, lower_ci, upper_ci, color='gray', alpha=0.2, label='95% Confidence Interval')

[paul-buerkner]

Thank you!

[paul-buerkner]

We have now improved some of the generative architecture including affine flows and also added back the diagnostics module. Woudl you mind checking your example again? We are still playing around with the clamping so I cannot guarantee yet that the current implementation will already be as good as you and we want it to be. But we will get there.

[paul-buerkner]

I have run over linear regression example and got good results if initialization worked. However, on some training runs, I got NaN losses after a few iterations. So I believe we need to work a bit more at least on the part that handles initialization of the coupling flows.

[ErikRingen]

I have run over linear regression example and got good results if initialization worked. However, on some training runs, I got NaN losses after a few iterations. So I believe we need to work a bit more at least on the part that handles initialization of the coupling flows.

Thanks for letting me know about the updates. Both my poisson and linear regression notebook examples look better now.

I have also run into issues with NaN losses in the past when experimenting with coupling flows. I addressed this by setting TerminateOnNaN in the tensorflow callbacks.

Below is my updated polynomial poisson regression example.

import numpy as np
import os
if "KERAS_BACKEND" not in os.environ:
    os.environ["KERAS_BACKEND"] = "jax"
import keras
import bayesflow as bf

def meta(batch_shape):
    N = int(10)
    return dict(N=N)

def prior():
    beta = np.random.normal([-1, 0, 0], [1, 0.25, 0.25])
    return dict(beta=beta)

def likelihood(beta, N):
    # x: predictor variable
    x = np.random.normal(0, 1, size=N)
    # y: response variable
    mu = np.exp(beta[0] + beta[1] * x + beta[2] * (x**2))
    y = np.random.poisson(mu, size=N)
    return dict(y=y, x=x)

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

adapter = (
    bf.Adapter()
    .broadcast("N", to="x")
    .as_set(["x", "y"])
    .standardize(exclude=["N"])
    .apply(include="N", forward=lambda n: np.sqrt(n), inverse=lambda n: n**2)
    .concatenate(["beta"], into="inference_variables")
    .concatenate(["x", "y"], into="summary_variables")
    .rename("N", "inference_conditions")
)

inference_network = bf.networks.CouplingFlow()
summary_network = bf.networks.DeepSet(depth=2)

approximator = bf.ContinuousApproximator(
   inference_network=inference_network,
   summary_network=summary_network,
   adapter=adapter,
)

epochs = 50
num_batches = 128
batch_size = 64
learning_rate = keras.optimizers.schedules.CosineDecay(5e-4, decay_steps=epochs*num_batches, alpha=1e-6)
optimizer = keras.optimizers.Adam(learning_rate=learning_rate, clipnorm=1.0)
approximator.compile(optimizer=optimizer)

history = approximator.fit(
    epochs=epochs,
    num_batches=num_batches,
    batch_size=batch_size,
    simulator=simulator,
)

f = bf.diagnostics.plot_losses(
    train_losses=np.array(history.history["loss"]).astype(np.float32)
)

# Set the number of posterior draws you want to get
num_samples = 2000

# Simulate validation data
val_sims = simulator.sample(300)
conditions = {k: v for k, v in val_sims.items() if k != "beta" and k != "sigma"}

pdraws = approximator.sample(conditions=conditions, num_samples=num_samples)

f = bf.diagnostics.plot_recovery(pdraws["beta"], val_sims["beta"])

f = bf.diagnostics.plot_sbc_ecdf(pdraws["beta"], val_sims["beta"], difference=True)

f = bf.diagnostics.plot_sbc_histograms(pdraws["beta"], val_sims["beta"])

[ErikRingen]

p.s. I am excited about the potential to use bayesflow for power/precision analysis. If anyone is actively working on this I'd be happy to chat with them.

[paul-buerkner]

Great, thanks! I have also thought about power analysis as a cool bayesflow application but haven't gotten around actually implementing it. If you are interested in chatting regardless, I would be happy to chat :-)

[paul-buerkner]

We the latest fixes, I believe the coupling flow is now stable enough. At least for the example models we have tried. I will close this issue now but feel free to reopen it if new issues along similar lines arise.