Expanding Flow Maps: Learning How Big Your Output Should Be

Paper: 2607.21585 Authors: Sophia Tang, Pranam Chatterjee Categories: cs.LG

The Gap

Flow-based generative models have become a serious alternative to diffusion models for fast, controllable generation. Methods like rectified flows, flow matching, and continuous normalizing flows learn a velocity field that transports samples from noise to data in few steps. They work beautifully — but with one stubborn assumption baked in: the input and output live in the same dimensional space. You start with a noise vector of size d, you end with a data vector of size d.

This is fine when generating fixed-size images (say, 256×256). But what about generating a molecular graph with 12 atoms? Or a protein sequence of 280 residues? Or a variable-length sentence? The output size is itself a meaningful degree of freedom, yet existing flow models either pad/truncate to a fixed canvas or simply can’t handle it.

Prior work has made partial moves — diffusion models for graphs, autoregressive decoders for sequences — but nobody has built a principled flow-based framework where the state space itself grows during generation. That’s the gap this paper fills.

Problem                              Assumption in prior work
Variable-size outputs                Input dim = Output dim
(graphs, sequences, molecules)       (fixed canvas throughout)
        |                                      |
        v                                      v
  Need flows that              Existing flows assume
  GROW the state space         static dimensionality
        |                                      |
        +------------------ gap ----------------+
                               |
                               v
                         EFlows / EFMs
                     (expand + transport)
                               |
                               v
                    Principled variable-size
                    generation via learned
                    expand operators
                               |
                               v
                        Experiments on
                    continuous & discrete
                       modalities

The Increment

One sentence: Before this paper, flow-based generative models could only produce fixed-size outputs; after, they can learn to expand the canvas itself as part of the generation process.

Core Mechanism

The framework has two key components: Expanding Flows (EFlows) and Expanding Flow Maps (EFMs).

An EFlow defines an interpolant that doesn’t just move probability mass from noise to data — it simultaneously grows the state space. Starting from a small initial state (e.g., a single token or a low-dimensional vector), the interpolant progressively augments it with new coordinates or tokens, each conditioned on what already exists. This is a continuous-time process: the dimensionality increases smoothly along the trajectory from t=0 to t=1.

EFMs take this continuous-time expanding flow and distill it into a practical few-step generative model. Each step between two timesteps factors into two operations. First, the expand operator augments the current state — adding new token positions, graph nodes, or coordinates — conditioned on the existing state. Second, the transport map pushes the expanded state forward along the interpolant toward the target distribution. Composing these operators over K steps yields a single map that jointly expands and denoises.

Step 1                    Step 2                    Step 3
x_0 (small)              x_1 (medium)              x_2 (large)
    |                         |                         |
    v                         v                         v
[Expand] -- add coords --> [Expand] -- add coords --> [Expand]
    |                         |                         |
    v                         v                         v
[Transport]               [Transport]               [Transport]
    |                         |                         |
    v                         v                         v
x_0'                      x_1'                      x_2' = output
(grown,                   (grown,                   (target
 denoised)                 denoised)                 size reached)

The crucial insight: when the expand operator is the identity (no growth), EFMs collapse to existing fixed-canvas flow maps. So this is a genuine generalization, not a replacement.

Structural Metaphor

Think of building a house on a lot where you don’t know the final floor plan in advance.

The expand operator is the foundation crew. Each morning, they arrive and pour new foundation slabs — but they decide how many and where based on the structure already standing. They’re not randomly extending; they look at the current frame and ask “what does this house need next?” Similarly, the expand operator adds new state dimensions conditioned on what already exists — the existing tokens guide where new ones should appear.

The transport map is the framing and finishing crew. Once new foundation exists, they bring the whole structure closer to habitable condition — framing walls, wiring, plumbing — pushing the state from “expanded skeleton” toward “finished house.” They transform the raw expanded state into something that better matches the target distribution.

The timestep schedule is the project manager. They decide how many rounds of expansion happen. Three rounds for a small cottage (short sequence), twelve rounds for a mansion (long sequence). The output size is not hardcoded; it emerges from how many expand-transport cycles the model runs.

