PaperlayerView exact revision on arXiv

Source-linked AI summary

Co-Evolving Paths and Flows via Path-Flow Alignment

Zeyu Michael Li, William Xingxu Chen, Xiang Cheng

arXiv:2610.08717v1cs.CVcs.LG

TL;DR

The paper asks how to learn a useful interpolation path jointly with a flow, since the path determines the states and velocity targets used for training. It introduces endpoint-preserving path–flow alignment and stochastic source perturbation; on image generation tasks, the regularized approach improves generation quality and extends to model-guidance training.

  • Problem

    Because the interpolation path determines both training states and velocity targets, the paper asks how to learn a good path jointly with the flow.

    Source: §1 · p. 1

  • Method

    The method jointly trains an endpoint-preserving path network and flow with an alignment objective, using stochastic source perturbation to regularize the learned path.

    Source: §1 · pp. 1–2, §4.2 · p. 9

  • Results

    The stochastic-regularized alignment improves generation quality on challenging image-generation tasks and extends to model-guidance training.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Across SiT and FlowDCN models, post-training improves generation quality under compute-matched comparisons without changing the inference architecture or sampler.

    Source: §7 · p. 12

    The entropy guarantees concern stochastic training-path marginals and need not transfer to imperfectly fitted learned samplers; the stated guarantees also rely on independent endpoint coupling.

    Source: §3.1 · p. 3, §4 · p. 7

Abstract

from arXiv · show

We study path-flow alignment as a unified training objective for flow matching. Instead of fixing the interpolation path and learning only the velocity field, we jointly train an endpoint-preserving path network and a flow network using the same alignment loss: the flow learns to match the path velocity, and the path learns to align its velocity to the current flow. Although every fixed learned path defines a valid flow-matching objective, the alignment loss alone is not a reliable criterion for path learning. We identify path overfitting, a failure mode in which the alignment loss decreases while sample quality worsens. We find that this failure is associated with low-entropy bottlenecks in the induced probability path, where the learned path routes samples through overly concentrated intermediate marginals. Motivated by this diagnosis, we introduce a stochastic path regularizer that hides part of the source information from the path network while preserving exact endpoints. The resulting regularization gives an explicit entropy floor for the stochastic training-path marginals and empirically suppresses the bottleneck in the learned sampler, making joint path-flow training effective. On ImageNet-256x256 with SiT backbones, our method consistently improves FID across model scales, extends to model-guidance training, and leaves the inference-time architecture and sampler unchanged. Code is available at https://github.com/lizeyu090312/traj_opt_paper

1 Introduction

The paper asks whether the interpolation path can be learned jointly with the flow, rather than fixed in advance. It proposes path–flow alignment, diagnoses path overfitting, and introduces stochastic regularization to address entropy collapse.

1 Introduction

  • Path–flow alignment jointly trains an endpoint-preserving path network and a flow network using the flow-matching objective as their shared alignment signal.The flow matches the path velocity, while the path aligns its velocity to the current flow.Source: §1 · p. 1
  • Alignment loss can keep improving while generation quality deteriorates, a failure the paper calls path overfitting.The authors associate this failure with intermediate-marginal entropy collapse and low-entropy bottlenecks.Source: §1 · p. 2
  • Gaussian perturbation of the path-network input is introduced to address entropy collapse while preserving exact path endpoints.Under the independent endpoint coupling used here, the regularization establishes an entropy floor for stochastic training-path marginals.Source: §1 · p. 2

1.1 Contributions

  • On ImageNet-256 × 256, the method consistently improves FID across SiT model sizes and also extends to model-guidance training.The paper reports that the model-guidance alignment template improves over its baseline.Source: §1.1 · p. 3, §1 · p. 3

2 Related Work

Prior flow-matching work changes couplings, trajectory geometry, or paths; this paper instead learns an endpoint-conditioned residual path as training-time supervision and regularizes intermediate marginals.

