ML surrogate training training/std

ML surrogate training training/std#

Description#

Train a neural-network surrogate for the CIF observation operator.

This mode wraps the CIF observation operator \(\mathcal{H}\) as a differentiable function, trains a neural-network surrogate against it, exports the surrogate to ONNX, and returns a zero-ML-dependency OnnxSurrogateOperator that is a drop-in replacement for the obsoperator in any subsequent CIF mode.

Pipeline#

  1. Wrapping — the CIF obsoperator is wrapped as a plain-numpy object exposing forward(x) -> y and adjoint(x, dy) -> dx.

  2. Architecture — a fully-connected multi-layer perceptron (MLP) is built automatically from the state and observation dimensions (read from controlvect.dim and obsvect.dim) with n_hidden_layers hidden layers of width hidden_size.

  3. Sampling — training inputs are drawn from the prior:

    \[\mathbf{x} \sim \mathcal{N}(\mathbf{x}_b,\, \boldsymbol{\sigma}_b)\]

    where \(\mathbf{x}_b\) is the background state and \(\boldsymbol{\sigma}_b\) is the background standard deviation stored in the control vector.

  4. Training — the surrogate \(\hat{\mathcal{H}}_\theta\) is trained online against the live CIF operator (Mode 2):

    \[\mathcal{L}_\mathrm{fwd}(\theta) = \bigl\|\hat{\mathcal{H}}_\theta(\mathbf{x}) - \mathcal{H}(\mathbf{x})\bigr\|^2\]

    When adjoint_consistency = True (Mode 3), a VJP-agreement term is added to enforce that the surrogate Jacobian matches the CIF adjoint:

    \[\mathcal{L}(\theta) = \mathcal{L}_\mathrm{fwd}(\theta) + \lambda_\mathrm{adj} \bigl\|\mathbf{J}_\theta(\mathbf{x})^\top \mathbf{v} - \mathbf{H}^\top \mathbf{v}\bigr\|^2\]

    where \(\mathbf{v}\) is a random test vector and \(\lambda_\mathrm{adj}\) is set by lambda_adj. This term guarantees that gradient-based DA solvers (e.g. 4D-Var) remain consistent when switching from the full to the surrogate operator.

  5. Export — the trained model is serialised to ONNX (file onnx_file inside workdir) via torch.onnx.export with opset 17 and a dynamic batch axis.

  6. Reload — the ONNX file is reloaded as an OnnxSurrogateOperator backed by onnxruntime. This object has zero dependency on PyTorch or JAX and exposes the same forward(x) interface as the original operator.

Output#

The mode returns an OnnxSurrogateOperator instance that can replace the obsoperator in any subsequent CIF mode, for example to accelerate a variational inversion:

surrogate = setup.mode.execute()        # runs the full pipeline
surrogate.forward(x)                    # fast inference, no ML deps

Dependencies#

  • torch (pip install torch) — required for the default "torch" backend and ONNX export.

  • onnxruntime (pip install onnxruntime) — required at reload time (the only ML dependency needed at inference).

  • jax, optax (pip install jax optax) — required only when backend = "jax".

YAML arguments#

The following arguments are used to configure the plugin. pyCIF will return an exception at the initialization if mandatory arguments are not specified, or if any argument does not fit accepted values or type:

run_mode : “fwd” or “tl”, optional, default “fwd”

CIF mode used to evaluate the observation operator during training. Use 'fwd' to train against the full nonlinear operator, or 'tl' to train against its linearisation.

  • “fwd”: full nonlinear forward operator

  • “tl”: tangent-linear operator (linearised around xb)

reload_results : bool, optional, default False

If True, skip CIF simulations that have already been computed and reload their results from disk. Useful to resume a training run interrupted mid-way.

backend : “torch” or “jax”, optional, default “torch”

ML framework used for the training loop and gradient computation. 'torch' also handles ONNX export directly; 'jax' requires the jax2tf + tf2onnx pipeline for export (see export_onnx()).

  • “torch”: PyTorch — pip install torch

  • “jax”: JAX + optax — pip install jax optax

n_steps : int, optional, default 5000

Number of gradient-descent steps for surrogate training.

adjoint_consistency : bool, optional, default False

Enable Mode 3 training: adds a VJP-agreement loss term \(\|\mathbf{J}_\theta^\top \mathbf{v} - \mathbf{H}^\top \mathbf{v}\|^2\) alongside the forward MSE loss. Recommended when the surrogate will be used inside a gradient-based DA solver (e.g. 4D-Var).

lambda_adj : float, optional, default 0.1

Weight \(\lambda_\mathrm{adj}\) of the adjoint-consistency loss term. Only used when adjoint_consistency = True. Increase if the surrogate Jacobian deviates noticeably from the CIF adjoint.

hidden_size : int, optional, default 256

Number of neurons in each hidden layer of the default MLP surrogate. Increase for more expressive surrogates at the cost of training time and inference latency.

n_hidden_layers : int, optional, default 2

Number of hidden layers in the default MLP (multi-layer perceptron) surrogate. The total depth is n_hidden_layers + 1 (including the output layer). Two hidden layers is a good starting point for most atmospheric operators.

onnx_file : str, optional, default “surrogate.onnx”

Filename of the exported ONNX surrogate, written inside workdir. The file is reloaded immediately after export as an OnnxSurrogateOperator (zero ML dependencies at inference time).

Requirements#

The current plugin requires the present plugins to run properly:

Requirement name

Requirement type

Explicit definition

Any valid

Default name

Default version

controlvect

ControlVect

True

True

standard

std

obsvect

ObsVect

True

True

standard

std

obsoperator

ObsOperator

True

False

standard

std

YAML template#

Please find below a template for a YAML configuration:

 1mode:
 2  plugin:
 3    name: training
 4    version: std
 5    type: mode
 6
 7  # Optional arguments
 8  run_mode: XXXXX  # fwd|tl
 9  reload_results: XXXXX  # bool
10  backend: XXXXX  # torch|jax
11  n_steps: XXXXX  # int
12  adjoint_consistency: XXXXX  # bool
13  lambda_adj: XXXXX  # float
14  hidden_size: XXXXX  # int
15  n_hidden_layers: XXXXX  # int
16  onnx_file: XXXXX  # str

See also

API reference (functions & code)