Speaker
Description
Neural Simulation Based Inference (NSBI) has emerged as a powerful statistical inference methodology for large datasets with high-dimensional representations. NSBI methods rely on neural networks to estimate the underlying, multi-dimensional likelihood distributions of the data at a per-event level. This approach significantly improves the inference performance over classical binned approaches by circumventing the need for summary statistics. In practice, NSBI tools remain computationally expensive due to the per-event statistical inference step and the training of many large neural networks in order to prevent biases. Implementations of NSBI in High Energy Physics (HEP) analyses therefore require reliable orchestration on heterogeneous resources, high-throughput ingestion, and processing of variable-length event data.
The NEEDLE project aims to meet these demands by providing a flexible framework for distributed training on computing infrastructure alongside a toolbox of powerful NSBI methods. The framework reduces the operational cost associated with NSBI by implementing core features such as orchestration, data ingestion and experiment tracking. First, the orchestration is performed with a directed acyclic graph (DAG) workflow manager for the machine learning training and evaluation, versioning and experimentation life cycle. Second, models and datasets are tracked with pytorch Lightning, allowing for flexible and reproducible experiments. Finally, common HEP storage formats such as ROOT and parquet are read dynamically using dask-based libraries for optimal memory management.
In this contribution, we present the design principles, software architecture, and performance characteristics of the NEEDLE framework. This will be demonstrated for Neural Likelihood Ratio approaches on HEP open datasets.