The High-Luminosity LHC (HL-LHC) will bring significant computing challenges from the need to simulate one of the largest particle physics datasets ever recorded. Accurate and fast simulations will be essential for the new physics sensitivity of the HL-LHC. Advances in using machine learning to estimate probability distribution functions could produce faster and hardware agnostic elements of the HL-LHC compute-intensive simulations. Argonne researchers aim to use new techniques in AI to improve the accuracy and computing performance of these simulations. Argonne researchers have estimated the momentum resolution (which is a probability distribution function that depends on many input variables) of jets, objects that are made up of a large number of particles, using a machine learning algorithm. The machine learning algorithm was successfully trained to replicate not only the momentum resolution of the objects in the majority of cases but also in rare cases. Argonne researchers are now investigating the use of an uncertainty quantification method to estimate the momentum resolution.