CNP-Flow adds a conditional noise predictor that maps a temporal condition
\(\mathcal{C}_\mathcal{T}\) to an isotropic Gaussian source distribution:
\[
p_\phi(\boldsymbol{x}_0 \mid \mathcal{C}_\mathcal{T}) :=
\mathcal{N}\!\left(\boldsymbol{x}_0;\; \boldsymbol{\mu}_\phi(\mathcal{C}_\mathcal{T}),\; \sigma^2_\phi(\mathcal{C}_\mathcal{T})\,\mathbf{I}\right)
\]
Stage 1: Pretrain the FM, then reverse it
The FM is pretrained in the standard way with a \(\mathcal{N}(\mathbf{0}, \mathbf{I})\) source.
Because the learned ODE flow is continuous, flow reversal maps nearby targets to
nearby regions of source space. For each target \(\boldsymbol{x}_1 \sim p_1(\cdot \mid \mathcal{C}_\mathcal{T})\)
we retrieve its source sample by integrating backward:
\[
\tilde{\boldsymbol{x}}_0 = \tilde{\Phi}_{\theta, 1 \to 0}(\boldsymbol{x}_1 \mid \mathcal{C}_\mathcal{T})
:= \boldsymbol{x}_1 - \int_0^1 \boldsymbol{v}_\theta(\boldsymbol{x}_\tau, \tau \mid \mathcal{C}_\mathcal{T})\,\mathrm{d}\tau
\]
This induces a per-condition source distribution \(\tilde{p}_0(\cdot \mid \mathcal{C}_\mathcal{T})\)
that is clustered by temporal condition rather than spread as a standard Gaussian. It serves as
the supervision signal for the CNP, sidestepping the intractable search for an optimal source.
Stage 2: Train the CNP
The CNP is trained to maximize the likelihood of the retrieved source samples, with a KL term
that keeps the predicted distribution anchored to the standard Gaussian prior:
\[
\mathcal{L}_{\mathrm{CNP}}(\phi) =
\underbrace{-\,\mathbb{E}_{\mathcal{C}_\mathcal{T},\, \tilde{\boldsymbol{x}}_0 \sim \tilde{p}_0(\cdot \mid \mathcal{C}_\mathcal{T})}
\big[\log p_\phi(\tilde{\boldsymbol{x}}_0 \mid \mathcal{C}_\mathcal{T})\big]}_{\mathcal{L}_{\mathrm{NLL}}}
\;+\; \lambda\,
\underbrace{\mathcal{D}_{\mathrm{KL}}\!\big(p_\phi(\cdot \mid \mathcal{C}_\mathcal{T}) \,\|\, \mathcal{N}(\mathbf{0}, \mathbf{I})\big)}_{\mathcal{L}_{\mathrm{KL}}}
\]
Proposition (Population minimizer of the NLL objective)
Over the isotropic Gaussian family, the minimizer of \(\mathcal{L}_{\mathrm{NLL}}\) has
\[
\tilde{\boldsymbol{\mu}}_0(\mathcal{C}_\mathcal{T}) = \mathbb{E}_{\boldsymbol{x}_1}\big[\tilde{\Phi}_{\theta,1\to0}(\boldsymbol{x}_1 \mid \mathcal{C}_\mathcal{T})\big],
\qquad
\tilde{\sigma}_0^2(\mathcal{C}_\mathcal{T}) = \tfrac{1}{d}\,\mathbb{E}_{\boldsymbol{x}_1}
\big\|\tilde{\Phi}_{\theta,1\to0}(\boldsymbol{x}_1 \mid \mathcal{C}_\mathcal{T}) - \tilde{\boldsymbol{\mu}}_0\big\|_2^2
\]
The CNP's optimal mean and variance coincide with the mean and scaled second central moment of
the flow-reversal-induced source distribution.
Theorem (Conditional source locality, informal)
If the predicted mean \(\boldsymbol{\mu}^*_\phi(\cdot)\) is \(L_\mu\)-Lipschitz in the temporal
condition, then for two conditions \(\mathcal{C}_\mathcal{T}, \mathcal{C}'_\mathcal{T}\),
independently drawn CNP source samples satisfy with high probability
\[
\|\boldsymbol{x}_0 - \boldsymbol{x}'_0\|_2 \le L_\mu \|\mathcal{C}_\mathcal{T} - \mathcal{C}'_\mathcal{T}\|_2
+ \mathcal{O}\!\left(\sqrt{d\big(\sigma^{*2}_\phi(\mathcal{C}_\mathcal{T}) + \sigma^{*2}_\phi(\mathcal{C}'_\mathcal{T})\big)}\right)
\]
Nearby temporal conditions yield nearby source centers; the predicted variances control the
stochastic dispersion around them.
Stage 3: Fine-tune the FM with the learned source
Finally, the FM is fine-tuned with the same objective as Stage 1, but with source samples drawn
from the CNP instead of the standard Gaussian:
\[
\mathcal{L}_{\mathrm{FT\text{-}FM}}(\theta) =
\mathbb{E}_{\mathcal{C}_\mathcal{T},\, \boldsymbol{x}_0 \sim p_\phi(\cdot \mid \mathcal{C}_\mathcal{T}),\, \boldsymbol{x}_1 \sim p_1(\cdot \mid \mathcal{C}_\mathcal{T}),\, t}
\left[\big\|\boldsymbol{v}_\theta(\boldsymbol{x}_t, t \mid \mathcal{C}_\mathcal{T}) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\big\|_2^2\right]
\]
At inference, draw \(\boldsymbol{x}_0 \sim p_\phi(\cdot \mid \mathcal{C}_\mathcal{T})\) and integrate
the vector field as usual. CNP-Flow only modulates the source supplied to the FM; the backbone
architecture is untouched, so it drops into existing conditional FM frameworks.
One formulation, three temporal tasks
The tasks differ only in how the temporal condition is constructed:
- Video prediction (extrapolation): at step \(i\), the previous frame
is the reference and a context frame \(\boldsymbol{x}^c\) is sampled from earlier observations,
\(\mathcal{C}^i_\mathcal{T} = (\boldsymbol{x}^{i-1}, \boldsymbol{x}^c, i-c)\).
- Video interpolation: the previous frame is the reference and
the final frame \(\boldsymbol{x}^N\) is a fixed future context,
\(\mathcal{C}^i_\mathcal{T} = (\boldsymbol{x}^{i-1}, \boldsymbol{x}^N, N-i)\).
- Motion planning: the whole joint-space trajectory is generated at once
from the initial state and target end-effector pose,
\(\mathcal{C}_\mathcal{T} = (\boldsymbol{x}^0, \boldsymbol{x}^N)\).
