Semigroup Riemannian Transport Flow for Few-Step Crystal Generation

Deep generative models are the default engine for inorganic crystal design, but each structure costs 50–5000 neural function evaluations, which dominates the cost of screening campaigns, closed-loop laboratories and inverse problems that need many samples per query. The obstacle is what the model learns: flow matching predicts an instantaneous velocity, whereas one large step needs the finite-horizon velocity—the average of that velocity over the interval the step covers. Supervising it is the hard part: distillation needs a converged teacher, and derivative-based self-consistency objectives differentiate through the model, requiring a Jacobian-vector product (JVP). The finite-horizon velocity is instead characterized by a composition law—one long transport step equals two composed shorter steps—together with a boundary condition that pins the field down as the step length goes to zero, and on the flat product manifold used for crystal generation the composition law holds exactly. It can be imposed with ordinary forward evaluations of a frozen copy of the model, with no teacher and no derivative through the model. This yields Semigroup Riemannian Transport Flow (SRTF): a semigroup consistency loss supplies the composition law, a standard flow-matching loss supplies the boundary condition, and the field is parameterized as a local velocity plus a correction proportional to the step length, so the two can be read and tested separately. We prove that the objective's population fixed point recovers the transport coefficients to leading order, so the correction is identified rather than merely fitted, and on trained networks the measured and predicted residual coefficients agree within 3%. On the four standard crystal-structure-prediction benchmarks SRTF exceeds the single-sample match rate of the published baselines considered here with only two to five sampling steps, 10–25 times fewer than the fastest of them, at flow-matching training memory, whereas our JVP implementation trains at only a third of the batch size on the same GPU memory.