PaperlayerView exact revision on arXiv

Source-linked AI summary

Flexible Spectral-Normalized Neural Gaussian Process for Dynamic Aperture Prediction

Yousra El-Bachir, Frederik Van der Veken, Davide di Croce, Carlo Emilio Montanari, Massimo Giovannozzi, Ekaterina Krymova, Tatiana Pieloni

arXiv:2609.08620v1physics.acc-phcs.LG

TL;DR

Reliable uncertainty disentanglement and scalable, well-calibrated methods remain open problems. The paper introduces an empirical Bayes flexible hetSNGP trained with hyperparameter optimization integrated into the training loop; on LHC tracking data, it achieves test-set results comparable to DDU, marginally outperforms MC Dropout, and produces more consistent, interpretable uncertainty estimates than DDU across three configurations.

  • Problem

    Reliably disentangling epistemic and aleatoric uncertainty and developing scalable, well-calibrated methods remain open problems.

    Source: §1 · p. 2

  • Method

    The flexible hetSNGP uses empirical Bayes to tune parameters during training and learns class-specific variances for class-dependent heteroscedasticity.

    Source: §2 · p. 5, §5 · p. 10

  • Results

    DDU and auto-hetSNGP achieve comparable test-set results and marginally outperform MC Dropout, while auto-hetSNGP gives more consistent, interpretable uncertainty estimates than DDU across three configurations.

    Source: §4 · pp. 9–10

  • Takeaways & Limitations

    The framework is presented as an efficient approach for large-scale simulation-driven scientific problems where costly validation or manual tuning are bottlenecks.

    Source: §5 · p. 10

    The current method uses a Laplace approximation; exploring more complex posterior inference is left for future work.

    Source: §5 · p. 10

Abstract

from arXiv · show

We address the challenge of scalable uncertainty quantification in large-scale scientific applications, where complex state-of-the-art machine learning methods are often computationally infeasible. Our primary contribution is a simple yet effective empirical Bayes method for automatically tuning the hyperparameters of a flexible, heteroscedastic Spectral-normalized Neural Gaussian Process. This approach retains the expressiveness and uncertainty-awareness of semi-Bayesian neural models while significantly reducing the computational burden by integrating hyperparameter learning directly into the training loop. We demonstrate the practical impact of our method on the task of estimating the dynamic aperture in circular particle accelerators, a fundamental problem in high-energy physics colliders and storage rings, using simulation data from the case of the Large Hadron Collider at CERN. Traditional approaches to DA estimation require extensive particle-tracking simulations, which are prohibitively time-consuming and resource-intensive. Our results show that the proposed method achieves competitive predictive performance and well-calibrated uncertainty estimates at much lower computational cost than state-of-the-art approaches. We stress that, beyond this application, the proposed empirical Bayes framework offers a general solution for training heteroscedastic neural models in situations where manual hyperparameter tuning is impractical. Accordingly, we anticipate that this framework can be applied to other domains that encounter comparable computational limitations.

1 Introduction

Reliable, scalable uncertainty estimation remains difficult, while dynamic-aperture prediction is important for beam stability and traditionally relies on costly particle-tracking simulations. The paper proposes an efficient approach that estimates total predictive uncertainty during training and predicts the stable region from simulation data.

1 Introduction

  • The paper targets open problems in disentangling epistemic and aleatoric uncertainty and developing scalable, well-calibrated methods.It instead estimates total predictive uncertainty directly within model training.Source: §1 · pp. 1–2
  • Deep ensembles can be computationally prohibitive, while SNGP requires manual hyperparameter tuning despite its uncertainty benefits under distributional shifts.The paper proposes integrating SNGP hyperparameter tuning into training to improve usability without compromising performance.Source: §1 · p. 2
  • DA is the phase-space region where particle trajectories remain bounded, and predicting it matters for beam stability and collider operation.Traditional LHC studies scan a high-dimensional parameter space using computationally intensive particle-tracking simulations.Source: §1 · p. 2
  • The task combines nonlinear stable-region structure, an ambiguous chaotic boundary, and simulations involving millions of particles and hundreds of control variables.The study merges the chaotic boundary with the unstable region and focuses on predicting the stable region.Source: §1 · pp. 2–3

2 The flexible auto-hetSNGP model

The flexible auto-hetSNGP model combines a spectrally normalized residual network with a Gaussian-process output and heteroscedastic uncertainty. It extends the model with learnable class-dependent variances and empirical Bayes hyperparameter tuning integrated into training.