2 Related Work

  • Prior approaches improve flow matching through endpoint couplings or weights, reflow, geometric paths, or learned paths and forward processes.The passage names learned curved paths, Neural Flow Diffusion Models, ALI-CFM, and Curly-FM.Source: §2 · p. 3
  • Unlike methods targeting straightness, likelihood, or observed intermediate marginals, this method optimizes path–flow alignment and regularizes intermediate marginals.After sampling endpoints, it learns an endpoint-conditioned residual path as a training-time supervision mechanism.Source: §2 · p. 3
  • Classifier-free guidance combines conditional and unconditional predictions at sampling, whereas model guidance incorporates a guided direction into training for single-evaluation sampling.The paper uses standard CFG evaluation in its main experiments and applies path–flow alignment to model-guided velocity objectives.Source: §2 · p. 3

3 Flow-Guided Interpolation

The paper defines flow-guided interpolation as learning endpoint-preserving paths and matching a flow to their velocities. Alternating alignment can overfit: loss falls even as generation quality worsens.

3.1 Problem setup and notation

  • The setup uses independently coupled source and target endpoints, with a Gaussian latent source and a label-conditional target distribution.The paper notes that alignment and population correctness extend to more general couplings under stated conditions, while entropy guarantees require an additional conditional-independence condition.Source: §3.1 · pp. 3–4

3.2 Flow matching and induced transport

  • For a sufficiently regular endpoint-preserving path, the population flow-matching target induces an ODE whose marginals match the path marginals and reach the target distribution.The target velocity averages the tangents of paths passing through the same state, time, and label.Source: §3.2 · p. 4

3.3 Endpoint-preserving neural paths

  • The learned path adds a residual term multiplied by t(1 −t) to the linear interpolation, so the residual can curve the interior while preserving both endpoints.The envelope vanishes at t = 0 and t = 1.Source: §3.3 · p. 5

3.4 Flow-guided path alignment and alternating optimization

  • Path and flow alternate optimization on the alignment objective, while a kinetic-energy term regularizes the path toward the straight line.The same alignment term supervises both phases; the path phase freezes the flow and the flow phase trains on the current path.Source: §3.4 · p. 5
  • In the toy mixture example, path–flow alignment produces more curved trajectories than standard flow matching, and stochastic regularization changes intermediate marginals.The path affects the training states and tangent targets; sampling uses only the learned velocity field.Source: §3.4 · p. 5

3.5 Path overfitting

  • Unregularized training lowers flow-matching loss while FID first improves and then rapidly deteriorates, showing that low loss alone does not ensure good generation.Figure 2 tracks flow-matching loss and Figure 3 tracks FID across alternating cycles; the authors call the mismatch path overfitting.Source: §3.5 · p. 6, Figure 2 · p. 6, Figure 3 · p. 6

4 Entropy Bottlenecks and Stochastic Regularization

The section diagnoses entropy bottlenecks in learned flows using divergence as an entropy-rate signal, then develops Gaussian source perturbation to regularize the training path. The entropy guarantees concern stochastic training-path marginals; their transfer to finite learned samplers is not guaranteed.

4.1 Divergence spikes reveal entropy bottlenecks

  • The overfit ρt = 1 run has a pronounced positive-divergence spike near t = 1, unlike the baseline and stochastic source-regularized run.The diagnostic compares per-dimensional binned divergence over time.Source: Figure 4 · p. 7, §4.1 · pp. 7–8
  • The overfit run’s cumulative normalized divergence dips below zero around t ≈ 0.9, supporting a bottleneck interpretation and compensating expansion.Values below zero indicate entropy below the generated endpoint distribution.Source: Figure 5 · p. 7, §4.1 · p. 8

4 Entropy Bottlenecks and Stochastic Regularization

  • The entropy guarantee applies to stochastic training-path marginals, not necessarily finite learned-sampler marginals, because finite optimization and capacity can separate them.Velocity-level control alone does not transfer the path-level entropy guarantee to the learned sampler.Source: §4 · p. 7

