Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Re-Envisioning Numerical Information Field Theory (NIFTy.re): A Library for Gaussian Processes and Variational Inference

Gordian Edenhofer, Philipp Frank, Jakob Roth, Reimar H. Leike, Massin Guerdi, Lukas I. Scheel-Platz, Matteo Guardiani, Vincent Eberle, Margret Westerkamp, Torsten A. Enßlin

TL;DR

Imaging from noisy data in astrophysics requires scalable, statistically rigorous inference over million- to billion-dimensional posteriors. The authors present NIFTy.re, a JAX-based rewrite of NIFTy, replacing a bespoke autodiff engine with JAX to accelerate GP priors and VI (MGVI and geoVI) and enable higher-order differentiation and GPU/TPU deployment. The approach standardizes priors via a latent Gaussian $\xi$ and builds flexible GP models (including Iterative Charted Refinement for irregular grids) within a modular Model class, facilitating complex forward models and likelihoods. The resulting library delivers substantial speedups (CPU and GPU) and interoperability with the broader JAX ecosystem, enabling scalable, uncertainty-aware imaging for astrophysical and non-astronomical applications.

Abstract

Imaging is the process of transforming noisy, incomplete data into a space that humans can interpret. NIFTy is a Bayesian framework for imaging and has already successfully been applied to many fields in astrophysics. Previous design decisions held the performance and the development of methods in NIFTy back. We present a rewrite of NIFTy, coined NIFTy.re, which reworks the modeling principle, extends the inference strategies, and outsources much of the heavy lifting to JAX. The rewrite dramatically accelerates models written in NIFTy, lays the foundation for new types of inference machineries, improves maintainability, and enables interoperability between NIFTy and the JAX machine learning ecosystem.

Re-Envisioning Numerical Information Field Theory (NIFTy.re): A Library for Gaussian Processes and Variational Inference

TL;DR

Imaging from noisy data in astrophysics requires scalable, statistically rigorous inference over million- to billion-dimensional posteriors. The authors present NIFTy.re, a JAX-based rewrite of NIFTy, replacing a bespoke autodiff engine with JAX to accelerate GP priors and VI (MGVI and geoVI) and enable higher-order differentiation and GPU/TPU deployment. The approach standardizes priors via a latent Gaussian and builds flexible GP models (including Iterative Charted Refinement for irregular grids) within a modular Model class, facilitating complex forward models and likelihoods. The resulting library delivers substantial speedups (CPU and GPU) and interoperability with the broader JAX ecosystem, enabling scalable, uncertainty-aware imaging for astrophysical and non-astronomical applications.

Abstract

Imaging is the process of transforming noisy, incomplete data into a space that humans can interpret. NIFTy is a Bayesian framework for imaging and has already successfully been applied to many fields in astrophysics. Previous design decisions held the performance and the development of methods in NIFTy back. We present a rewrite of NIFTy, coined NIFTy.re, which reworks the modeling principle, extends the inference strategies, and outsources much of the heavy lifting to JAX. The rewrite dramatically accelerates models written in NIFTy, lays the foundation for new types of inference machineries, improves maintainability, and enables interoperability between NIFTy and the JAX machine learning ecosystem.
Paper Structure (10 sections, 2 figures)

This paper contains 10 sections, 2 figures.

Table of Contents

  1. Summary
  2. Statement of Need
  3. Core Components
  4. Gaussian Processes
  5. Building Up Complex Models
  6. Variational Inference
  7. Performance of compared to
  8. Conclusion
  9. Acknowledgements
  10. References

Figures (2)

  • Figure 1: Data (left), posterior mean (middle), and posterior uncertainty (right) for a simple toy example.
  • Figure 2: Median evaluation time of applying the Fisher metric plus the identity metric to random input for NIFTy.re and NIFTy on the CPU (one and eight core(s) of an Intel Xeon Platinum 8358 CPU clocked at 2.60G Hz) and the GPU (A100 SXM4 80 GB HBM2). The quantile range from the 16%- to the 84%-quantile is obscured by the marker symbols.