Decomposition Solver

The decomposition approach is a bi-level scheme that decouples the network-parameter update from the DAE solve. It alternates between:

1. Inner step: fix $\boldsymbol{\theta}$ and solve the per-scenario DAE NLP with a Grey-Box Model (GBM) of the network, by default through POUNCE’s grey-box interface (nlp_solver='cyipopt' selects cyipopt instead).

2. Outer step: compute $\nabla_{\boldsymbol{\theta}} \Phi$ by implicit differentiation of the inner KKT conditions and take an Adam step on $\boldsymbol{\theta}$.

This follows the bi-level formulation of Lueg et al. (2025), applied here to a physics-informed setting.

Algorithm¶

The outer problem minimizes

where $\tilde{\mathbf{x}}^{(s)}(\boldsymbol{\theta})$ is the solution of the inner DAE NLP for scenario $s$ at the current weights, and $r(\boldsymbol{\theta}) = \tfrac{1}{2}\|\boldsymbol{\theta}\|_2^2$ with weight $\alpha_r$ (param_reg_coef). Each outer iteration takes an Adam step,

Inner subproblem. For fixed $\boldsymbol{\theta}$, the network constraint is relaxed with non-negative $\ell_1$ slack variables $\boldsymbol{\Delta}^{\pm}$,

penalized in the inner objective with weight slack_coef. The relaxation keeps the inner NLP feasible even when $\mathbf{f}_{NN}$ is a poor approximation early in training.

Outer gradient. Differentiating $\Phi$ gives

The sensitivity $\nabla_{\boldsymbol{\theta}} \tilde{\mathbf{x}}^{(s)}$ comes from implicit differentiation of the inner KKT system: the linearized KKT matrix is factorized once per outer iteration with FERAL (a pure-Rust sparse symmetric-indefinite LDL$^{\top}$ solver), and the network terms enter through vector-Jacobian products evaluated by automatic differentiation. FERAL is the default; pass linear_solver='ma27' or linear_solver='scipy' to train_decomp to select an alternative (see Solvers).

The inner NLP is solved with POUNCE by default (nlp_solver='pounce'; 'cyipopt' or 'ipopt' are selectable). For POUNCE the decomposition defaults mu_strategy='adaptive', which converges robustly on the cold inner solves; pass your own solver_options to override it.

Usage¶

The high-level wrapper drives this method with HybridDAE(method="decomposition", ...); the stage-level calls are:

from sindae.algorithms.pretrain import PretrainConfig, pretrain_mlp
from sindae.algorithms.decomp.train import DecompConfig, train_decomp

# (Optional) supervised pretrain on smoother output
mlp = pretrain_mlp(mlp, smoother_data, PretrainConfig(epochs=400, reg_coef=0.1))

cfg = DecompConfig(
    n_steps=300,
    lr=5e-3,
    init_slack_coef=10.0,       # initial L1 slack penalty
    slack_scale=2.0,             # multiply slack_coef every slack_update_interval steps
    slack_update_interval=50,
    max_slack_coef=1e3,
    param_reg_coef=1e-3,
    patience=20,                 # early stopping (set 0 to disable)
)

trained_m, mlp, history = train_decomp(
    problem=problem,
    mlp=mlp,
    cfg=cfg,
    data=smoother_data,           # normalization statistics
    smoother_model=smoother_m,    # warm-start the inner NLP
    solver_options={'tol': 1e-6, 'max_iter': 300},
)

# trained_m is the solved NLP; recover trajectories with:
trained_data = extract_instance_data(problem, trained_m)

history is a dict with keys: obj_history, data_fit_history, grad_norm_history, diag_history, pouncetiming_history.

DecompConfig Reference¶

MPI Parallelism¶

For multi-trajectory training, trajectories are distributed across MPI ranks. Each rank maintains its own TrajectoryBatchSubproblem and gradients are All-Reduced before the Adam step:

from mpi4py import MPI
comm = MPI.COMM_WORLD

trained_m, mlp, history = train_decomp(problem, mlp, cfg, data, mpi_comm=comm)

Under MPI, trained_m is the rank-local model (holding only that rank’s trajectories); the returned mlp is identical across ranks (rank-0 authoritative).

Run with: mpirun -n 4 python train.py

When to Use¶

• Large networks where the full Hessian is too expensive.

• Many independent training trajectories (MPI parallelism).

• When you want fine-grained control over the training schedule.

For small-to-medium networks with few trajectories, the Simultaneous Solver is simpler and often faster.

API Reference¶

See Decomposition Solver for DecompConfig, train_decomp, and build_decomp_model. See Pre-training for PretrainConfig and pretrain_mlp.