Normalizing Flows to Reconstruct Pseudo-PDFs

Technical Summary: Normalizing Flows to Reconstruct Pseudo-PDFs

Problem Statement
Parton distribution functions (PDFs) are fundamental to understanding hadron structure but remain difficult to calculate directly due to the non-perturbative nature of quark and gluon interactions at low energy scales. Recent lattice frameworks propose investigating these quantities by computing non-local matrix elements (Ioffe-time distributions, ITDs) and solving an inverse problem to reconstruct the pseudo-PDF, $P(x, z^2)$. This inverse problem is ill-posed: the available lattice data is limited to a small number of Ioffe-time points ($N_\nu \sim 10$), while the PDF must be reconstructed over a continuous, discretized momentum fraction grid ($N_x = 129$). Previous approaches, including Bayesian methods and Gaussian Process (GP) regression, have addressed this by encoding smoothness and physical constraints, but they require careful, often manual, hyperparameter selection to control uncertainty in extrapolation regions where lattice information is absent.

Methodology
This work proposes a framework combining Gaussian Process priors with Invertible Neural Networks (INNs), a class of normalizing flows, to learn the full posterior distribution over PDFs. The methodology consists of three primary components:

1. Data Generation and GP Prior:
   Training data is generated by sampling PDFs from a Gaussian Process prior defined over a discretized $x$-grid. The GP kernel and mean are constrained via Lagrange multipliers to enforce physical boundary conditions: normalization ($\int q(x)dx = 1$) and the vanishing condition at $x=1$ ($q(1)=0$). These sampled PDFs are transformed into synthetic ITD data using a discrete linear operator derived from second-order finite element approximations of the Fourier transform relations.

2. Invertible Neural Network Architecture:
   The core of the framework is an INN based on Real NVP (Affine Coupling) layers. The network establishes a bijective mapping $f$ between the high-dimensional space of PDFs ($\mathbb{R}^{N_x}$) and a space containing the low-dimensional ITD observables and latent variables ($\mathbb{R}^{N_\nu + N_z}$).

   • Dimensionality Handling: Since $N_x > N_\nu$, the mapping introduces latent variables ($N_{noise}$) and zero-padding ($N_{pad}$) to ensure the input and output dimensions match ($N_x = N_\nu + N_{noise} + N_{pad}$).
   • Structure: The architecture comprises six stacked affine coupling layers. Each layer partitions the input and applies scaling and translation functions parameterized by feed-forward neural networks (3-layer MLPs with LeakyReLU activations).
   • Normalization: Input samples are normalized using the prior mean and point-wise variance to ensure stability across the divergent functional forms of PDFs.

3. Training Objective (Maximum Mean Discrepancy):
   The network is trained using the Maximum Mean Discrepancy (MMD) loss, a kernel-based metric that measures the distance between probability distributions. The loss function is bidirectional:

   • It minimizes the discrepancy between the network's predicted PDFs (inverse mapping) and the ground-truth training PDFs.
   • It minimizes the discrepancy between the forward-mapped ITD data (plus latent variables) and their prior distributions.
   • Additional terms enforce the forward mapping accuracy ($L_1$ norm) and ensure the padded output dimensions remain zero ($L_2$ norm).

Key Results
The authors performed a "closure test" using parametric PDFs to evaluate the framework:

• Latent Dimension Dependence: Reconstructions were consistent with the ground truth in the region close to $x=1$. However, in the extrapolation regions (low $x$ and high Ioffe time $\nu$), the quality of reconstruction depended on the dimensionality of the latent noise space. This suggests the latent variables effectively regulate the regularization of the inverse problem.
• Data Availability: As the range of available ITD data ($\nu$) increased, the reconstruction accuracy improved, a behavior consistent with other methods like GP regression and Bayesian inference.
• Constraint Preservation: The architecture successfully maintained physical constraints, specifically the normalization and the boundary condition $q(1)=0$, without explicit enforcement during the inference phase beyond the training data generation.

Significance and Claims
The paper positions this work as a proof of concept for mapping the hyperparameter selection steps of traditional GP approaches to a normalizing-flow architecture. The authors claim that this approach:

• Retains the flexibility of Gaussian Processes while utilizing the expressive power of invertible neural networks.
• Successfully learns a posterior distribution consistent with limited Ioffe-time data.
• Preserves physical constraints and extrapolation properties inherent to the problem.
• Offers a data-driven alternative to existing reconstruction methods, demonstrating that the specified calibration procedures found in GP regression can be learned by the network.

The authors remain modest regarding immediate applications, noting that the current method is constrained by the specific GP prior used for training. They suggest future work should focus on training on broader sets of prior samples, applying the method to real lattice data, incorporating Mellin moment measurements, and extending the framework to higher-dimensional quantities like generalized parton distribution functions.