# ott.initializers.nn.initializers.MetaInitializer.update#

MetaInitializer.update(state, a, b)[source]#

Update the meta model with the dual objective.

The goal is for the model to match the optimal duals, i.e., $$\hat f_\theta \approx f^\star$$. This can be done by training the predictions of $$\hat f_\theta$$ to optimize the dual objective, which $$f^\star$$ also optimizes for. The overall learning setup can thus be written as:

$\min_\theta\; {\mathbb E}_{(\alpha,\beta)\sim{\mathcal{D}}}\; J(\hat f_\theta(a, b); \alpha, \beta),$

where $$a,b$$ are the probabilities of the measures $$\alpha,\beta$$, $$\mathcal{D}$$ is a meta distribution of optimal transport problems,

$-J(f; \alpha, \beta, c) := \langle f, a\rangle + \langle g, b \rangle - \varepsilon\left\langle \exp\{f/\varepsilon\}, K\exp\{g/\varepsilon\}\right\rangle$

is the entropic dual objective, and $$K_{i,j} := -C_{i,j}/\varepsilon$$ is the Gibbs kernel.

Parameters
• state (TrainState) – Optimizer state of the meta model.

• a (Array) – Probabilites of the $$\alpha$$ measure’s atoms.

• b (Array) – Probabilites of the $$\beta$$ measure’s atoms.

Return type

Tuple[Array, Array, TrainState]

Returns

The training loss, $$f$$, and updated state.