4DVarNet — Learned 4DVar

4DVarNet Fablet et al., 2021 keeps the shape of classical 4DVar — minimise a variational cost that balances a data term against a prior — but makes two of its components learnable:

1. the Gaussian background $\tfrac{1}{2}\|\mathbf{x}-\mathbf{x}_b\|^2_{\mathbf{B}^{-1}}$ becomes a learned prior $\boldsymbol{\phi}_{\boldsymbol{\theta}}$ (an autoencoder), able to capture nonlinear, multimodal manifold structure a Gaussian cannot;

2. the inner gradient-descent solver becomes a learned iterative update $\boldsymbol{\Phi}_{\boldsymbol{\psi}}$ (a ConvLSTM / MLP / attention network).

The whole pipeline is differentiable end-to-end and is trained on $(\text{observation}, \text{target})$ pairs. The “Net” is a method name (Fablet et al.), not a separate paradigm — with the right (trivial) choices it reduces to classical 4DVar and to OI.

The notation follows the DA notes: state $\mathbf{x}$, observations $\mathbf{y}$ (here on a subdomain / mask $\boldsymbol{\Omega}$), observation operator $H$. The learnable parameters are collected as $\boldsymbol{\Theta} = (\boldsymbol{\theta}, \boldsymbol{\psi})$ — prior parameters $\boldsymbol{\theta}$ and solver parameters $\boldsymbol{\psi}$.

Setup¶

We observe a noisy, gappy field — the truth corrupted by noise and seen only on a subdomain $\boldsymbol{\Omega}$ (the observed pixels):

The goal is a state $\mathbf{x}$ that matches the observations on $\boldsymbol{\Omega}$ and fills the gaps on the complement $\bar{\boldsymbol{\Omega}}$. Classical 4DVar fills gaps using a Gaussian background; 4DVarNet fills them using a learned prior.

The Variational Cost¶

The inner cost balances an observation term (restricted to $\boldsymbol{\Omega}$) against a learned-prior term:

where $\|\cdot\|_{\boldsymbol{\Omega}}^2$ is the quadratic norm evaluated only on the observed subdomain. The prior term measures how far $\mathbf{x}$ is from its own learned reconstruction $\boldsymbol{\phi}_{\boldsymbol{\theta}}(\mathbf{x})$: it is small on the manifold of plausible states and large off it. This replaces — not complements — the Gaussian background.

State vs. Parameters: a Bi-Level Problem¶

This is the conceptual heart of 4DVarNet, and where it differs from classical 4DVar. There are two distinct kinds of unknown, optimised in two nested loops:

Inner problem — state estimation. For fixed parameters $\boldsymbol{\Theta}$ and a given observation $\mathbf{y}$, run the solver to produce the analysis

Outer problem — parameter learning. Choose the parameters that make the inner solver’s output match the truth, averaged over the dataset:

The Learned Inner Solver¶

Classical 4DVar takes plain gradient steps on the cost, $\mathbf{x}^{(k+1)} = \mathbf{x}^{(k)} - \lambda \nabla_{\mathbf{x}} J_{\boldsymbol{\theta}}(\mathbf{x}^{(k)})$. 4DVarNet replaces the fixed step with a learned update that reads the cost gradient and a recurrent hidden state $\mathbf{h}^{(k)}$:

$\boldsymbol{\Phi}_{\boldsymbol{\psi}}$ (a ConvLSTM, MLP, or attention network) learns adaptive step sizes, momentum, and problem-specific preconditioning. After each step the observed pixels are reset to the data and only the gaps are updated — the $\boldsymbol{\Omega}$-projection:

Three flavours of inner update, from hand-built to fully learned:

import einx
from jaxtyping import Array, Bool, Float

def project_step(
    x:    Float[Array, "... D"],
    y:    Float[Array, "... D"],
    mask: Float[Array, "... D"],     # Ω indicator (1 = observed)
    phi,                              # learned reconstruction φ_θ
) -> Float[Array, "... D"]:
    x_hat = phi(x)
    return einx.multiply("... D, ... D -> ... D", mask, y) \
         + einx.multiply("... D, ... D -> ... D", 1.0 - mask, x_hat)

import jax
from jaxtyping import Array, Float

def energy(x, y, mask, theta, *, a_obs=0.99, a_prior=0.01) -> Float[Array, ""]:
    loss_obs   = jax.numpy.mean(mask * (x - y) ** 2)         # ‖H(x)-y‖²_Ω  (H = identity here)
    loss_prior = jax.numpy.mean((phi(x, theta) - x) ** 2)    # ‖x - φ_θ(x)‖²
    return a_obs * loss_obs + a_prior * loss_prior

def gradient_step(x, y, mask, theta, alpha: float) -> Float[Array, "... D"]:
    return x - alpha * jax.grad(energy)(x, y, mask, theta)

import jax
from jaxtyping import Array, Float, PyTree

