Learning Deep Hybrid Models with Sharpness-Aware Minimization

Hybrid modeling, the combination of machine learning models and scientific mathematical models, enables flexible and robust data-driven prediction with partial interpretability. However, effectively the scientific models may be ignored in prediction due to the flexibility of the machine learning model, making the idea of hybrid modeling pointless. Typically some regularization is applied to hybrid model learning to avoid such a failure case, but the formulation of the regularizer strongly depends on model architectures and domain knowledge. In this paper, we propose to focus on the flatness of loss minima in learning hybrid models, aiming to make the model as simple as possible. We employ the idea of sharpness-aware minimization and adapt it to the hybrid modeling setting. Numerical experiments show that the SAM-based method works well across different choices of models and datasets.

💡 Research Summary

The paper addresses a fundamental challenge in hybrid modeling, where a scientific (mechanistic) model is combined with a deep learning component. In many settings the neural network is so expressive that it can fit the data alone, rendering the scientific model effectively unused and making the hybrid approach pointless. Traditional solutions rely on handcrafted regularizers that penalize the neural network’s complexity (e.g., L2 norm, functional norm), but such regularizers are difficult to design for arbitrary hybrid architectures, especially when the models are nested or involve complex compositions.

The authors propose to tackle this issue by focusing on the flatness of the loss landscape rather than on explicit regularization of the neural network. Flat minima—regions where the loss remains low under small parameter perturbations—are associated with simpler models, better generalization, and lower description length. By encouraging the learning process to converge to flat minima, the neural network part can be kept “simple” while the scientific model parameters become more identifiable.

To operationalize this idea, the paper adapts Sharpness‑Aware Minimization (SAM), a recent optimizer that minimizes the worst‑case loss within a radius ρ around the current parameters. The key modification is that the SAM perturbation is applied only to the neural network parameters ϕ, leaving the scientific parameters θ untouched. This reflects the intuition that we want the loss to be flat with respect to ϕ (so that ϕ does not over‑fit) but not necessarily flat with respect to θ, because a flat loss in θ would imply that many different scientific parameter values achieve the same performance, i.e., the scientific model remains unidentifiable.

Algorithm 1 details the training loop: for each minibatch the loss is computed, a SAM perturbation ε*ϕ is generated (using the original SAM formula or its adaptive variants such as ASAM or Fisher‑SAM), gradients are evaluated at the perturbed point, and the original parameters are updated using these gradients. The method is architecture‑agnostic and can be combined with any hybrid composition hθ,ϕ(x)=gϕ(x,fθ(x)), including additive, nested, or even neural‑ODE based hybrids.

The theoretical discussion connects flatness to the KL‑divergence decomposition of a Bayesian posterior over (θ,ϕ). When the posterior over θ spreads widely (poor identifiability), the posterior over ϕ must adapt to each θ, leading to high mutual information between θ and ϕ and a more complex model. By seeking flat minima in ϕ, the mutual information term is implicitly reduced, encouraging a posterior that concentrates around a single θ value, i.e., better identification of the scientific parameters.

Empirical evaluation spans multiple domains—weather forecasting, robotics, healthcare—and includes several hybrid architectures (additive, compositional, nested). The authors compare SAM‑based training against standard regularization techniques (L2 penalties, explicit norm constraints) and demonstrate consistent improvements in both predictive accuracy and the recovery error of θ. Notably, the SAM variants (standard, Adaptive SAM, Fisher SAM) all perform favorably, indicating robustness to the choice of perturbation scaling.

In summary, the paper introduces a versatile, regularizer‑free approach for training deep hybrid models by leveraging loss‑flatness via SAM. The method requires no domain‑specific design of regularizers, works across diverse model structures, and empirically improves the identifiability of scientific parameters while maintaining or enhancing predictive performance. This contribution broadens the practical applicability of hybrid modeling in scientific and engineering contexts where preserving interpretability and physical fidelity is essential.