Without this decomposition, you’d need a single monolithic contractor who simultaneously decides the house size AND builds it — which is what fixed-canvas flow models try to do, and why they break when the output size isn’t known in advance.

Key Concepts

• Expanding Interpolant: A path between distributions that doesn’t just transport mass — it grows the space the mass lives in. Imagine pouring water into a balloon: the water is the probability mass, and the balloon stretches (expands) as you pour. The interpolant defines *how the balloon stretches at each moment, not just where the water goes. This is the continuous-time recipe that EFMs learn to approximate in few steps.

• Factorization into Expand + Transport: Instead of learning one big function that does everything, the method separates “how big should this be?” from “what should the content be?” This is like separating the question “how many paragraphs should this essay have?” from “what should each paragraph say?” — two distinct decisions that interact but can be learned independently. The factorization is what makes variable-size generation tractable.

• Discrete Simplex Extension: The framework extends to probability distributions over categorical variables (like tokens in a sequence or node types in a graph). Here, “expanding” means adding new rows to a probability simplex — new token positions whose initial distributions are learned. This bridges continuous flows with the discrete world where most practical applications live.

Framework Shift

Before (mainstream approach):          After (this paper):
                                      
  Noise (fixed dim)                    Noise (small dim)
      |                                    |
      v                                    v
  [Flow: transport]                   [Expand + Transport]
      |                                    |
      v                                    |
  Data (same fixed dim)                   v
                                      [Expand + Transport]
  "One canvas, one size"                  |
                                          v
                                     Data (variable dim)
                                      
                                      "Canvas grows with content"

From fixed-dimensional transport to expanding transport, the core shift is treating output size as a learned generation decision rather than a hardcoded assumption.

Expert Assessment

Problem choice: This is a real gap. The fixed-dimension assumption in flow models is a genuine constraint that limits applicability to graphs, variable-length sequences, and molecules. The field has been quietly working around it with padding tricks and post-hoc truncation, which are inelegant and wasteful. The paper positions itself cleanly at the intersection of flow matching and variable-structure generation — a sweet spot with growing interest.

Method maturity: Clever insight rather than brute force. The expand-transport factorization is clean and principled, and the observation that fixed-canvas flows are a special case is mathematically satisfying. However, the expand operator itself is somewhat hand-wavy in the paper — the details of *how it learns to add dimensions (architecture, training objective for the expansion decision) deserve more scrutiny. There may be simpler baselines that achieve similar results through adaptive padding with learned masks, though these would lack the theoretical elegance.

Experimental integrity: The experiments span both continuous and discrete settings — variable-size graph generation and variable-length sequence generation — which is good for a framework paper. However, the baselines could be more comprehensive. Comparing against strong autoregressive baselines for sequence tasks and against equivariant graph diffusion models would strengthen the claims. The few-shot efficiency gains are promising but the evaluation metrics for variable-size generation are notoriously tricky; the paper should be more explicit about how they handle size-conditioned quality assessment.

Writing quality: The paper is well-structured and the notation is reasonably clean. The weakest section is the discrete simplex extension — it reads like an afterthought grafted onto the main continuous framework. Rewriting Section 4 with a self-contained, intuitive treatment of the discrete case would significantly improve accessibility. The related work section is also thin; a deeper comparison with diffusion-based graph generation (e.g., EDP-GNN, GDSS) would help readers situate the contribution.

Verdict: weak accept — The expand-transport factorization is an elegant idea that addresses a real limitation, but the experiments need stronger baselines and the discrete extension needs more polish to be fully convincing.

Takeaways

Steal this: The factorization of generation into “how big?” and “what content?” is a transferable design pattern. If you’re building any generative system where output structure varies — code generation, layout synthesis, UI generation — consider separating the size decision from the content decision as independent learned operations.

Steal this too: The “expanding interpolant” idea suggests you can define flows that don’t just move between fixed distributions but *grow the state space during interpolation. This could inspire new approaches to hierarchical generation, coarse-to-fine modeling, or any setting where the natural representation has variable dimensionality.

Watch out for: The expand operator’s training dynamics are the make-or-break component. If the model learns poor expansion decisions early on, the downstream transport is fighting an uphill battle. Practitioners adopting this framework should invest heavily in debugging and regularizing the expansion step.