2 The flexible auto-hetSNGP model

  • The model represents each class logit with a predictive mean and a zero-mean stochastic deviation that captures heteroscedastic uncertainty.The stochastic term combines class-specific noise with shared global noise, allowing correlated noise across classes.Source: §2 · pp. 3, 5
  • The residual network feeds a Gaussian-process output layer approximated by random Fourier features, which provides distance awareness and sensitivity to out-of-distribution samples.Its spectrally normalized layers are intended to preserve Lipschitz continuity and geometric consistency in the feature space.Source: §2 · p. 4
  • The structured noise uses a low-rank global component and a class-specific component, with the latter helping keep the latent covariance strictly positive definite.The low-rank condition is R ≪ K.Source: §2 · p. 5
  • Temperature scaling can be absorbed into learnable noise parameters when d(x_i) and V(x_i) are linear, but the formulation retains explicit scaling for broader settings.The model also introduces learnable class-dependent variance terms.Source: §2 · p. 5
  • Class-dependent variances for β_c and empirical Bayes tuning are integrated into training, with the variances enabling class-specific regularization.The framework tunes the additional parameters during training to enhance scalability.Source: §2 · p. 5

3 Empirical Bayes optimization

The empirical Bayes framework learns deterministic hyperparameters during training and uses posterior sampling over stochastic parameters at prediction time. This double-optimization design reduces repeated neural-network computation while supporting flexible predictive uncertainty.

3 Empirical Bayes optimization

  • The method separates deterministic parameters ψ, learned by maximizing the log-marginal likelihood, from stochastic parameters θ, sampled from their posterior at inference.The learned ψ is held fixed for inference, and posterior sampling uses a Laplace approximation.Source: §3 · p. 6, §3.1 · p. 6

3.1 Double maximization

  • The double maximization first optimizes the penalized likelihood over θ for fixed ψ, then uses the resulting Laplace approximation to sample from θ’s posterior.A preconditioned Newton–Raphson minimizer is used for the penalized likelihood, with Hessian adjustments to support stable optimization.Source: §3.1 · pp. 6–7
  • Separating ψ from θ reduces computational overhead by training the neural network once and sharing one forward pass across marginal- and penalized-likelihood evaluations.The penalized-likelihood optimization is performed once at inference time for predictive evaluation.Source: §3.1 · p. 7

3.2 Posterior predictive inference with Monte-Carlo sampling

  • Posterior predictive inference draws Monte Carlo samples of θ and predicts the class with the highest probability.The approach requires tuning only the spectral-normalization factor and network architecture, similar to DDU.Source: §3.2 · p. 8

4 Results for LHC simulated tracking data

On LHC simulated tracking data, auto-hetSNGP and DDU achieved comparable test-set results and marginally outperformed MC Dropout. Auto-hetSNGP produced more consistent, interpretable uncertainty estimates on localized configurations, while DDU performed slightly better overall on the full test set.

  • The training set contained 1’065’273 particles, and the test set contained 131’593 particles; the study compared auto-hetSNGP with MC Dropout and DDU.The training set included 475’362 stable and 589’911 unstable particles; the test set included 57’918 stable and 73’675 unstable particles.Source: §4 · p. 8
  • Deep ensembles were infeasible under the computational constraints, and each method was evaluated with a single training run, so seed-to-seed variance was not assessed.Source: §4 · pp. 8, 10
  • Auto-hetSNGP and DDU achieved comparable test-set results, marginally outperforming MC Dropout; DDU was slightly better overall, while auto-hetSNGP excelled on localized configurations.The comparison used the best models evaluated on the test set.Source: Table 2 · p. 9, §4 · pp. 9–10
  • Auto-hetSNGP produced more consistent, interpretable uncertainty estimates than DDU on three configurations, aligning more closely with expected uncertainty behavior.DDU showed higher uncertainty in the stable region for configurations 1186 and 29876, and unreliable estimates for configuration 29472.Source: §4 · p. 10
  • The test set included three localized configurations where auto-hetSNGP outperformed DDU, consistent with its class-dependent noise modeling versus DDU’s shared class Gaussian.The paper relates this difference to auto-hetSNGP’s ability to capture heteroscedastic uncertainty across input regions.Source: §4 · p. 10

5 Discussion

The discussion presents an empirical-Bayes hetSNGP for scalable uncertainty-aware dynamic-aperture prediction, reporting competitive performance, calibrated uncertainty, and training efficiency. It also outlines possible applications beyond accelerator physics and future extensions.

