Research map / Probabilistic and scientific modeling / Variational and generative sampling
Path Gradients after Flow Matching
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Fine-tunes Flow-Matching-trained continuous normalizing flows (Boltzmann Generators) with path gradients on target samples (forward KL), raising ESS and NLL at similar compute on LJ13/LJ55/alanine dipeptide.
Problem
Boltzmann Generators trained with Flow Matching do not exploit first-order target energy information, and sampling efficiency (ESS) of CNFs remains limited for molecular systems.
Approach
Hybrid training: pre-train a CNF with Flow Matching, then fine-tune with path gradient estimators of the forward KL on existing target samples (reusing precomputed forces), using the augmented adjoint method for constant memory.
Key results
- On LJ55, hybrid fine-tuning with path gradients nearly triples ESS_p (2.97% to 13.71%) at similar wall-time.
- On alanine dipeptide XTB, the hybrid approach roughly doubles ESS_q (0.74% to 0.25%? no: TBG XTB ESS_q 3.38% to 7.30%).
- Fine-tuning leaves flow trajectory length largely unchanged on LJ13 while NLL improves (e.g., -16.01 to -16.22 for OT FM).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune the OT FM-pretrained CNF with Path Gradients (ablation) | Boltzmann Generators · LJ13 | ESS_q (%): 54.36 ± 5.43 -> 82.97 ± 0.40 | Helped |
| Fine-tune the OT FM-pretrained CNF with Path Gradients (ablation) | Boltzmann Generators · LJ55 | ESS_p (%): 2.97±0.08 -> 13.71 ± 3.15 | Helped |
| Fine-tune the OT FM-pretrained CNF with Path Gradients (ablation) | Boltzmann Generators · Alanine dipeptide - XTB | ESS_q (%): 0.74± 0.33 -> 0.25 ± 0.37 | Hurt |
| Fine-tune the OT FM-pretrained CNF with Path Gradients (ablation) | Boltzmann Generators · Alanine dipeptide - Classical | ESS_q (%): 4.88 ± 0.42 -> 2.86 ± 3.05 | Hurt |
| Fine-tune the OT FM-pretrained CNF with Path Gradients (ablation) | Boltzmann Generators · Alanine dipeptide - Classical | NLL: -110.14±0.01 -> -110.30 ± 0.13 | Helped |
| Fine-tune the transferable Boltzmann Generator with Path Gradients (ablation) | Transferable Boltzmann Generator · Unseen dipeptides (average) | ESS_q (%): 7.96 -> 9.79 | Helped |
| Fine-tune the transferable Boltzmann Generator with Path Gradients (ablation) | Transferable Boltzmann Generator · Unseen dipeptides (average) | NLL: -100.73 -> -100.93 | Helped |
| Fine-tune the FM-trained flow with Path Gradients (ablation) | LJ13 flows trained with Standard FM, OT FM, or EQ OT FM · LJ13 | ESS_q (%): 60.45 ± 0.89 -> 85.78 ± 0.69; 57.16 ± 0.88 -> 85.91 ± 0.03; 58.11 ± 1.93 -> 87.77 ± 0.24 | Helped |
| Fine-tune the FM-trained flow with Path Gradients (ablation) | LJ13 flows trained with Standard FM, OT FM, or EQ OT FM · LJ13 | Trajectory Length: 3.76 ± 0.00 -> 3.86 ± 0.00; 2.80 ± 0.01 -> 2.78 ± 0.04; 2.11 ± 0.00 -> 2.19 ± 0.00 | No clear effect |
| Switch from FM pre-training to Path Gradient fine-tuning (ablation) | 2D Gaussian Mixture Model CNF · 2D GMM | Forward KL | Helped |
| Apply Path Gradient fine-tuning after Flow Matching (ablation) | 2D Gaussian Mixture Model CNF · 2D GMM | MSE Flow Matching loss | No clear effect |
Surprising findings
- Path gradient fine-tuning does not reliably improve the standard FM model on AD2 with XTB (NLL slightly worse, ESS_q lower)
- Flow Matching loss is essentially unchanged by path gradient fine-tuning, and FM gradient estimators do not vanish at the optimum in a toy example
Limitations
- Requires differentiable energy function and unbiased target samples with precomputed forces
- Benefit only when the flow is already close to the target; gains are small or negative otherwise
- Higher per-iteration compute and training instability (gradient clipping needed on AD2)
Experimental setup
Datasets: LJ13 LJ55 Alanine dipeptide classical (Amber ff99SBildn MD) Alanine dipeptide XTB (GFN2-xTB relaxed) 2AA dipeptide dataset (timewarp) 2D Gaussian mixture model
Benchmarks: LJ13 LJ55 Alanine dipeptide (classical and XTB) 2AA dipeptides transferability
Metrics: ESS_p ESS_q NLL trajectory length MSE flow matching loss
Baselines: Only FM (OT Flow Matching) Maximum likelihood training Standard FM EQ OT FM
Builds on: Flow Matching Path Gradients augmented adjoint method Boltzmann Generators Equivariant CNF (EGNN)
Closest papers by meaning
- Conditional Normalizing Flows for Active Learning of Coarse-Grained Molecular Representations (ICML 2024 · similarity 0.89)
- Transferable Boltzmann Generators (NeurIPS 2024 · similarity 0.88)
- Scalable Equilibrium Sampling with Sequential Boltzmann Generators (ICML 2025 · similarity 0.88)
- Temperature-Annealed Boltzmann Generators (ICML 2025 · similarity 0.87)
- BoltzNCE: Learning likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation (NeurIPS 2025 · similarity 0.87)
- Energy-Based Flow Matching for Generating 3D Molecular Structure (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.