Learning Informed Prior Distributions with Normalizing Flows for Bayesian Analysis

This paper demonstrates that normalizing flow models trained on previous posteriors can serve as effective, flexible priors for sequential Bayesian inference in high-dimensional spaces, provided that target distributions are unimodal and robust sampling algorithms are employed.

The Gist

The Big Idea: Teaching a Robot to "Remember" What It Learned

Imagine you are a detective trying to solve a complex crime. You have a suspect list (parameters) and a pile of evidence (experimental data). Your goal is to figure out which suspect is most likely guilty.

In the world of physics, specifically High-Energy Nuclear Physics, scientists do this constantly. They try to figure out the properties of the Quark-Gluon Plasma (a super-hot soup of particles that existed right after the Big Bang) by comparing computer models to real-world collision data.

The problem? The math is incredibly hard, and the "suspect list" is huge (dozens of variables). Running the full investigation from scratch every time you get a new piece of evidence takes forever and costs a fortune in computer power.

This paper introduces a clever shortcut using Normalizing Flows (NF). Think of an NF as a smart, shape-shifting robot that learns to mimic the "guilt profile" of a suspect based on past investigations.

---

The Problem: The "One-Size-Fits-All" Mistake

Traditionally, when scientists start a new investigation, they assume every suspect is equally likely to be guilty until proven otherwise. This is called a Uniform Prior. It's like walking into a room and guessing everyone is equally suspicious.

But what if you already did a deep investigation yesterday? You know that Suspect A is definitely innocent, and Suspect B is very likely guilty. If you ignore that knowledge and start with a "blank slate" today, you are wasting time.

The challenge is: How do you take the complex, messy results from yesterday's investigation and use them as a starting point for today's?

Yesterday's results aren't simple. They aren't just a single peak (like a bell curve). They might have:

• Multiple peaks: Two different suspects could both be guilty (Multi-modality).
• Weird shapes: The evidence might be skewed or stretched out.
• Secret connections: If Suspect A is guilty, Suspect C is also likely guilty (Correlations).

Standard math tools struggle to describe these weird shapes easily.

The Solution: The "Shape-Shifting Robot" (Normalizing Flow)

The authors trained a Normalizing Flow (NF) model. Here is how to think about it:

1. The Training Phase: Imagine you have a bag of marbles representing the results of a previous experiment. Some are red, some blue, clumped in weird shapes. You feed these marbles into the robot (the NF).
2. The Transformation: The robot learns a magical map. It figures out how to stretch, squeeze, and twist a simple, perfect circle of marbles (a standard Gaussian distribution) until it looks exactly like your weird, clumpy bag of results.
3. The Result: Now, instead of dealing with the messy bag, the robot can instantly generate new marbles that look exactly like the old results, preserving all the weird shapes and connections.

Why is this useful?
In a Sequential Bayesian Analysis, you use the results of Experiment A as the "Prior" (starting knowledge) for Experiment B.

• Old way: You try to approximate the messy results of Experiment A with a simple bell curve. You lose information.
• New way (This paper): You use the robot to perfectly mimic the messy results of Experiment A. You feed this perfect mimic into Experiment B.

The Experiment: Testing the Robot

The authors tested this on a real physics problem involving J/ψ particle production in collisions. They had two sets of data:

1. Data A: Collisions with a single proton (γ + p).
2. Data B: Collisions with a heavy lead nucleus (γ + Pb).

They wanted to see if they could analyze Data A, turn the results into a robot-prior, and then analyze Data B, to see if they got the same answer as if they had analyzed both datasets together at once (the "One-Shot" method).

The Results: It Depends on the Shape

Scenario 1: The Smooth Hill (Success)
When the data formed a nice, single "hill" (unimodal), the robot worked perfectly.

• Analogy: Imagine the suspect is clearly guilty. The robot remembers the shape of the guilt perfectly. When you add new evidence, the robot updates the location of the guilt accurately.
• Outcome: The sequential method matched the "One-Shot" method almost exactly.

Scenario 2: The Double Peak (Failure)
Sometimes, the data has two distinct peaks (bimodal). Maybe the physics allows for two very different scenarios to both be true.