5 Discussion

  • The proposed model achieves competitive predictive performance, well-calibrated uncertainty, and training efficiency similar to the simplest state-of-the-art method.It is more computationally tractable than deep ensembles and shows more consistent uncertainty behavior than DDU on configuration-specific test cases.Source: §5 · p. 10
  • Separating deterministic parameters ψ from stochastic parameters θ enables efficient training and Laplace-approximation inference over θ at test time.A preconditioned Newton–Raphson procedure makes posterior inference tractable because θ is low-dimensional.Source: §5 · p. 10
  • The framework is proposed as applicable to climate science, biological modeling, and engineering design under similar computational constraints.Future work could use more complex posterior inference or couple the method with active-learning or Bayesian-optimization acquisition functions to focus effort on uncertain regions.Source: §5 · p. 10

A Performance on different configurations

The comparison considers low-, medium-, and high-DA configurations and expects uncertainty to concentrate near the stable–unstable boundary.

A Performance on different configurations

  • The analysis considers configurations corresponding to low, medium, and high dynamic aperture.The passage introduces these configurations without specifying their individual outcomes.Source: A Performance on different… · p. 13
  • The stable–unstable boundary is expected to have higher uncertainty than the complementary regions.This is the stated qualitative expectation for predicted boundaries.Source: A Performance on different… · p. 13
  • A well-performing boundary prediction should combine few errors around the DA region with uncertainty concentrated reasonably around the predicted boundary.Source: A Performance on different… · p. 13

A.1 Comparison for configuration 1186. Medium DA.

For medium-DA configuration 1186, the methods perform reasonably well overall, but their uncertainty behavior differs.

A.1 Comparison for configuration 1186. Medium DA.

  • All methods perform reasonably well on medium-DA configuration 1186, although MC Dropout overestimates the DA for larger angles.Source: §A.1 · p. 13
  • MC Dropout and auto-hetSNGP show uncertainty qualitatively consistent with the expected behavior in configuration 1186.Source: §A.1 · p. 13
  • DDU assigns high uncertainty within the stable-particle region, contrary to the expected uncertainty pattern.Source: §A.1 · p. 13

A.2 Comparison for configuration 29472. Low DA.

For low-DA configuration 29472, DDU and auto-hetSNGP have good overall accuracy, while MC Dropout overestimates the aperture; uncertainty patterns also differ.

A.2 Comparison for configuration 29472. Low DA.

  • DDU and auto-hetSNGP demonstrate good overall accuracy on low-DA configuration 29472, whereas MC Dropout significantly overestimates the DA.Source: §A.2 · p. 13
  • MC Dropout uncertainty is high around the incorrectly predicted class boundary, as expected.Source: §A.2 · p. 13
  • DDU has unexpectedly high uncertainty in the unstable region, while auto-hetSNGP highlights the uncertain boundary reasonably.Source: §A.2 · p. 13

A.3 Comparison for configuration 29876. High DA.

For configuration 29876, DDU and auto-hetSNGP show good overall accuracy, while MC Dropout significantly underestimates the dynamic aperture. Their uncertainty patterns differ around the class boundary and stable region.

A.3 Comparison for configuration 29876. High DA.

  • DDU and auto-hetSNGP show good overall accuracy, while MC Dropout significantly underestimates DA.The comparison is for configuration 29876.Source: §A.3 · p. 13, Figure 5 · p. 16
  • MC Dropout uncertainty concentrates around the incorrectly predicted class boundary, while DDU uncertainty is unexpectedly high in the stable region.Source: §A.3 · p. 13
  • Auto-hetSNGP reasonably highlights the uncertain boundary.Source: §A.3 · p. 13

B 2D example

The 2D example uses concentric-circle data with instance-dependent label noise to examine uncertainty estimates. Models perform comparably, while auto-hetSNGP has marginally higher accuracy; the DDU and DE comparison is qualified by limited tuning.

B 2D example

  • The synthetic benchmark samples inner- and outer-circle inputs with isotropic Gaussian perturbations and instance-dependent label noise.Source: B 2D example · p. 17
  • Labels are flipped with probability depending on input orientation, producing an observed-label posterior that combines true-label probability and flip probability.Source: B 2D example · p. 17
  • DE and SNGP-based models use predictive logit variance, while DDU uses estimated feature density for uncertainty; Figure 7 presents these measures.For SNGP-type methods, the expectation is over posterior noise; for DE, it is over ensemble-member randomness.Source: B 2D example · p. 17, Figure 7 · p. 18
  • Models exhibit comparable performance, with auto-hetSNGP marginally higher accuracy on the synthetic benchmark.Table 3 reports test accuracy across methods and flip scales.Source: B 2D example · p. 17, Table 3 · p. 17
  • Because DDU and DE were not extensively tuned on this synthetic benchmark, their accuracy gap should not be over-interpreted as a general performance ranking.Source: B 2D example · p. 17
Loading 2609.08620v1…