4.1 Divergence spikes reveal entropy bottlenecks

  • Expected flow divergence equals the rate of change of differential entropy, so late expansion after earlier contraction indicates an overly concentrated intermediate marginal.The section calls this a low-entropy bottleneck.Source: §4 · pp. 6–7, §4.1 · p. 7

4.2 Regularization via Gaussian perturbation

  • Gaussian perturbation hides part of the source input from the residual path network while preserving the true endpoint interpolation and exact endpoints.The residual network receives a source-like input with reduced correlation to the true source endpoint; the perturbation does not change the endpoints.Source: §4.2 · pp. 8–9
  • Under the stated independent-coupling condition, source perturbation Gaussian-smooths intermediate training-path marginals and gives them an entropy floor.The bound becomes vacuous as t approaches 1, where endpoint preservation requires smoothing to vanish.Source: §4.2 · pp. 9–10
  • With Gaussian perturbation, kinetic-energy control complements smoothing by bounding how quickly the path-induced marginal can decrease in entropy.The bound recovers the standard linear-path scaling when ρt is bounded away from 1, and weakens as ρt approaches 1.Source: §4.2 · p. 10

5 Generalizing Beyond Flow Matching

Path–flow alignment extends to model-guidance training by regressing a guided effective velocity to the path tangent. The extension improves 50k-FID over standard model-guidance fine-tuning on both tested SiT models.

5 Generalizing Beyond Flow Matching

  • The alignment template applies whenever a velocity-like target is regressed to the path tangent, including model guidance.For guidance scale ω, the paper defines an effective velocity by folding classifier-free guidance into training.Source: §5 · p. 11
  • The model-guidance extension uses the same alternating updates, replacing vϕ with the effective guided velocity in both phases.The stochastic source perturbation and endpoint preservation remain unchanged.Source: §5 · p. 11
  • Model-guidance path–flow alignment improves 50k-FID over standard MG fine-tuning on both SiT-B-MG and SiT-XL-MG.The method replaces the standard velocity with an effective guided velocity in both alternating-update phases; stochastic source perturbation and endpoint preservation are unchanged.Source: §5 · p. 11

6 Experiments

On ImageNet-256×256, regularized path–flow alignment improves best 50k-FID at every tested SiT scale under a 12k flow-update budget, and also improves results for model-guidance checkpoints. Additional reported checks include compute-matched continuation, generation-seed sets, and FlowDCN-B.

ImageNet results

  • Path–flow alignment improves best 50k-FID at every SiT scale under the same 12k flow-update budget.The largest relative FID improvements are 19.1% on SiT-L and 12.7% on SiT-XL; standard fine-tuning does not improve the starting SiT-XL checkpoint.Source: §6 · p. 11, Table 1 · p. 11
  • Each run starts from the same pretrained SiT checkpoint; the baseline continues standard flow-matching training for 12k steps.The proposed method uses three alternating cycles of 3k path updates followed by 4k flow updates, and reports the best 50k-FID within the 12k flow-update budget.Source: §6 · p. 11, Table 1 · p. 11

Model guidance

  • Best FID improves from 7.940 to 5.585 on SiT-B-MG and from 1.494 to 1.322 on SiT-XL-MG.These comparisons use standard MG fine-tuning as the baseline and start from FID-converged checkpoints.Source: §6 · p. 11

Additional results

  • FID improvements persist against compute-matched continued fine-tuning and across three disjoint generation-seed sets.The ImageNet results also improve recall and Inception Score; gains are reported on FlowDCN-B, alongside toy trajectory-curvature and qualitative image-quality comparisons.Source: §6 · p. 11

7 Conclusion

The conclusion presents endpoint-preserving joint path–flow learning as promising when paired with stochastic source perturbation. The notation guide distinguishes generic population-flow symbols from learned-path cases.

