Expanding Flow Maps

🌊 Expanding Generative Flows

The key limitation of standard flows is that the state space is fixed at initialization. We overcome this by defining generative flows between distributions of different dimensionality, specifically where the source has lower dimensionality than the target. We replace the fixed dimension $d$ with a non-decreasing schedule $d(t):[0,1]\to\mathbb{N}$, so the marginal $p_t\in\mathcal{P}(\mathbb{R}^{d(t)})$ lives on a state space whose dimensionality grows along the trajectory.

Because the source $\boldsymbol{x}_s\in\mathbb{R}^{d(s)}$ and target $\boldsymbol{x}_t\in\mathbb{R}^{d(t)}$ live in spaces of different dimension, no fixed-dimension velocity field can transport $p_s$ to $p_t$. We resolve this by augmenting the source with noise drawn from a tractable conditional distribution $p_{\epsilon|s,t}\in\mathcal{P}(\mathbb{R}^{d(t)-d(s)})$, lifting it to the target dimension before transport.

Because the state jumps to a higher dimension at each insertion, the two-stage expand-then-transport process is not a classical flow. We instead unify expansion and transport as a single piecewise-deterministic Markov process (PDMP) on the disjoint union of spaces $\mathcal{M} := \bigsqcup_t \mathbb{R}^{d(t)}$, whose generator combines a transport term for smooth denoising with a jump kernel for dimensionality increases.

---

🌊 Expand Operator

The expand operator $\mathcal{E}_{s,t}$ combines the existing coordinates of $\boldsymbol{x}_s$ with the latent coordinates of the augmented noise $\boldsymbol{\epsilon}$ according to a placement scheme.

It admits several instantiations depending on where new coordinates are placed.

• Concatenation. Append new coordinates to the end of the state, natural for autoregressive sequence generation or extending a point cloud.
• Positional Insertion. Interleave new entries at a subset of positions within the augmented state, suitable for any-length sequence generation or infilling.
• Child Expansion. Attach each new coordinate to a parent entry of $\boldsymbol{x}_s$ and initialize it relative to that parent, $\boldsymbol{x}_s^\epsilon[i] = \boldsymbol{x}_s[\mathrm{pa}(i)] + \sigma\boldsymbol{\epsilon}_i$, suitable for coarse-to-fine generation on structured states (e.g. resolving hydrogens around a heavy-atom backbone).

Since newly inserted dimensions are instantiated at different times, each component $i$ carries a local time coordinate $t_i = (t - t_i^{\text{ins}})/(1 - t_i^{\text{ins}}) \in [0,1]$, initialized at $t_i = 0$ at its insertion time $t_i^{\text{ins}}$.

---

🌊 Expanding Flow Map

The Expanding Flow Map (EFM) transports between two states along the expanding interpolant in a single learned step, jointly increasing dimensionality and pushing the flow forward in time. It is the composition of the expand operator and a fixed-dimensional flow map in the augmented space.

where $\Phi_{s,t}:\mathbb{R}^{d(s)}\to\mathbb{R}^{d(t)}$ and $v_{s,t}$ is the average velocity in the augmented state space. The EFM satisfies the same tangent and consistency (Lagrangian, Eulerian, semigroup) identities as fixed-dimensional flow maps, with the only difference being that velocities act on the expanded state $\mathcal{E}_{s,t}(\boldsymbol{x})\in\mathbb{R}^{d(t)}$. These translate into Lagrangian (LSD), Eulerian (ESD), and progressive/semigroup (PSD) self-distillation objectives that jointly train the expand operator $\hat{\mathcal{E}}_{s,t}$ and the velocity $\hat{v}_{s,t}$.

A notable property is that the EFM defines a stochastic flow map: the randomness from the conditional noise injection induces a distribution over trajectories, recovering the posterior $p_{t|s}(\cdot\mid\boldsymbol{x}_s) = (\Phi_{s,t}(\boldsymbol{x}_s))_{\#}\, p_{\epsilon|s,t}$ and enabling direct lookahead to all possible future states rather than a single deterministic trajectory.

---

🌊 Discrete Expanding Flow Maps

The framework is naturally positioned for any-length discrete generation, where a sequence or graph expands from a single token or node into a full sample. We define an expanding interpolant whose intermediate states grow in length via an insertion schedule $\alpha_t$, and learn a two-time insertion expectation $\hat{\mathcal{I}}_{s,t}(\boldsymbol{x}_s)[i]$, the expected number of tokens inserted into gap $i$ over $(s,t]$, trained with a Poisson-NLL insertion loss on both diagonal and off-diagonal time pairs. At inference, the discrete expand operator inserts $\ell_i$ latent tokens per gap, drawn from the exact conditional law of the interpolant, a Binomial over the remaining token budget.

For many-step sampling the Binomial converges to the Poisson jump law of the PDMP, while for few-step sampling, where the operator must produce the whole sequence in one jump, the Binomial concentrates on the exact count with vanishing variance, avoiding the unbounded variance of the Poisson limit. Substituting the expanding mean denoiser into the flow map yields the discrete expanding flow map, which jointly inserts and denoises tokens toward the target one-hot distribution.

Since the mean denoiser $\hat{\psi}\in\Delta^{V-1}$, the discrete consistency identities are enforced via cross-entropy analogs of the Lagrangian, Eulerian, and semigroup objectives, jointly training $\hat{\mathcal{E}}$ and $\hat{\psi}$ for few-step variable-length generation.