Hardware-Accelerated Variational Inference using Inverse Autoregressive Flows (JAX/Equinox)

[DJLacombeTTU]

Hi everyone,

I’ve been working heavily with JAX-backend PyMC to scale up some large-N and densely correlated models, and I wanted to share an IAFOptimizer (Inverse Autoregressive Flow) implementation that might be useful for others in the community facing MCMC bottlenecks.

I want to make it absolutely clear up front: as an applied econometrician, not a computer scientist, I am not the inventor of this methodology. I simply refactored and adapted the incredible code and concepts originally provided in this recent paper. My goal was to package their work into a robust, auto-tuning inference engine that integrates cleanly with PyMC and ArviZ for applied research.

The Neural Network Architecture

To map standard Gaussian noise into highly complex posterior distributions, this optimizer leverages a deep neural network under the hood.

The architecture is built using the Equinox library and consists of a Masked Autoencoder for Distribution Estimation (MADE) coupled with a Deep Sigmoidal Bijector. Because we are training a neural network to maximize the Evidence Lower Bound (ELBO), ensuring stable gradient dynamics is critical. The network architecture relies on several deep learning optimizations:

• GELU Activations: Provides smoother continuous derivatives within the neural network, preventing sharp cliffs in the log-determinant Jacobian calculations.
• Identity Initialization: The final layer of the neural conditioner initializes at absolute zero. This guarantees the flow starts as a safe Identity function ($y=x$) rather than a wildly distorted manifold.
• Optax Schedulers & Regularization: Uses a Warmup Cosine Decay schedule, AdamW (weight decay), and global norm gradient clipping to prevent extreme tail samples from destroying the neural network's learned geometry.
• IWAE Objective: Includes a toggle for an Importance Weighted Autoencoder (IWAE) loss, which shifts the network's optimization landscape from "mode-seeking" to "mass-covering."

Benchmark Results

Using an Optuna integration to find the optimal neural network depth, mixture components, and learning rate, I ran the flow against four standard topologies.

The results highlight exactly where these neural-network-backed Normalizing Flows scale brilliantly (massive datasets and dense covariance matrices) and where they are still strictly bound by proper parameterization hygiene (hierarchical funnels).

Benchmark Model	Dimensions / Size	Objective	Pareto $k$ Statistic	IS-ESS (out of 5000)	Result
Massive Vectorized Logistic	100,000 observations	Standard ELBO	0.4143	4960	Success (Perfect tail coverage)
Dense High-Dim Correlation	100 correlated parameters	Standard ELBO	0.1002	3911	Success (Bypasses $O(D^3)$ mass matrix inversion)
Centered Funnel of Hell	50 parameters	IWAE ($k=16$)	1.1689	6	Failure (Network cannot warp into infinite neck)
Uncentered Funnel (Isotropic)	50 parameters	Standard ELBO	0.2821	4939	Success (Flow effortlessly maps the reparameterized space)

Note: Diagnostics generated via arviz.psislw.

For applied users working with massive datasets or spatial matrices where NUTS leapfrog evaluations become paralyzed, this has been an incredible pre-processor and alternative inference tool.

The full refactored codebase, neural network architecture, and benchmarking suite can be found in my repository here: https://github.com/djlacombettu/iaf-pymc-optimizer

I would love to hear any feedback from the core devs or anyone experimenting with JAX-based VI on whether this neural architecture is useful.