Sharpness-Aware Hybrid Model Learning for Architecture-Agnostic Parameter Estimation

Technical Summary: Sharpness-Aware Hybrid Model Learning for Architecture-Agnostic Parameter Estimation

1. Problem Statement

Hybrid modeling, which combines scientific mathematical models (mechanistic/physics-based) with machine learning (ML) models, offers a promising avenue for achieving flexible, robust, and partially interpretable predictions. However, a critical challenge arises when the ML component is significantly more flexible than the scientific model. In such scenarios, the ML part can absorb all data variations, effectively rendering the scientific model's unknown parameters unidentifiable. This phenomenon, often referred to as the "unidentifiability" of scientific parameters, occurs because the optimization process minimizes prediction error by adjusting only the ML parameters ($\phi$), leaving the scientific parameters ($\theta$) arbitrary.

Existing solutions typically rely on regularization terms designed to suppress the flexibility of the ML component. However, formulating these regularizers is highly architecture-dependent. For simple additive models ($y = f(x) + g(x)$), one might penalize the norm of $g$. Yet, for complex, nested compositions ($y = g(x, f(x))$ or deeper functional chains), designing a generic, effective regularizer that does not require specific domain knowledge or architectural assumptions is non-trivial.

2. Methodology

The paper proposes an architecture-agnostic method to estimate scientific parameters by leveraging the flatness of loss minima, grounded in the Occam's razor principle. The core hypothesis is that a "correct" scientific parameter $\theta$ should leave a "simple" residual task for the ML component $g_\phi$. In the context of loss landscapes, a simple model corresponds to a flat minimum, whereas a model relying on a specific, finely-tuned ML configuration to compensate for a wrong $\theta$ corresponds to a sharp minimum.

Theoretical Foundation:
The authors connect hybrid model learning to the concept of free energy (or variational inference). They define a loss objective $\bar{J}(\theta)$ that minimizes the expected loss over a distribution of ML parameters $q(\phi)$, regularized by the Kullback-Leibler (KL) divergence from a prior $p(\phi)$:
$$\bar{J}(\theta) := \min_{q} \left( \mathbb{E}_{\phi \sim q} L(\theta, \phi) + \text{KL}(q(\phi) \| p(\phi)) \right)$$
Minimizing this objective favors $\theta$ values where the loss $L(\theta, \phi)$ remains low across a wide range of $\phi$ (i.e., a flat region), rather than requiring a specific, sharp $\phi$ to achieve low loss.

Algorithmic Implementation (SAM for Hybrid Models):
Directly optimizing the free energy is computationally intractable due to the normalization constant. Instead, the authors adapt Sharpness-Aware Minimization (SAM) [15].

• Standard SAM: Minimizes the worst-case loss within a neighborhood of the parameters: $\min_w \max_{\|\epsilon\| \le \rho} L(w + \epsilon)$.
• Proposed Modification: The authors modify SAM to perturb only the ML parameters ($\phi$) while leaving the scientific parameters ($\theta$) unperturbed during the ascent step.
  • Rationale: The goal is to find a $\theta$ such that the loss is flat along the $\phi$-axis. Flatness along the $\theta$-axis is not desired, as it would imply that different $\theta$ values yield similar losses, making $\theta$ unidentifiable.
• Algorithm 1: At each training step, the algorithm computes the ascent direction $\epsilon^*_\phi$ based on the gradient of the loss with respect to $\phi$, perturbs $\phi$, computes the gradient at the perturbed point, and updates both $\theta$ and $\phi$ using this gradient.

3. Key Contributions

1. Architecture-Agnostic Parameter Estimation: The paper introduces a method that does not require designing specific regularizers for different hybrid model architectures (e.g., additive vs. nested compositions).
2. Theoretical Connection: It establishes a link between loss flatness (typically associated with generalization) and the identifiability of scientific parameters in hybrid models. It posits that flat minima in the ML parameter space correspond to scientifically meaningful parameters.
3. Modified SAM Algorithm: The authors propose a specific variant of SAM that perturbs only the ML parameters, theoretically justified by the need to enforce simplicity in the ML component's residual task.
4. Empirical Validation: The method is validated across diverse synthetic and real-world datasets, demonstrating effectiveness where traditional regularization methods fail or require complex, architecture-specific tuning.

4. Experimental Results

The authors evaluated the proposed SAM-based methods (SAM, Adaptive SAM, Fisher SAM) against baselines: Empirical Risk Minimization (ERM), Parameter Regularization ($L_2$ on $\phi$), and Functional Regularization (architecture-specific penalties).

• Tasks: Experiments included four synthetic tasks (Damped Pendulum Time-Series, Reaction-Diffusion, Duffing Oscillator, Pendulum Images) and two real-world tasks (Wind Tunnel pressure prediction, Light Tunnel image prediction).
• Performance:
  • Parameter Estimation ($\theta$-error): SAM-based methods consistently achieved lower estimation errors for scientific parameters compared to ERM. In many non-additive cases (e.g., Wind Tunnel, Light Tunnel, Pendulum Images), baseline regularization methods (p-reg, f-reg) performed poorly or required careful tuning, whereas SAM methods provided robust estimation without architecture-specific design.
  • Prediction Accuracy ($y$-error): SAM methods achieved prediction errors comparable to or slightly higher than the best baselines. The authors note that if necessary, $\phi$ can be re-trained or fine-tuned after fixing $\theta$ to recover any minor prediction loss.
  • Specific Findings:
    • In the Duffing Oscillator task (additive), functional regularization failed to estimate $\theta$ well, while SAM succeeded.
    • In non-additive tasks (e.g., Light Tunnel), the proposed method significantly outperformed baselines that struggled to define appropriate regularizers.
    • Hyperparameter Sensitivity: The method was found to be relatively robust to the perturbation radius $\rho$ within standard ranges, though strong regularization could degrade prediction performance.
    • Perturbation Scope: Experiments comparing the modified SAM (perturbing only $\phi$) against ordinary SAM (perturbing all parameters) showed that the modification generally improved or maintained $\theta$ estimation accuracy, particularly in tasks with fewer scientific parameters.

5. Significance and Claims

The paper claims to open a new pathway for architecture-agnostic parameter estimation in hybrid models. By shifting the focus from designing complex, architecture-specific regularizers to exploiting the geometric properties of the loss landscape (flatness), the method offers a versatile tool for scientific parameter identification.

Modesty and Limitations:
The authors explicitly state that their method does not guarantee the identification of the "correct" scientific parameter in all cases. As with existing methods, the "correctness" of $\theta$ depends on the underlying assumptions about the data-generating process and model selection principles. The paper acknowledges that without specific assumptions about the data distribution or model architecture, the proposed method cannot universally guarantee the recovery of ground-truth parameters. It positions the work as a practical step toward solving the unidentifiability problem in general hybrid settings, while noting that establishing a rigorous, universal theory of identification for general hybrid models remains an open challenge.