Training instability with spline coupling flows in BayesFlow 2.0.7

General
1

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!

2

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.