def learned_step(
    x:    Float[Array, "... D"],
    y:    Float[Array, "... D"],
    mask: Float[Array, "... D"],
    theta: PyTree,                 # prior parameters
    grad_mod,                      # solver network Φ_ψ (carries its own state)
    h:    PyTree,                  # recurrent hidden state h_k
    alpha: float,
) -> tuple[Float[Array, "... D"], PyTree]:
    g       = alpha * jax.grad(energy)(x, y, mask, theta)   # ∇ₓ J
    update, h = grad_mod(g, h)                              # learned modulation
    return x - update, h

Iterating (5) for $K$ steps and applying (6) at each step is the inner solver $\mathbf{x}^\star(\boldsymbol{\Theta}; \mathbf{y})$ of (3).

Why We Need optax¶

The inner loop takes “gradient steps,” but those steps optimise the state $\mathbf{x}$ for one example — they are part of the forward pass. The thing we actually train is the parameters $\boldsymbol{\Theta}$, in the outer loop (4), and that is an ordinary supervised deep-learning problem: minimise a loss over a dataset with respect to network weights. That is precisely what a gradient-based optimiser library like optax is for — Adam/SGD with momentum, schedules, clipping, and optimiser state — so we do not hand-roll it.

The outer gradient $\nabla_{\boldsymbol{\Theta}}\mathcal{L}$ flows through the $K$ inner solver iterations (see adjoint choices below); optax then consumes it to update the parameters:

import jax
import jax.numpy as jnp
import optax

optimizer = optax.adam(1e-3)
opt_state = optimizer.init(params)            # params = Θ = (θ, ψ)

@jax.jit
def train_step(params, opt_state, y, mask, x_gt):
    def loss_fn(params):
        x_star = inner_solve(params, y, mask)               # K learned steps -> x*(Θ; y)
        return jnp.mean((x_star - x_gt) ** 2)               # outer reconstruction loss
    loss, grads = jax.value_and_grad(loss_fn)(params)       # ∇_Θ L, through the inner solve
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

Training¶

The reconstruction loss (4) comes in two flavours, mirroring the data you have:

Differentiating Through the Inner Solver¶

Because $\mathbf{x}^\star$ is produced by $K$ iterations, the outer gradient must backpropagate through them. Three strategies trade memory against exactness:

The implicit strategy differentiates the fixed point $\mathbf{x}^\star = \boldsymbol{F}(\mathbf{x}^\star, \boldsymbol{\Theta})$ via the implicit function theorem, avoiding storage of the iterates entirely:

The Prior and Solver Families¶

The two learnable pieces are swappable. The prior $\boldsymbol{\phi}_{\boldsymbol{\theta}}$:

The gradient modulator $\boldsymbol{\Phi}_{\boldsymbol{\psi}}$:

When to Use 4DVarNet¶

Posterior¶

A Laplace approximation works as in the classical methods, but the Hessian at the MAP uses the learned-prior curvature, which may be poorly conditioned outside the training manifold. Treat 4DVarNet posteriors with caution and validate them; when full uncertainty matters, an ensemble method is safer.

Perspectives: Explicit vs. Implicit¶

The learned map can be framed three ways, which clarifies what gets penalised:

• Explicit — a function predicts the truth, $\mathbf{x}_{\text{gt}} = \boldsymbol{f}(\mathbf{x}_{\text{init}})$, with a model term added to the loss, $\mathcal{L} = \mathcal{L}_{\text{model}} + \mathcal{L}_{\text{data}}$.

• Conditional explicit — the map reads observations too, $\mathbf{x}_{\text{gt}} = \boldsymbol{f}(\mathbf{x}_{\text{init}}, \mathbf{y})$; no extra loss term is needed.

• Implicit (fixed point) — the truth is the solution of $\boldsymbol{g}(\mathbf{x}, \mathbf{y}) = \boldsymbol{f}(\mathbf{x}, \mathbf{y}) - \mathbf{x} = \boldsymbol{0}$, i.e. $\mathbf{x}_{\text{gt}} = \boldsymbol{f}(\mathbf{x}_{\text{gt}}, \mathbf{y})$. This is the iterative solver view, and the regime where implicit differentiation shines.

A probabilistic head replaces the squared error with a Gaussian negative log-likelihood, predicting a mean and variance, $p(\mathbf{x}^{(k+1)} \mid \mathbf{x}^{(k)}) = \mathcal{N}\big(\mathbf{x}^{(k)} \mid \boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}^{(k)}), \boldsymbol{\sigma}^2_{\boldsymbol{\theta}}(\mathbf{x}^{(k)})\big)$, giving calibrated per-pixel uncertainty when trained with the NLL loss.

History¶

4DVarNet originated with Fablet et al. Fablet et al., 2021Fablet et al., 2023 for satellite sea-surface-height reconstruction; the original implementation was ocean-specific (PyTorch). The formulation here is framework-agnostic (jaxtyping + JAX autodiff + optax) and positions 4DVarNet as one method among many in the DA hierarchy — a learnable generalisation of, not a replacement for, classical variational assimilation.