Technical Summary: Differentiable Parameter Optimization for DAEs with State-Dependent Events

Problem Formulation

The paper addresses the challenge of performing gradient-based parameter learning for semi-explicit Differential-Algebraic Equations (DAEs) that include state-dependent events. Unlike standard Ordinary Differential Equations (ODEs), DAEs involve algebraic constraints ($0 = g(t, x, z, p)$) that implicitly define algebraic variables ($z$) alongside differential states ($x$). In hybrid systems, these continuous dynamics are interrupted by discrete events triggered when guard functions ($\phi_e(t, x, z, p) = 0$) are satisfied. These events induce mode changes, impacts, or state reinitializations via reset maps ($\Psi$).

The core difficulty in optimizing parameters ($p$) for such systems is that the mapping from parameters to outputs is not a smooth function. It is complicated by:

1. Implicit Algebraic Variables: $z$ must be solved at every integration step and after every reset.
2. Parameter-Dependent Event Times: Changing $p$ shifts the instant a guard condition is met, altering the integration intervals.
3. Discontinuities: Reset maps introduce jumps in the state trajectory.

The learning problem is formulated as a constrained least-squares problem: minimizing the mismatch between model outputs and observed data, subject to the DAE dynamics, algebraic constraints, event guard conditions, and reset equations.

Methodology

The authors propose and compare two complementary strategies for computing gradients in this setting:

1. AD-Through-Simulation (Automatic Differentiation)

This approach treats the event-driven DAE simulator as a differentiable program.

• Algebraic Differentiation: Instead of differentiating through every step of the Newton iteration used to solve $g(t, x, z, p) = 0$, the method uses the Implicit Function Theorem. It computes a custom Jacobian-Vector Product (JVP) that propagates gradients through the algebraic solve directly.
• Event Handling: The simulation is implemented as segmented integration. The solver integrates between events, detects the first zero-crossing of a composite guard function, applies the reset map, and restarts integration.
• Implementation: Built using JAX and the diffrax library. It employs a fixed-length scan over a maximum number of segments to ensure compatibility with Just-In-Time (JIT) compilation. Target-time evaluation uses right-continuous selection (assigning an event time to the post-event segment).
• Nature: This method differentiates the specific numerical algorithm executed by the simulator, including solver tolerances and root-finding steps.

2. Explicit Discrete-Adjoint Method

This approach avoids differentiating through the solver's internal operations. Instead, it constructs a discrete residual system based on the forward trajectory.

• Formulation: The forward simulation is treated as a set of equality constraints. Smooth segments are constrained by trapezoidal DAE residuals, while event blocks enforce guard conditions, reset equations, and continuity.
• Lagrange Multipliers: The residuals are treated strictly as equality constraints, not heuristic penalty terms. The method solves for the Lagrange multipliers (adjoint variables) associated with these constraints.
• Backward Sweep: Gradients are computed via a reverse sweep.
  • Smooth Segments: Solve local adjoint systems for the trapezoidal residuals.
  • Event Blocks: Solve event-adjoint equations. Crucially, because event times define the boundaries of adjacent segments, the method resolves an affine pending-event adjoint. It computes a scalar coefficient that satisfies the stationarity condition with respect to the event time, coupling the event block with the preceding segment.
• Nature: This method provides explicit control over the gradient calculation and does not require the forward solver to be differentiable, allowing the use of mature external solvers (e.g., SUNDIALS/IDA).

Key Contributions

1. Unified Formulation: The paper formulates parameter learning for hybrid DAEs as a constrained optimization problem, explicitly distinguishing between algebraic constraints, event guards, and reset maps.
2. AD-Through-Simulation Implementation: It presents a practical implementation using JAX/diffrax that differentiates through algebraic solves via the implicit function theorem and handles events through segmented integration.
3. Explicit Discrete-Adjoint Derivation: It derives a discrete-adjoint formulation where residuals are equality constraints. It details the mechanism for handling event-time stationarity via an affine representation of multipliers and a backward sweep.
4. Comparative Analysis: The paper provides a detailed comparison of the two methods regarding gradient interpretation, event-time handling, implementation complexity, and local validity.

Experimental Results

The authors evaluated both methods on two benchmarks:

1. Cauer Ladder Circuit: A 7-parameter electrical circuit with a state-dependent reinitialization event.
   • Results: The AD-through-simulation method achieved a lower final loss ($1.69 \times 10^{-6}$) compared to the discrete-adjoint method ($4.00 \times 10^{-4}$). The AD method recovered parameters with higher accuracy, though both methods struggled with weakly identifiable parameters (e.g., $C_5$).
   • Observation: The discrete-adjoint method's accuracy was limited by the density of the adjoint mesh (trapezoidal steps) rather than the forward solver tolerance.
2. Multi-Ball Bouncing System: A planar system of $N$ balls ($N \in \{3, 7, 15\}$) with elastic and inelastic collisions.
   • Scaling: The JAX-based AD-through-simulation pipeline demonstrated superior scaling with $N$ (flat per-iteration cost) compared to a PyTorch baseline, which scaled steeply due to graph rebuilding overhead.
   • Accuracy vs. Cost: For small $N$, the JAX-AD method dominated in both speed and accuracy. For large $N$ ($N=15$), the PyTorch baseline achieved slightly higher accuracy but at a significantly higher computational cost (37x). The discrete-adjoint method showed increasing gradient bias as $N$ grew, attributed to the fixed-shape JIT regime and uniform target grid.

Significance and Claims

The paper claims to provide a unified treatment of differentiable parameter learning for DAEs with state-dependent events, bridging the gap between differentiable ODE solvers and classical DAE sensitivity analysis.

• Local Validity: The authors explicitly state that both methods compute gradients that are local to the event path selected by the forward simulation. They rely on assumptions of transversal guard crossings and fixed event ordering under small parameter perturbations. They acknowledge that gradients may be invalid or nonsmooth near grazing events, simultaneous events, or when the event sequence changes.
• Residual Interpretation: A key theoretical clarification is that in the discrete-adjoint method, residuals are equality constraints, not penalty terms. This distinguishes the approach from heuristic penalty-based optimization.
• Implementation Trade-offs:
  • AD-through-Simulation: Offers implementation simplicity and tight coupling with modern ML frameworks (JAX/PyTorch) but requires retaining the computational trace of the simulation, which can be memory-intensive.
  • Discrete-Adjoint: Offers control over the gradient calculation and compatibility with external solvers but requires explicit construction of the residual system and solving for multipliers.

The authors position their work as a foundation for differentiable hybrid simulation, noting that future work should address robustness near nonsmooth transitions and the integration of advanced DAE compilation techniques (e.g., tearing) to reduce system size.