Wavelet Flow Matching

1 · A lossless multi-scale representation

The discrete wavelet transform (DWT) decomposes a field into components that are localized simultaneously in space and in scale. For an orthonormal wavelet $\psi$ with scaling function $\phi$,

$$ f \;=\; \sum_{n} a_J[n]\,\phi_{J,n} \;+\; \sum_{j=1}^{J}\sum_{n} d_j[n]\,\psi_{j,n}, \qquad \psi_{j,n}(x) = 2^{-j/2}\,\psi\!\left(2^{-j}x - n\right). $$

Applied to a state $x \in \mathbb{R}^{C\times H\times W}$, each scale yields one approximation (LL) and three oriented detail sub-bands (LH, HL, HH) on a progressively coarser grid,

$$ w_j = \mathcal{W}_j(x) \in \mathbb{R}^{C \times 4 \times H_j \times W_j}, \qquad H_j \approx H/2^{j},\;\; W_j \approx W/2^{j}, $$

and the inverse transform reconstructs the field exactly, $x = \mathcal{W}^{-1}(w)$. We use Daubechies wavelets ($\mathrm{haar}$, $\mathrm{db}2$, $\mathrm{db}4$, $\mathrm{db}6$): no training, no information loss.

2 · One flow per scale, one shared network

Instead of modelling $p(x^{t+1}\mid x^t)$ in pixel space, WFM learns $p\bigl(w^{t+1} \mid w^t, \{x^{t-l}\}_{l=1}^{L}, \kappa\bigr)$. At every scale $j$ we build a straight optimal-transport path from scale-matched Gaussian noise to the target coefficients, with a shared interpolation time:

$$ w_{j}^{\tau} = \tau\, w_{j}^{t+1} + (1-\tau)\,\varepsilon_{j}, \qquad \tau \sim \mathcal{U}([0,1]), \quad \varepsilon_j \sim \mathcal{N}(0, I). $$

Because the basis is orthonormal, sampling noise directly in wavelet space is equivalent in distribution to transforming pixel-space noise, $\mathcal{W}_j(\varepsilon) \overset{d}{=} \varepsilon_j$. The target velocity along each path is constant:

$$ u_j = \frac{d\,w_j^{\tau}}{d\tau} = w_j^{t+1} - \varepsilon_j . $$

The velocity network is a U-Net whose encoder–decoder hierarchy mirrors the DWT. The coefficients $z_j^{\tau} = [\,w_j^{t},\, w_j^{\tau}\,]_C$ of each scale are injected at the matching encoder level through a dedicated stem, and $J$ decoder heads predict $\hat u_j$ at their native resolution. Because all heads share one backbone, the $J$ flows are coupled at every integration step.

3 · Variance-normalized multi-scale objective

Approximation and detail coefficients differ in energy by orders of magnitude. To prevent the LL sub-band from dominating, each per-scale loss is normalized by the spatial variance of its target velocity:

$$ \ell_j(\theta) = \mathbb{E}\!\left[\frac{1}{BC}\sum_{b,c} \frac{\bigl\langle (\hat u_j - u_j)^2 \bigr\rangle} {\bigl\langle (u_j - \langle u_j\rangle)^2 \bigr\rangle + \epsilon}\right], \qquad \mathcal{L}(\theta) = \sum_{j=1}^{J} \frac{\lambda_j}{\sum_{j'}\lambda_{j'}}\,\ell_j(\theta). $$

4 · Sampling and autoregressive rollout

Sampling happens entirely in wavelet space. Starting from $w_j^{0} = \varepsilon_j$, we integrate the learned field with $N$ Euler steps,

$$ w_j^{\tau_{n+1}} = w_j^{\tau_n} + \Delta\tau\; u_\theta\!\left(\bigl[w_j^{t}, w_j^{\tau_n}\bigr]_C,\, \tau_n,\, \kappa,\, \{x^{t-l}\}_{l=1}^{L}\right), $$

and recover the physical state with a single inverse transform, $\hat x^{t+1} = \mathcal{W}^{-1}\bigl(\{\hat w_j^{t+1}\}_{j=1}^{J}\bigr)$, which is fed back as the next input.