Expanding Flow Maps

Core Idea§

Expanding Generative Flows (EFlows) define a continuous normalizing flow between distributions of increasing dimensionality. The key construction is an expanding interpolant: instead of a fixed-dimensional path $x(t)$ for $t\in[0,1]$, we define $x(t) \in \mathbb{R}^{d(t)}$ where $d(t)$ grows monotonically. This is achieved by augmenting the state with conditional noise: at each time $t$, we sample new coordinates from a conditional distribution given the current state, effectively "expanding" the representation.

From EFlows, the authors propose Expanding Flow Maps (EFMs), a distillation into few-step generative models. Each EFM factors the map between two timesteps into two learnable operations: 1. Expand operator: $E_\theta(x, t) = (x, \tilde{y})$ where $\tilde{y} \sim p_\theta(\cdot | x, t)$ adds new coordinates. 2. Transport map: $T_\phi(x, t) = x + v_\phi(x, t)$ which pushes the expanded state forward along the interpolant (e.g., via a residual block).

Composing these yields a map that jointly expands and denoises. The framework recovers standard fixed-dimensional flows when the expand operator is identity.

Discrete Extension§

The construction is extended to the probability simplex for discrete data. For variable-length sequences, the expand operator adds a new token from a categorical distribution conditioned on the current sequence; the transport map then refines the entire sequence via a diffusion-like process over the simplex.

Implementation Sketch§

Below is a simplified PyTorch implementation of an EFM block for continuous data:

import torch
import torch.nn as nn

class ExpandOperator(nn.Module):
    def __init__(self, d_in, d_out, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in + 1, hidden),
            nn.ReLU(),
            nn.Linear(hidden, d_out - d_in)  # output new coordinates
        )
    def forward(self, x, t):
        # x: (B, d_in)
        t = t.unsqueeze(-1)  # (B, 1)
        inp = torch.cat([x, t], dim=-1)
        new_coords = self.net(inp)  # (B, d_out-d_in)
        return torch.cat([x, new_coords], dim=-1)

class TransportMap(nn.Module):
    def __init__(self, d, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d + 1, hidden),
            nn.ReLU(),
            nn.Linear(hidden, d)
        )
    def forward(self, x, t):
        t = t.unsqueeze(-1)
        inp = torch.cat([x, t], dim=-1)
        delta = self.net(inp)
        return x + delta

class EFMBlock(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.expand = ExpandOperator(d_in, d_out)
        self.transport = TransportMap(d_out)
    def forward(self, x, t):
        x = self.expand(x, t)
        x = self.transport(x, t)
        return x

Training of an EFM proceeds by sampling noise $x_0$ from a base distribution and then applying learned EFM blocks to transform it to data $x_1$, with a loss that matches the transported distribution to the true data distribution at each expansion step.

Key Results§

• EFlows can be trained via standard continuous flow matching, but the expanding interpolant requires careful null-space parameterization.
• EFMs achieve comparable quality to fixed-dimensional flows on image generation (e.g., MNIST→FashionMNIST variable resolution) and outperform fixed-dimensional baselines on molecule graph generation (variable number of atoms) and text generation (variable length).
• The framework is principled, covering continuous and discrete modalities, and opens the door to generative models where output size is a learnable degree of freedom.