Improving High-Energy Unfolding in Omnifold via Learned Pre-Training Weights

In particle physics, reconstructing the true distribution of events from detector-distorted data is a necessary but historically challenging task. Traditionally, this had been accomplished using Iterative Bayesian Unfolding (IBU), a machine learning that predicts the underlying distribution from the detector data. However, IBU has limitations such as requiring the binning of data. Recently, the field has seen major improvements through AI. In 2020, researchers introduced the Omnifold method, which uses neural networks to iteratively learn weights that transform detector-level data into an estimate of the true underlying distribution event-by-event.

However, like many machine learning systems, Omnifold struggles when faced with highly imbalanced datasets. In particular, it tends to overfit to densely populated regions while neglecting bins with sparse data.

My project aims to mitigate this imbalance by introducing pre-training weights to re-emphasize underrepresented regions before training begins, generated by another nueral network. I plan to test how this strategy affects training stability, accuracy across multiple datasets, and performance under different reweighting multipliers and iteration counts.

Ultimately, the goal is to develop a reusable method to improve Omnifold’s performance on steeply falling distributions such as momentum spectra of hadrons within jets. Our findings aim to contribute toward more reliable unfolding techniques, especially in datasets where certain events are rare but scientifically significant.