General Proximal Flow Networks

Alexander Strunk; Roland Assam

cs.LGFeb 28, 2026

General Proximal Flow Networks

Authors: Alexander Strunk, Roland Assam

Organizations: Evercot AI

Abstract

This paper introduces General Proximal Flow Networks (GPFNs), a generalization of Bayesian Flow Networks that broadens the class of admissible belief-update operators. In Bayesian Flow Networks, each update step is a Bayesian posterior update, which is equivalent to a proximal step with respect to the Kullback-Leibler divergence. GPFNs replace this fixed choice with an arbitrary divergence or distance function, such as the Wasserstein distance, yielding a unified proximal-operator framework for iterative generative modeling. The corresponding training and sampling procedures are derived, establishing a formal link to proximal optimization and recovering the standard BFN update as a special case. Empirical evaluations confirm that adapting the divergence to the underlying data geometry yields measurable improvements in generation quality, highlighting the practical benefits of this broader framework.

Figures & tables

Panel of Figure 1 — Figure 1: The GPFN generation process.

Panel of Figure 2 — Figure 2: The GPFN training procedure.

Panel of Figure 3 — Figure 3: Generated MNIST samples (64 images per panel, NFE = 20).

NFE	GPFN-stoch			GPFN-det			BFN-stoch			BFN-det
	SWD $\!\downarrow$	aFID $\!\downarrow$	IS/P/R/D/C/Div $\!\uparrow$	SWD $\!\downarrow$	aFID $\!\downarrow$	IS/P/R/D/C/Div $\!\uparrow$	SWD $\!\downarrow$	aFID $\!\downarrow$	IS/P/R/D/C/Div $\!\uparrow$	SWD $\!\downarrow$	aFID $\!\downarrow$	IS/P/R/D/C/Div $\!\uparrow$
	( $\times 10^{-2}$ )			( $\times 10^{-2}$ )			( $\times 10^{-2}$ )			( $\times 10^{-2}$ )
5	4.06	193	7.5/0.91/0.83/0.73/0.78/0.41	3.80	166	7.7/0.87/0.86/0.71/0.78/0.42	14.82	2910	1.2/1.00/0.00/0.48/0.01/0.31	20.47	3423	1.0/1.00/0.00/0.40/0.00/0.00
10	2.75	95	8.3/0.86/0.90/0.73/0.84/0.40	2.68	86	8.3/0.84/0.92/0.67/0.83/0.40	11.35	2139	1.8/1.00/0.00/0.76/0.02/0.32	21.95	3614	1.0/1.00/0.00/0.20/0.00/0.00
20	2.51	74	8.5/0.86/0.93/0.72/0.85/0.41	2.55	67	8.6/0.84/0.95/0.65/0.84/0.40	8.72	1513	2.8/0.99/0.01/0.67/0.06/0.32	23.85	3798	1.0/1.00/0.00/0.20/0.00/0.00
40	2.66	72	8.5/0.87/0.94/0.75/0.88/0.40	2.75	69	8.6/0.83/0.96/0.63/0.83/0.41	6.34	1141	3.7/0.93/0.14/0.51/0.10/0.33	24.26	3825	1.0/1.00/0.00/0.20/0.00/0.00

Table 1: Generation metrics for GPFN and BFN across different NFE budgets. Reported are Sliced Wasserstein Distance (SWD

\downarrow

), approximate FID (aFID

\downarrow

), Inception Score (IS

\uparrow

), Precision (P

\uparrow

), Recall (R

\uparrow

), Density (D

\uparrow

), Coverage (C

\uparrow

), and Intra-set Diversity (Div

\uparrow

). Best values per row for SWD and aFID are in bold .

Figure 4: Generation metrics vs. NFE budget. GPFN achieves superior performance at a fraction of the steps required by BFN.

Explore similar work

Proximal Policy Optimization for Amortized Discrete Sampling
Jun 14, 2026Anna Zykova-Myzina, Timofei Gritsaev, Daniil Tiapkin +1Molecular Graph Generation Generative Flow Networks
HSE University · Constructor University · CMAP, CNRS, École polytechnique, IPP
Variational Bayesian Flow Network for Graph Generation
Jan 30, 2026Yida Xiong, Jiameng Chen, Xiuwen Gong +3Graph Generation Molecular Graph Generation
School of Computer Science, Wuhan University, Wuhan, China · Department of Computing, Macquarie University, Sydney, Australia · Griffith University, Gold Coast, Australia +1
Stable GFlowNets with Probabilistic Guarantees
May 3, 2026Zengxiang Lei, Ananth Shreekumar, Jonathan Rosenthal +6Generative Flow Networks Flow-Based Generative Modeling
Lyles School of Civil and Construction Engineering, Purdue University, IN, USA · Department of Computer Science, Purdue University, IN, USA · Computer Science and Engineering, University of California, Santa Cruz, USA

Jun 14, 2026J/K move · Enter open · S save

Proximal Policy Optimization for Amortized Discrete Sampling

Anna Zykova-Myzina, Timofei Gritsaev, Daniil Tiapkin, Nikita Morozov

HSE University · Constructor University · CMAP, CNRS, École polytechnique, IPP

Molecular Graph Generation Generative Flow Networks Policy Gradient Methods Proximal Policy Optimization

Open paper AskPDF arXiv

This paper explores policy gradient algorithms for training stochastic policies to sample from structured discrete probability distributions under the Generative Flow Network (GFlowNet) framework. Building on extensive theoretical connections between GFlowNets and entropy-regularized reinforcement learning, we derive equivalents of standard policy gradient algorithms for training GFlowNets, as well as experimentally explore their various methodological aspects, including baseline training and advantage estimation. Most importantly, our work is the first to derive and successfully apply proximal policy optimization to GFlowNets, showing its improved convergence speed and data efficiency compared to standard GFlowNet training objectives on benchmarks ranging from synthetic energies to molecular graph generation.

Abstract

Explore similar work

Proximal Policy Optimization for Amortized Discrete Sampling

Variational Bayesian Flow Network for Graph Generation

Stable GFlowNets with Probabilistic Guarantees