Training instability with spline coupling flows in BayesFlow 2.0.7

Hi,

I am using BayesFlow 2.0.7 with the PyTorch backend and have repeatedly encountered training instabilities with spline coupling flows. Training initially proceeds normally, but after a variable number of epochs, the loss and gradients suddenly become NaN.

My flow architecture uses five inference variables, three coupling layers, spline transformations with four bins, and MLP subnetworks with two hidden layers of 96 units each. ActNorm is disabled.

I have verified that the adapted input data and inference variables are finite. I have also monitored the gradients, and the non-finite gradients appear to originate in the flow transformation layers rather than the summary network.

Are there any known numerical-stability issues with this type of architecture in BayesFlow 2.0.7? I noticed that coupling flows and their MLP subnetworks were substantially revised in later versions, particularly in PR #671. I mainly want to rule out a known BayesFlow issue that may already have been addressed in a newer release.

Thank you!

Hi Simon, please, update to version 2.0.13, as there were many improvements (too many to list here).

ActNorm is highly recommended for normalizing flows. By fine, do you mean that your parameters and observables are properly standardized (e.g., passing standardize="all" to the workflow)? Finally, expanding the default_domain of splines to let’s say (-5., 5., -5., 5) can help if the targets have heavy tails. More aggressive gradient clipping can also help.

Unless you have a specific question tied to normalizing flows, I recommend using Flow Matching / Diffusion as your inference network for NPE.

Hi, thank you for the reply.

By “finite,” I mean that I globally scale my observations using fixed constants so that they remain approximately within the [−1,1] range. However, I don’t standardize the observations (only the targets). I will test if this could solve my problem and also turn ActNorm on.

I’m also already setting the default domain to (-4., 4., -4., 4).

What seemed strange to me is that I have two models using the same simulator and summary-network architecture. One infers only one target variable, while the other infers four. The first model trains without issues, whereas the second repeatedly becomes unstable and crashes.

The latter is admittedly a much harder inference problem. However, I would have expected the model to simply perform poorly or fail to recover the parameters, rather than become numerically unstable. I have generally encountered quite a few stability issues when using spline transformations.

Also, is there a reason you recommend switching to flow-matching / Diffusion net over a NF net?

Thanks for your help :slight_smile:

It’s possible that one of the four parameters has a heavy tail (Spline flow can be unstable for heavy-failed distributions · Issue #721 · bayesflow-org/bayesflow · GitHub). I’ve also observed it for time-varying models and increasing the default domain to a factor of 5 solved the issue.

Diffusion / FM (free-form) are just more performant, more flexible, more stable, and less fragile to model misspecification (some properties we explore in our tutorial paper: [2512.20685] Diffusion Models in Simulation-Based Inference: A Tutorial Review)

@LarsKue Of course, it would be nice to bake in the stability guardrails for RQS flows.

1 Like