Few Step Variable Length Sampling with Branching Flow Maps

Few step "flow maps" target the same problem space as diffusion and flow matching, but at vastly lower inference cost, which is becoming increasingly important with the large-scale use of inference-time reward alignment strategies. However, most standard diffusion, flow matching, and flow map models assume the length of the output is known a priori, which limits their usefulness, especially in scientific domains such as conditional small molecule or protein design. While variable-length flow matching approaches exist, they are inefficient at inference time, requiring many neural function evaluations. Few-step flow maps operate on continuous spaces, and it is unclear how to readily extend these to length variation, which is structurally a discrete property. Here we describe Branching Flow Maps, which enables few-step variable length generation on both discrete and continuous state spaces via element splitting and deletion. We demonstrate that our construction can compose with various established flow map approaches, affording efficient variable length generation with dramatically fewer neural function evaluations. We characterize the behavior of Branching Flow Maps on variable-length unconditional and conditional sampling tasks in synthetic and scientific domains, including small molecules and proteins, and we demonstrate accelerated reward-tilted sampling, with a sampling process that is now able to explore length as an additional dimension.