Bi-Level optimization for interpolation-based parameter estimation of differential equations

1. Problem Statement

The paper addresses the inverse problem of parameter estimation for Ordinary Differential Equations (ODEs) and Delay Differential Equations (DDEs). The goal is to estimate unknown model parameters ($p$) by minimizing the mismatch between model-predicted states and experimental measurements.

Key Challenges:

• Computational Cost: Traditional methods like Sequential Optimization (Single-Shooting) require numerical integration of ODEs in every iteration. Calculating gradients (sensitivities) for the optimizer involves differentiating through the numerical integrator, which is computationally expensive and prone to instability, especially for stiff or chaotic systems.
• Convergence Issues: Single-shooting methods often fail to converge for long trajectories, poor initial guesses, or highly nonlinear dynamics.
• Limitations of Existing Interpolation: Previous interpolation-based methods were largely restricted to systems where parameters appeared linearly or required user-provided initial state estimates close to the truth.
• Model Discovery: Estimating both the structure of the differential equations and their parameters simultaneously is even more challenging due to the non-convex nature of the search space.

2. Methodology

The authors propose a Bi-Level Optimization Framework that leverages interpolation to decouple the state estimation from the parameter optimization, reducing computational complexity and improving stability.

Core Concept: Bi-Level Structure

The problem is formulated as a nested optimization:

1. Inner Problem (Convex): Optimizes linear parameters ($p$) given a fixed set of nonlinear parameters ($\phi$).
   • The state $X(t)$ is approximated using an interpolation function $\Psi(X)(t)$ derived from measured data.
   • The ODE residual is minimized. Because the parameters $p$ appear linearly in the ODE function $\Theta$, and the interpolation is fixed for a given $\phi$, this sub-problem is convex.
   • Constraints (equality/inequality) are assumed to be affine/convex with respect to $p$.
2. Outer Problem (Non-Convex): Optimizes the nonlinear parameters ($\phi$).
   • The objective is the loss function evaluated using the optimal solution ($p^*$) from the inner problem.

Gradient Calculation via Implicit Function Theorem

To solve the outer problem, gradients of the optimal inner solution ($p^*$) with respect to the nonlinear parameters ($\phi$) are required.

• The authors use the Implicit Function Theorem applied to the Karush-Kuhn-Tucker (KKT) conditions of the inner convex problem.
• This allows for the derivation of forward-mode sensitivities ($\frac{dp^*}{d\phi}$) without differentiating through the entire numerical integration history.
• The KKT matrix factorization from the inner solver is reused to accelerate derivative calculations.

Handling Sensitivities and Interpolation

• Interpolation: Instead of integrating the ODE to find states, the method uses interpolation (e.g., splines) on the measured data to approximate states. This avoids the instability of integrating ODEs with incorrect parameters.
• Sensitivity Equations:
  • Standard approach: $\frac{dS}{dt} = \frac{\partial \Theta}{\partial X}S + \dots$ (Expensive).
  • Proposed approach: Since $X$ is replaced by $\Psi(X)$, the sensitivity equation simplifies to $\frac{dS}{dt} = \frac{\partial \Theta}{\partial p^*}\frac{dp^*}{d\phi} + \frac{\partial \Theta}{\partial \phi}$. This removes the need to compute $\frac{\partial \Theta}{\partial X}$, significantly reducing cost and increasing stability.

Implementation Details

• Tools: JAX (for automatic differentiation), IPOPT (nonlinear solver), SciPy (interpolation), Pyomo (collocation).
• Model Discovery Extension: For unknown model structures, the method uses a Sequential Thresholding algorithm (Algorithm 1). It iteratively solves the bi-level problem, applies a threshold to eliminate small coefficients (spurious terms) from a library of basis functions, and refines the model structure.

3. Key Contributions

1. Generalization of Interpolation: Extends interpolation-based parameter estimation from linearly separable parameters to a general class of ODEs where parameters can appear non-linearly, handled via a bi-level framework.
2. Efficient Gradient Computation: Derives a method to compute gradients through the optimal solution of a convex inner problem using the Implicit Function Theorem, avoiding the high cost of sensitivity integration.
3. Robustness to Stiffness and Delays: Demonstrates that the method handles stiff systems (where explicit solvers fail in single-shooting) and delayed differential equations (DDEs) without requiring specialized stiff solvers during the optimization loop.
4. Unified Framework for Discovery: Successfully applies the framework to model discovery, simultaneously identifying the functional form of the ODE and its parameters.
5. Handling Partial Observability: Extends the method to systems with partially observed states (latent variables) by treating unobserved state trajectories as decision variables interpolated via splines.

4. Results

The method was tested on seven benchmark problems spanning diverse dynamics:

• Systems: Calcium ion dynamics, Mendes biochemical model, Ethanol fermentation, Kermack-McKendrick (DDE), Belousov reaction (Stiff ODE), CSTR (Partially observed), and Esterification of carboxylic acid (Model Discovery).
• Comparison: Compared against Single-Shooting and Orthogonal Collocation.

Key Findings:

• Convergence: The proposed bi-level method converged to the global optimum for all benchmark problems. In contrast, Single-Shooting and Orthogonal Collocation frequently failed to converge or got stuck in local minima.
• Accuracy: The Mean Squared Error (MSE) for the proposed method was orders of magnitude lower than the baselines (e.g., $1.42 \times 10^{-2}$ vs $3.67 \times 10^7$ for the Belousov reaction).
• Iterations: The method required significantly fewer optimization iterations than Single-Shooting due to the convexity of the inner problem.
• Wall Time: While the method required more total wall-clock time per run (due to solving the inner convex problem and KKT factorization), it was the only method to successfully solve the stiff and highly nonlinear cases.
• Stiffness Handling: For the Belousov reaction (stiff), Single-Shooting failed completely because the explicit solver could not handle the stiffness with poor initial guesses. The interpolation-based method remained stable because it queried the data directly rather than integrating forward from a potentially divergent guess.

5. Significance and Limitations

Significance:

• This work bridges the gap between data-driven interpolation and rigorous optimization, offering a robust alternative to traditional shooting methods for complex chemical engineering and biological systems.
• It enables the estimation of parameters in stiff and delayed systems where traditional gradient-based methods often fail.
• It provides a scalable path for dynamic model discovery, allowing researchers to learn both the equations and parameters from data without prior structural knowledge.

Limitations:

• Computational Cost: The method is computationally heavier per iteration than single-shooting due to the nested optimization and KKT matrix operations.
• Data Requirements: The interpolation approach relies on sufficiently dense and frequent sampling. Sparse data can lead to inaccurate integral approximations.
• Noise Sensitivity: Interpolation over noisy measurements can lead to inaccurate derivatives. The authors suggest smoothing techniques as a mitigation strategy.
• Initial Guess: While more robust than single-shooting, the outer non-convex problem still requires reasonable initial guesses for nonlinear parameters to ensure convergence to the global optimum.

In conclusion, the paper presents a mathematically rigorous and practically effective framework for parameter estimation that overcomes the convergence and stability limitations of traditional methods, particularly in the context of stiff, delayed, and partially observed systems.