Expanding Flows for Fast and Flexible Generation Beyond the Fixed Canvas
This talk introduces Expanding Generative Flows (EFlows) and Expanding Flow Maps (EFMs), a novel framework for flow-based generative models that overcome the limitation of fixed-canvas generation. It proposes decomposing the generative process into an 'expand' operation (to increase dimensionality with new coordinates/tokens) and a 'transport' map, allowing for fast and flexible generation of variable-sized outputs across both continuous and discrete state spaces. The framework leverages local time clocks, piece-wise deterministic Markov processes, and specific training objectives for consistency, demonstrated through applications in conformer generation, molecular graph generation, and language modeling.
Sophia Tang introduces a novel framework for generative modeling that addresses the limitations of fixed-canvas (fixed-dimension or fixed sequence length) flow-based models. The core idea is to treat output size as a learned degree of freedom, enabling flexible and fast generation of variable-sized data.
The framework begins with Expanding Generative Flows (EFlows), which extend standard continuous flow matching to handle increasing dimensionality. EFlows operate on an expanding interpolant, where the state space dimension grows from DS to DT during the flow. This expansion is achieved by augmenting the source distribution with extra noise coordinates, sampled from a conditional distribution. A critical component is the Expand Operator, which lifts an existing state Xs to a higher dimensionality DT by incorporating these new noise coordinates. Placement schemes for these coordinates can vary, including concatenation (like autoregressive appending), positional insertion, or hierarchical expansion (e.g., coarse-to-fine molecular structures). The noise itself can be deterministic or learned. To account for varying insertion times, a local clock mechanism is introduced, where each coordinate’s time trajectory starts at zero upon insertion. The Transport Operator then integrates the flow over this expanded, locally-timed state space. To reconcile discrete jumps in continuous flows, EFlows are formulated as Piece-wise Deterministic Markov Processes (PDMPs), combining continuous denoising with an expansion kernel governed by an insertion intensity.
Building upon EFlows for fast generation, Tang then introduces Expanding Flow Maps (EFMs). EFMs distill the continuous EFlows into efficient, few-step generative models. An EFM maps directly from state s to state t by first applying the Expand Operator to increase dimensionality and then utilizing a mean velocity (Transport Map) to jump across the augmented state space. This approach effectively recovers fixed-dimensional flow maps as a special case when the Expand Operator is an identity function. While insertions introduce discontinuities, the framework focuses on active coordinates to maintain consistency. EFMs satisfy analogous tangent and consistency conditions (Lagrangian, Eulerian, Semigroup) as standard flow maps, adapted for the expanding state space. Training involves a diagonal objective (for instantaneous velocity, often via distillation or self-distillation) and a consistency loss, typically enforcing the semigroup property. EFMs can also be interpreted as stochastic flow maps if the augmented noise is sampled probabilistically.
The framework is extended to discrete state spaces to enable variable-length sequence generation through token insertions. An insertion schedule (alpha_t) dictates the probability of a token being inserted by a given time. The Expand Operator inserts noisy tokens, often with local clocks reset to zero at insertion. The number of insertions during inference is typically sampled from a binomial distribution, which is bounded and has vanishing variance for few-step generation, preventing arbitrary sequence length growth, unlike an unbounded Poisson distribution. The discrete EFM leverages a mean denoiser on the augmented state space, and training uses cross-entropy losses for both a diagonal denoising objective and a consistency objective.
Experimental validation demonstrates the framework’s generalizability across diverse tasks: (1) Continuous state space conformer generation on QM9 and Gome drugs, where EFM achieves performance comparable to state-of-the-art diffusion models with significantly fewer steps (1-4 steps vs. ~20 for EFlows). (2) Discrete molecular graph generation on QM9, where EFM (with a learned expand operator) outperforms existing multi-step discrete flow models and Categorical Flow Maps, particularly in one-step generation, showing lower FCD and higher validity. (3) Variable language modeling on LM1B, where multi-step EFlows are competitive with other discrete generation methods, and few-step EFMs produce coherent samples. These results underscore EFlows and EFMs as a principled, unified, and highly applicable approach for generative problems where output size is a dynamic, learned, or controllable parameter.