7 Conclusion

  • The paper introduces joint learning of an endpoint-preserving training path and flow model, and identifies path overfitting under naive joint optimization.Stochastic source perturbation provides Gaussian-smoothing and entropy guarantees for the training path, while experiments show suppression of the associated late-time divergence bottleneck in finite learned samplers.Source: §7 · pp. 11–12
  • Across SiT and FlowDCN models, the post-training procedure improves generation quality under compute-matched comparisons without changing the inference architecture or sampler.The theory does not establish that entropy bottlenecks cause generation-quality changes or guarantee transfer of path-level entropy control to an imperfectly fitted flow.Source: §7 · p. 12

A Notation guide

  • Generic γ and v⋆γ denote population flow-matching statements, while γψ and ˆγψ denote learned-path cases used by the algorithm and stochastic regularizer.Table 2 summarizes notation for paths, induced marginals, and velocity fields.Source: A Notation guide · p. 14, Table 2 · p. 14

B Discussion and Limitations

The evaluation and theory have practical and scope boundaries: experiments focus on image generation, while entropy guarantees do not directly control finite learned-sampler marginals.

B Discussion and Limitations

  • Evaluation remains restricted to class-conditional image generation; text-to-image models and non-image domains are untested.The experiments focus on image generation using FID as the key metric.Source: B Discussion and Limitations · p. 14
  • Path–flow alignment adds training compute and memory for the path network, and alternating optimization adds engineering complexity.Compute-matched comparisons preserve reported gains but do not remove this overhead; experiments study continuation from pretrained or converged checkpoints rather than training from scratch.Source: B Discussion and Limitations · p. 14
  • The entropy guarantees control stochastic training-path marginals and path-induced population velocity, not finite learned-sampler marginals.Transfer may fail under finite optimization or model capacity; the observed association between late-time divergence, entropy bottlenecks, and FID is empirical rather than causal.Source: B Discussion and Limitations · p. 14

C Path network and subtractor implementation

The implementation learns endpoint-preserving residual paths with a dual-stem path network and a frozen subtractor. Stochastic source perturbations affect the residual branch’s inputs, not the path’s true endpoints.

C Path network and subtractor implementation

  • The path network and frozen subtractor define a residual around the linear path, giving exact linear startup when both are initialized from the same checkpoint.The subtractor is frozen, and both networks are also frozen during flow training.Source: C Path network and subtractor… · pp. 15–16
  • The residual branch receives perturbed inputs, while the outer path uses the true endpoints and a boundary envelope preserves exact endpoint identities.The source perturbation schedule is the stochastic regularizer; perturbations are applied only to the path-network and subtractor inputs.Source: C Path network and subtractor… · p. 15
  • The dual-stem path network embeds the linear state and endpoint displacement separately before processing them with the SiT transformer body.The path network shares the transformer body with the subtractor but replaces its single spatial input stem with two stems.Source: C Path network and subtractor… · p. 15

D Diagnostic correlations for path overfitting

Across the SiT-S diagnostic sweep, late-time divergence tracked FID degradation more clearly than the measured speed and acceleration summaries. The plots are qualitative diagnostics, not a statistical correlation study.

D Diagnostic correlations for path overfitting

  • The plots cover a small set of coupled SiT-S checkpoints and are qualitative diagnostics rather than a statistical correlation study.Checkpoints span 4k, 8k, and 12k flow-training steps; the FID axis is truncated to [8, 11].Source: D Diagnostic correlations for path… · p. 16
  • Large positive final-bin divergence accompanies FID-degrading checkpoints, while stochastic source-regularized checkpoints avoid the large late-time spike and maintain low FID.Acceleration regularization can reduce some geometric quantities but does not reliably remove the late positive-divergence behavior.Source: D Diagnostic correlations for path… · p. 16
  • Among the measured quantities, the late-time divergence statistic is the clearest visual diagnostic of FID degradation, unlike path or flow speed and acceleration.The divergence statistic is more directly tied to the entropy-rate identity.Source: D Diagnostic correlations for path… · p. 16

E.1 Proof of Proposition 1

