Loss Landscape Analysis for Reliable Quantized ML Models for Scientific Sensing

TL;DR

The paper tackles reliability of quantized ML models in scientific sensing by introducing a loss-landscape analysis framework that integrates visualization, CKA similarity, Hessian metrics, and mode connectivity. This approach enables a priori assessment of robustness to input and weight perturbations without exhaustive retraining, highlighting how quantization and regularization shape the loss surface. Key findings show that gently-shaped, flat minima correlate with robustness, that regularizers can improve or sometimes degrade robustness depending on the model, and that increasing precision does not universally enhance robustness. The work provides actionable guidance for including robustness in Pareto optimization, supporting more reliable and adaptive scientific sensing at the edge with efficient design workflows.

Abstract

In this paper, we propose a method to perform empirical analysis of the loss landscape of machine learning (ML) models. The method is applied to two ML models for scientific sensing, which necessitates quantization to be deployed and are subject to noise and perturbations due to experimental conditions. Our method allows assessing the robustness of ML models to such effects as a function of quantization precision and under different regularization techniques -- two crucial concerns that remained underexplored so far. By investigating the interplay between performance, efficiency, and robustness by means of loss landscape analysis, we both established a strong correlation between gently-shaped landscapes and robustness to input and weight perturbations and observed other intriguing and non-obvious phenomena. Our method allows a systematic exploration of such trade-offs a priori, i.e., without training and testing multiple models, leading to more efficient development workflows. This work also highlights the importance of incorporating robustness into the Pareto optimization of ML models, enabling more reliable and adaptive scientific sensing systems.

Figures (9)

• Figure 1: Workflow of the models in this study. (a) The ECON-T model workflow econ, demonstrating the lossy data compression pipeline designed for deployment in the high-radiation environment of the Large Hadron Collider (LHC). (b) The Fusion model workflow Wei:2023mma, illustrating active feedback control in magnetic confinement fusion devices.
• Figure 2: Comparison of 3D loss landscape visualization methods: (a) uses the top-2 eigenvectors of model parameters, while (b) uses two random orthogonal directions.
• Figure 3: Comparison of loss curves computed by perturbing the models along the top eigenvector of the Hessian matrix (i.e., varying parameter $\alpha$ in Eq. \ref{['eqn:loss-land']}, reported on the x-axis, while keeping $\beta=0$). (a) and (b) compare the loss line of models trained with different precision respectively for ECON-T and Fusion models, while (c) and (d) compare models fine-tuned with different regularization techniques respectively for ECON-T and Fusion models.
• Figure 4: Analysis of loss landscape metrics for ECON-T (left column) models and Fusion models (right column) fine-tuned with different regularization strategies across varying precision levels. Subplots show: (a) and (b) CKA similarity, which evaluates representational alignment among models; (c) and (d) Hessian trace, capturing the overall curvature of the loss landscape where the model is converged; and (e) and (f) mode connectivity, indicating the presence of barriers among different minima. Regularization methods include Baseline (no regularization), Jacobian regularization, and orthogonal regularization.
• Figure 5: Evaluation of ECON-T models (top row) and Fusion models (bottom row) robustness under different input or weight perturbations. Each subplot represents performance benchmarks on specific scenarios: (a and e) clean data, (b and f) perturbed data with Gaussian noise, (c and g) perturbed data with salt-and-pepper noise, and (d and h) flipping the five most sensitive bits. The models are trained with three regularization methods: Baseline (no regularization), Jacobian regularization, and orthogonal regularization.