• Analogy: Imagine the evidence points to two different suspects being guilty, but they are in different rooms.
• The Trap: If the first experiment (Data A) accidentally focuses on the "left room" and misses the "right room," the robot learns to only generate marbles in the left room.
• The Disaster: When you feed this biased robot into the second experiment (Data B), it can't find the "right room" because it was trained to ignore it. The robot gets stuck in a local trap.
• Outcome: The sequential method failed to find the full truth. The "One-Shot" method (looking at all data at once) found both rooms, but the step-by-step method missed one.

The Lesson: Don't Skip the Hard Parts

The paper highlights two major takeaways:

1. Robots are great, but they need good training: The authors found that training the robot using a specific mathematical metric called KL Divergence worked best. It's like teaching the robot with a very strict grading system that penalizes it for missing any part of the shape.
2. The Detective Tool Matters: They compared two "search engines" (MCMC samplers) used to explore the results.
   • emcee: A standard, reliable search engine.
   • pocoMC: A high-tech, turbo-charged search engine.
   • Result: When the data was tricky (multi-modal), the standard engine got lost. The turbo-charged engine found the hidden peaks. This proves that if you are using advanced AI (the robot), you also need advanced search tools to navigate the results.

Summary

This paper is about building a better memory for scientific investigations.

• The Goal: Save time and computer power by reusing knowledge from past experiments.
• The Tool: A "Normalizing Flow" robot that learns the complex, messy shapes of past data and turns them into a flexible starting point for new data.
• The Catch: It works beautifully when the answer is clear and simple. However, if the answer is complex (with multiple possibilities), you have to be very careful. If your "memory" misses a possibility in step one, you will never find it in step two.

In short: It's a powerful new way to do science, but it requires smart tools and a cautious approach to ensure you don't accidentally forget the "other suspects" hiding in the shadows.

Technical Summary

1. Problem Statement

In high-energy nuclear physics, Bayesian inference is the standard framework for constraining model parameters ($\theta$) against experimental data ($y_{exp}$). However, conducting sequential Bayesian analysis—where results from one dataset are used as priors for the next—presents significant challenges:

• Complex Priors: Posterior distributions from previous analyses are often non-Gaussian, multimodal, and contain complex parameter correlations that cannot be easily represented by standard uniform or Gaussian priors.
• Sampling Limitations: Using raw MCMC samples from a previous run as a prior is impractical in high-dimensional spaces (20–50 parameters) due to the discrete nature of the samples and the difficulty of interpolating continuous distributions.
• Computational Cost: New inference tasks often involve expensive theoretical models. Shrinking the prior parameter space using informative priors is crucial for efficiency, but only if the prior accurately captures the true distribution.

The core problem is how to construct a flexible, continuous, and informative prior distribution that accurately reproduces complex posterior structures from previous analyses to facilitate efficient sequential inference.

2. Methodology

The authors propose using Normalizing Flows (NF) as a generative model to learn these informative priors.

A. Normalizing Flow Framework

• Concept: An NF constructs a bijective mapping ($F$) between a simple reference distribution (typically a multivariate Gaussian, $p_G(\omega)$) and a complex target distribution ($p(\theta)$).
• Transformation: $\theta = F(\omega)$. The probability density is transformed via the change of variables formula involving the Jacobian determinant: $p(\theta) = p_G(\omega) |\det(\partial \omega / \partial \theta)|$.
• Architecture: The model uses a Real NVP (Real-valued Non-Volume Preserving) architecture with affine coupling layers and a final scale-and-shift layer.

B. Training Strategies

The paper compares two distinct training approaches to learn the NF from MCMC samples:

1. Supervised Learning:
   • Uses samples $\{\theta_i\}$ and their associated log-likelihoods (or posterior weights).
   • Loss Functions: Minimizes Jeffreys' Divergence or Kullback–Leibler (KL) Divergence between the true posterior and the NF-approximated distribution.
2. Unsupervised Learning:
   • Uses only the sample vectors $\{\theta_i\}$ without explicit probability weights (useful when only samples are available, e.g., from PCA).
   • Objective: Maximizes the log-likelihood of the samples under the model by learning the distribution shape directly.

C. Sequential Bayesian Workflow

The workflow involves a two-stage process:

1. Stage 1: Perform Bayesian inference on Dataset $D_1$ to obtain a posterior $P(\theta|D_1)$.
2. Training: Train an NF model on the samples from $P(\theta|D_1)$.
3. Stage 2: Use the trained NF as the prior $P(\theta)$ for inference on Dataset $D_2$.
4. Validation: Compare the sequential posterior $P(\theta|D_1, D_2)$ against a "one-shot" joint inference (where $D_1$ and $D_2$ are analyzed simultaneously) to verify consistency.

D. Sampling Algorithms

The study utilizes two MCMC samplers:

• pocoMC: An advanced, preconditioned MCMC sampler suited for high-dimensional, complex distributions (no gradient required).
• emcee: A standard affine-invariant ensemble sampler used as a baseline for comparison.

3. Key Contributions

• Extension to Unsupervised Learning: The authors extend the NF framework to unsupervised scenarios where posterior densities (weights) are unavailable, demonstrating its viability for learning distributions from raw samples.
• Systematic Comparison of Loss Functions: They rigorously compare Jeffreys' divergence, KL divergence, and log-likelihood objectives, finding that KL divergence (supervised) and log-likelihood (unsupervised) yield the most accurate reproductions.
• Sequential Inference Benchmark: They provide a systematic verification of sequential Bayesian analysis against joint inference in a high-dimensional (7-parameter) nuclear physics context.
• Algorithmic Sensitivity: The study highlights the critical importance of robust MCMC samplers (like pocoMC) over standard ones (emcee) when dealing with multimodal posteriors in sequential workflows.

4. Results

The methodology was applied to a study of diffractive $J/\psi$ production in $\gamma+p$ and $\gamma+Pb$ collisions (based on Color Glass Condensate theory).

• NF Training Performance:

  • NF models successfully reproduced unimodal, well-behaved distributions (e.g., from $\gamma+p$ data) with very low KL divergence ($\langle D_{KL} \rangle \approx 0.025$).
  • For more complex distributions with boundary peaks and heavy tails (e.g., $\gamma+Pb$ data), the NF still captured the features reasonably well, though with slightly higher divergence ($\langle D_{KL} \rangle \approx 0.052$).
  • KL loss consistently outperformed Jeffreys' divergence in supervised training. Unsupervised training achieved comparable accuracy to supervised KL training.

• Sequential vs. Joint Inference:

  • Case 1 (Unimodal/Compatible): When the first stage ($\gamma+p$) provided a broad prior that covered the modes of the second stage, the sequential inference (using NF prior + pocoMC) reproduced the joint posterior with high fidelity ($\langle D_{KL} \rangle \approx 0.1$).
  • Case 2 (Multimodal/Tension): When the order was reversed (starting with $\gamma+Pb$), the first-stage posterior was multimodal. The NF prior failed to capture a specific mode ($\Lambda_{QCD} \approx 0.1$ GeV) that was favored by the joint analysis but suppressed in the first stage. Consequently, the sequential inference failed to recover this mode, resulting in a massive KL divergence ($\langle D_{KL} \rangle = 6.482$).
  • Sampler Comparison: Using the standard emcee sampler in the second stage (even with a good NF prior) resulted in a complete failure to reproduce the joint posterior ($\langle D_{KL} \rangle = 2.966$), whereas pocoMC succeeded.

• Simplified Case (Appendix): In a simplified scenario without multimodal structures (splitting $\gamma+Pb$ into integrated and differential cross-sections), the sequential inference worked perfectly regardless of the order, confirming that the failure in the main case was due to mode collapse in the first stage, not the NF method itself.

5. Significance and Conclusion

• Efficiency: NF-based priors offer a practical, memory-efficient way to compress complex posterior information into a generative model, enabling efficient sequential Bayesian analysis in high-dimensional spaces.
• Cautionary Insight: The paper establishes a critical limitation: Sequential inference is order-dependent when multimodality or dataset tension exists. If an early stage of analysis misses a relevant mode (collapses the distribution), that information is lost forever in subsequent stages, regardless of the sophistication of the NF prior.
• Algorithmic Necessity: The results underscore that advanced MCMC samplers (like pocoMC) are essential for exploring complex, multimodal posterior spaces. Standard samplers may get trapped in local modes, rendering even a perfect prior useless.
• Future Outlook: This framework provides a systematic method for reusing information across datasets in nuclear and particle physics, reducing computational costs for expensive simulations while avoiding the oversimplification of uniform priors.