The proof sections connect population flow matching to path-induced probability transport, entropy evolution, and Gaussian smoothing bounds. The accompanying experiments report lower FID from path–flow training than longer standard continuation across three model sizes.

E.1 Proof of Proposition 1

  • Under the stated Lipschitz and linear-growth assumptions, the ODE generated by the population-minimizing conditional-mean velocity shares the path marginals and reaches the data endpoint distribution.The proof obtains the population minimizer by conditioning the path tangent on the current state, time, and label.Source: §E.1 · pp. 17–19

E.2 Proof of Proposition 2

  • The entropy proof relates differential-entropy change along a probability flow to the expected divergence of its velocity field.The proof differentiates entropy and integrates by parts under a decay assumption.Source: §E.2 · p. 19, §E.3 · p. 19

E.3 Proof of Theorem 1

  • The stochastic training-path state has a Gaussian-convolution representation, yielding an entropy floor from its independent Gaussian component.Conditioning on the non-Gaussian component leaves independent Gaussian randomness; removing the conditioning can only increase entropy.Source: §E.3 · p. 19

E.4 Proof of Proposition 3

  • The Proposition 3 proof bounds entropy change by combining control from path kinetic energy with Gaussian smoothing.The proof applies Cauchy–Schwarz and Jensen’s inequality to the path-velocity term.Source: §E.4 · p. 20

F Experimental details

  • The ImageNet experiments post-train pretrained SiT checkpoints, with standard continuation and path–flow training compared under fixed flow-update budgets.The path and flow updates alternate in cycles; FID and Inception Score are evaluated on 50k images using Heun at 250 NFE.Source: F Experimental details · pp. 20–21

F.1 Compute-matched baseline continuation

  • Path–flow training gives lower FID than longer standard continuation for all three tested model sizes, with inference architecture, NFE, and sampling cost unchanged.The comparison extends the standard baseline to the larger of GPU-hour- and FLOP-matched budgets.Source: §F.1 · p. 22, Table 7 · p. 22

F.2 CFG selection and additional metrics

Additional evaluations show a recall–precision trade-off, gains on FlowDCN-B and at higher NFE on CIFAR-10, and qualitative refinements in matched SiT-XL samples. The CIFAR-10 reproduction of the comparison method was not exact.

F.2 CFG selection and additional metrics

  • At the selected CFG, recall increases and precision decreases for all four SiT sizes compared with baseline.For SiT-S/B/L/XL, recall changes from 0.354/0.525/0.616/0.663 to 0.416/0.578/0.655/0.688, while precision changes from 0.872/0.848/0.822/0.809 to 0.857/0.824/0.798/0.795.Source: §F.2 · p. 23, Table 9 · p. 23
  • Across three disjoint generation-seed sets, path–flow training has lower mean FID at every model size, with each gap at least 5.8 times the larger observed sample standard deviation.The comparison evaluates each fixed checkpoint on three 50k-image seed ranges.Source: §F.2 · p. 23, Table 11 · p. 24

F.3 Applying our method in more scenarios

  • On FlowDCN-B, path–flow training reaches 5.473 FID, a 7.0% improvement over the best FLOP-matched continuation result of 5.888.The model uses a purely convolutional backbone.Source: §F.3 · p. 24, Table 13 · p. 24
  • On CIFAR-10, Learning Straight Flows performs best at 2 NFE, while path–flow alignment performs best at 40 NFE.The authors attribute the higher-NFE advantage to more curved trajectories; they could not reproduce Learning Straight Flows exactly because original code was unavailable and some hyperparameters were missing.Source: §F.3 · p. 24, Table 14 · p. 25

F.4 Qualitative comparisons with our method

  • Matched-class, matched-seed samples generally retain similar composition, while several examples show clearer local structure with path–flow alignment.Reported changes include a more recognizable triumphal arch and cleaner separation of a rubber eraser from surrounding material.Source: §F.4 · p. 25

Related papers

Loading 2610.08717v1…