Structured Flow Autoencoders: Learning Structured Probabilistic Representations with Flow Matching
Authors: Yidan Xu
Yixin Wang
XuanLong Nguyen
Venue: ICLR 2026 Oral (2026)
Keywords: Flow Matching
Probabilistic Model
Representation Learning
Probabilistic Graphical Model
Autoencoder
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: A framework that composes any probabilistic graphical model with flow matching, jointly learning structured latent representations and high-fidelity generative models through a single objective.

Abstract: Flow matching is a powerful approach for high-fidelity density estimation, but it often fails to capture the latent structure of complex data. Probabilistic models like variational autoencoders (VAEs), on the other hand, learn structured representations but underperform in sample quality. We propose Structured Flow Autoencoders (SFA), a family of probabilistic models that augments graphical models with conditional continuous normalizing flow (CNF) likelihoods, enabling flow-matching-based structured representation learning. At the core of SFA is a novel flow matching objective that explicitly accounts for latent variables, allowing joint learning of the CNF likelihood and posterior. SFA applies broadly to graphical models with continuous and mixture latents, as well as latent dynamical systems. Empirical studies across image, video, and RNA-seq data show that SFA consistently outperforms VAEs and their structured extensions in generation quality, representation utility, and scalability to large datasets. Compared to generative models like latent flow matching (LatentFM), SFA also produces more diverse samples, suggesting better coverage of the data distribution.

STRUCTURED FLOW AUTOENCODERS: LEARNING STRUCTURED PROBABILISTIC REPRESEN- TATIONS WITH FLOW MATCHING
Yidan Xu, Yixin Wang Long Nguyen
Department of Statisties
University of Michigan (yt danxu, yixinw, xuan ong) eumi edu
ABSTRACT
Flow matching a powerful approach for high-fidelity density estimation, but it often fails to capture the latent structure of complex data. Probabilistic models like variational autoencoders (VAEs), on the other hand, learn structured rep- resentations but underperform in sample quality. We propose Structured Flow Autoencoders (SFA), a family of probabilistic models that augments graphical models with conditional continuous normalizing flow (CNF) likelihoods, enabling fow-matching-bass structured representation learning. At the core of SFA is a novel flow matching objective that explicitly accounts for latent variables, allowing joint learning of the CNF likelihood and posterior. SFA applies broadly to graphical models with continuous and mixture latents, as well latent dynamical systems. Empirical studies across image, video, and RNA-seq data show that SFA con- sistently outperforms VAEs and their structured extensions in generation quality, representation utility, and scalability to large datasets. Compared to generative models like latent flow matching (LatentFM), SFA also produces more diverse samples, suggesting better coverage of the data distribution.
INTRODUCTION
Generative modeling has become foundational pillar of modern machine learning, offering powerful tools for capturing complex data distributions and generating high-quality samples. Among recent advances, diffusion models (Ho et al., 2020; Nichol Dhariwal, 2021; Song et al., 2020; 202 Ib;a; Austin et al., 2021) and flow-based methods (Lipman al., 2022: Liu et al., 2022; Gat al., 2024; Tong et al., 2024; Isobe et al., 2024) have shown remarkable performance as neural density estimators, excelling likelihood estimation and figh-fidliity sample generation. In particular, flow matching has emerged as a scalable and efficient approach, aligning vector fields of probability paths using optimal transport principles, enabling efficient and scalable generative modeling with exact likelihood evaluation (Lipman et al., 2022; Liu et al.. 2022: Gat et al., 2024).
Despite their success in generation quality, neural density estimators like flow matching often fall short in structured representation learning, failing to capture or expose the rich latent structures underlying complex data. This limitation is especially salient in scientific and structured domains such as computational biology, where interpretable fow-dimensional representations are essential for downstream tasks, analysis, and control. Recent work has revealed both empirical evidence of implicit low oow-dimensional structures in pretrained diffusion models (Wang & Vastola, 2023; Chen al., 2024) and theoretical guarantees of their adaptivity to such structures (Wang et al., 2024; Li & Yan, 2024). However, these models neither explicitly model latent structure during training nor produce readily interpretable representations, limiting their utility beyond sample generation.
In contrast, probabilistic latent-variable models such as variational autoencoders (VAEs) (Kingma & Welling, 2013; Johnson et al., 2016) are explicitly designed to capture latent structure through architectures. These models leamn structured probabilistic represen-
probabilistic encoder-decode tations that can be leveraged for conditional generation and downstream tasks. However, VAEs typically underperform in data modeling and generation fidelity compared to modern flow-hased

---

Table 1:
Comparing generated samples to data samples Wi metric (Earth Mover's Distance). metric is evaluated with samples from marginal data distribution p(æ1) and that generated from (x:) I (())==)))(())dd SFA and FM achieve comparable performance marginal density estimations.
VAE
VAmp
Mixture-SVan -VAE
FM
LatentFM(w/VAE)
SFA
Mixture- -FA
W1 0.119
0.081
0.457
0.025
0.496(0.145)
0.024
0.046
S
SR
(a) Truth
(b) FM
(c) LatentFM
(d) LatentFM
VAE
S
S
R
(c) VAE
(f) Mixture SVAE
(g) SFA
(h) Mixture SFA
Figure I: Generated samples on the Pinwheel dataset with clusters. Color (a) indicates class membership, which not provided during training. Color (c) indicates latent distribution learned deterministic autoencoder. Color in (d)-(h) indicates the generated posterior value 2 ~ q(21 (a)) given the generated sample We use continuous latent variable in (g)) and a mixture latent variable in (I).
models, limiting their utility in high-resolutin or diverse generative tasks. This gap in generative fidelity also raises concerns about the reliability and expressiveness of their learned representations. This performance gap raises the question: Can we build models that retain the structured latent representations of VAES while achieving the high fidelity and scalability of flow matching?
Main idea. We propose structured flow autoencoders (SFA), new family of probabilistic models that augments graphical models with conditional Continuous Normalizing Flow (CNF) likelihoods. This family aims to combine the strengths of both approaches: the high-fidelity data modeling of neural density estimators and the structured representation learning capabilities of graphical models.
We motivate with simple latent variable model where continuous latents Rº generate observa- tions x Rd, with 0 < p < d:
i..dd 2; ~ p(z),
vrcd. ~ p(x(z).
x; i2
This standard latent variable framework enables structured representation learning through the posterior p(z]x). To enable high-fidelity data modeling in SFA, we parametrize p(œ]2) using conditional CNFs, achieving the expressivity of modern neural density estimators while maintaining structured latents. However, both the likelihood and the posterior are longer available in explicit forms. To address this challenge, we propose the Structured Conditional Flow Matching (SCFM) objective, a training objective that jointly learns both the conditional flow p(x z) and an approximate posterior g(z) x). Unlike standard flow matching that only models p(x). SCFM explicitly accounts for the conditional structure p(x (x) and posterior p(z|x). x). As illustrated in Fig. 1, this decomposition enables SFA to capture interpretable latent variables while maintaining high-fidelity generation, providing structured representation learning unavailable standard flow matching.
Contributions. (1) We introduce Structured Flow Autoencoders (SFA), family of generative models that augments graphical models with conditional Continuous Normalizing Flow (CNF) likelihoods. SFA bridges the gap between high-fidelity neural density estimation and structured representation learning, improving upon both VAEs and latent flow-bate based models. (2) We propose Structured Conditional Flow Matching (SCFM), novel training objective that extends flow matching
2

---

to explicitly incorporate latent variables. SCFM explicitly learns the conditional probability flows in the graphical model while preserving the marginal density information. SCFM enables joint learning of the likelihood and posterior, supporting both generative modeling and structured representation learning within a unified framework. (3) We demonstrate the flexibility of SFA across diverse domains, including image, video, and RNA-seq data, and modeling scenarios with continuous, finite mixture, and dynamical latent variables. SFA achieves high-fidelity sample generation, increased sample diversity, and enhances structured representation learning, while remaining computationally efficient on ligh-diiensional datasets.
Related work. Simultaneous high-fidelity generation and structured representation learning has been an important task (Grathwohl et al., 2018; Mittal et al., 2023; Dao et al., 2023; Davtyan et al., 2023), drawing particular interest in scientific domains (Bashiri al., 2021; Xu et al., 2023; Kapoor et al., 2024). Variational autoencoders (VAH) (Kingma & Welling, 2013) is one such probabilistic model that learns both generative model p(=) 1)) and inference model p(z æ)) simultaneously, typically using neural networks to parameterize exponential families. Grathwohl al. (2018); Chen et al. (2020) extended the VAE to families normalizing flows, which improves the flexibility of density estimation together with latent space learning. While appealing, VAEs fall short of modern generative
models in data modeling fidelity.
Recent work has explored combining neural density estimators with encoder-decoder frameworks (Mittal et al., 2023; Dao et al., 2023; Davtyan et al., 2023; Vahdat et al., 2021), typically mapping observations to low-dimensional latent spaces where the latent marginal distribution learnt via flows or diffusion models before decoding back to observation space. While these neural prior methods excel at dimensionality reduction, they often constrain encoders and decoders to simple parametric families (e.g., Gaussians) that inadequately capture complex data distributions. Our approach differs by making the entire likelihood and posterior flexible through conditional flows while incorporating structured latent dependencies. Another line of work, including Wang et al. (2023); Preechakul et al. (2022), focuses on modeling the likelihood with flexible diffusion models while maintaining simple Gaussian posteriors. However, these approaches lack mechanisms for structured dependencies between latent variables, limiting interpretability in complex domains. Most closely related to our work, structured variational autoencoders (SVAEs) (Johnson al., 2016; Lin et al., 2018) incorporate graphical model structure into VAEs to capture hierarchical dependencies. However, SVAES are constrained by parametric assumptions that limit expressiveness. In addition, extending SVAEs to more expressive density models like CNFs faces significant challenges: direct extensions suffer from numerical instability and computational inefficiency due to the need of likelihood evaluation every training step (Liu et al., 2022). Our flow matching approach circumvents these issues while enabling both structured representations and expressive data density modelling.
2 PRELIMINARIES: FLOW-BASED GENERATIVE MODELING
We start by reviewing continuous normalizing flows and the flow matching learning objective, laying the groundwork before introducing structured flow autoencoders (SFAs).
Notations. We follow the notations in Lipman et al. (2022) and denote the time indexed vector field
by u(.,)): : [0. 1] Edd d and equivalently, (-) R2 + Rd for t E [0, 1]. The path of probability
densities is denoted by p.(.) : Rd + R+, and the flow (.).) Rdd + Rd for E [0, 1]. In addition,
21 ~ Pdata represents
observed sample, and zu ~ Po sample from a chosen base distribution. We further denote the conditional vector field as u(:.:.)) [0. 1] x Rd + Rd, equivalently as we(:, z): Rdd Rd indexed by / E [0, 1). The path of conditional probability densities is denoted by
Pe(.(2) (z): IRd + R and the conditional flow by de (-12): IRd IRd, 1 [0, 1).
2.1 CONTINUOUS NORMALIZING FLOW
Continuous normalizing flows (CNFs) describe probability distributions by the evolution of some probability density path. Denote the observed data by x E Rd. Further, assume there exists time- dependent vector field U. : Rd > Rd t E [0, 1] that describes the evolution of a probability density path P.: Rd + R indexed by t [0. 1]: will use i to describe the density of æ. The path then solves the continuity equation o.p; = - V (v.P.), which is the Bockerr-Plande equation with zero diffusion. Due to the probabilistic representation theorem in Ambrosio et al. (2008, Theorem 8.2.1),
3

---

(a)
Figure 2: Examples of graphical models incorporated into SFAs: (a) latent continuous variable model; (b) latent finite mixture model; (c) latent linear dynamical system.
the continuity equation admits a representation formulated as solution of the ODE,
(x) v)()(())),
(0(x) 30
where di RR + Rd is forw-forward map sending Mi to μ4 Lt. This map is called flow in the machine learning literature (Chen et al., 2018; Grathwohl et al., 2018; Lipman et al., 2022). The log likelihood S(1) po (2)()) at any point be obtained by solving the instantaneous change-of-variahe formula forward in time, with initial conditions = log po (o (2))), and the target f(1) log (o(1
d (x(a)) = (- 2r (x())((:) di f(t)
(b((z)) f(0)
-(2)
(V-((((0(()))))
(3)
2.2 FLOW MATCHING
Liu et (2022) and Lipman al. (2022) concurrently introduced similar training objective for learning flexible flow-based generative models, with NN parameterized vector field it.
inf E, Paata (mn (m. ||ug (a,;; ; 0) - 4e (z: x11) x=)11²
0
where 24(0, 1]. Pdata (x1), and now æ; Pt (æ; x1). We refer to this objective Flow Matching (FM). Flow matching resembles diffusion model with score matching except that the steps of noising (with conditional vector field ua(2r (211 )) and denoising (with marginal vector field et) are deterministic. Solving Eq. 2 forward in time allows for generation from the learnt model.
The family of conditional vector field u, that governs the conditional probability path P. (2, ax) is design choice. Lipman et al. (2022) considered particular example of the conditional probability path, Pt (xx [x)) = ((r or ( I where f4, and a, are time-dependent functions, with end and (x)
such that (x) d (xx; Ia). Therefore, the probability path = Pr is induced by the map (x)) (e(1)) + ( )xx which is the solution of the characteristic ODE d
points Po( (x1)
Ze (x)) = ur (x))(())) A special example includes linear interpolation in the Wasserstein space, (x:) (1 t)ar + tar. For this choice, the corresponding conditional vector
field is uz (xxx))) 1-t
xt for t E [0.
3 STRUCTURED FLOW AUTOENCODERS
In this section, we augment probabilistic graphical models with CNF likelihoods to design structured flow autoencoders (SFAs), family of structured flow-based probabilistic generative models.
From marginal vector field to conditional vector field. To enable probabilistic graphical modeling using flow-based models, we rely on key insight arising from Bayes formula: the marginal vector field can be equivalently derived as the expectation of conditional vector field U2 (x(x)) (z) over an unobserved latent variable z,
vx(x) = [<] pr(m)z)p.(z) dz EEp, (<)===/)))) (x] (=))
(5)
which also resembles the posterior predictive distribution. We formally state this result below in Theorem 3.1, which shows that E. (v) (xx(=)) is indeed the vector field that generates the path of
4

---

pixia)
etz (x)
1-0
adi
dit ag)
(a) SFA framework.
(b) Sampled trajectories {a+} (top) and {
(bottom).
Figure Overview of Structured Flow Autoencoders (SFAs). (a) SFA framework showing co-evolving conditional probabilities for latent and observed æ with conditional CNFs. In the SCFM objective. we compute
a convolution of conditional vector field e for æ((:, zi) with respect to gi gz(z (z)) when æ; = (1 - () + ta For conditional generation, sampling follows
p1 (æt 21]; latent representation involves 21 Qı (z) (= (b) 1D PCA projected sampled trajectories MNIST. Upon conditioning 21, paths of
Pi (z1),
different sampled digits do not cross for t E [0, 1). confirming that the Iearned posterior encodes low-dimensional structure.
marginal probability distributions pr (z:) The proof is in Appendix A.1, which proceeds by verifying that Eq. 5 satisfies the continuity equation (Lipman al., 2022).
Theorem 3.1. Given conditional vector field (x 2) that generates the path [pe(z]z)} of prob- ability kernel for p(z) a.e. 2. Then (x(x) (= lu (x]22] is the marginal vector field that generates the marginal probability path pe(æ)) = f pt (mx 2) p(z) dz over E [0. 1] under regularity
conditions. Conversely, the marginal probability path (r(x) induced by De vx(x)) via the continuity equation coincides with f pe(alz) p(z)(d)) dz.
Theorem 3.1 allows to gain flexibility and interpretability in flow-based generative modeling by introducing latent structure to the otherwise marginal vector field data distribution. This realization is the key to uncovering rich latent structure while ensuring marginal distribution is captured faithfully.
Structured Flow Autoencoders (SFA). This motivates us to design structured flow autoencoders (SFAs). This family of probabilistic autoencoders consists of co-evolving probability paths (p, (-12,) and p+(-)a)) across time E [0, 1); and these paths are connected to the observed data distribution []] (Fig. 3a). At f 1, the probability pı (-1m2) and p1 (-12)) cor- responds to the model likelihood and posterior for the observed data; at f = O, the probabilities correspond the marginal base distributions that are easy sample and evaluate. In addition, Theorem 3.1 allows the specification of any posterior family, which could have inbuilt structure according to a graphical model. We defer three representative examples the next subsections.
Pa(x)) through (x.(*l=))(: (v: (x)=)
Structured Conditional Flow Matching (SCFM). To learn SFAs, we match the marginal path (e(2,) induced from (z)m)) [== a preselected path as in FM objective, and optimize for the conditional vector field v; and posterior approximation q E Q= = {(t, x) + qe(zlæ), (t,z))
[0, 1] x X). We formalize this approach through the Structured Conditional Flow Matching (SCFM) objective:
R(0,g) EE
us (x, x1)
(6)
me ())p((m( (0,
where the outer expectation is w.r.t. the conditional flow trajectory given observed samples στ; the inner expectation is w.r.t. the trajectory of the posterior distribution; the reference vector field u, is chosen a priori, which defines the desired trajectory connecting observed data to the base distribution
5

---

Po- Intuitively, SCFM solves "de-mixing" problem, decomposing the observed signal to (1) the data generation model and (2) the latent structure components.
Posterior approximation. The choice of approximating family Q must be sufficiently expressive to capture the complexity of the true posterior, while not so complex as to destabilize the training. We discuss specific choices for each latent structure in Fig. 2 in the following subsection. Now, with learned posterior approximation, the corresponding marginal distribution (prior) in the latent space can be derived post-hoc. Motivated by empirical Bayes, this can be achieved via integration over the observation marginal: gr (z+) (r(2; æz ) xx, )dz,, In practice, separate model can be used to learn the marginal after training (Wang al., 2023; Preechakul al., 2022).
VAE vs SFA. SFA shares the motivation of structured VAEs for learning probabilistic encoders and decoders, but replaces the likelihood- of marginal probability paths (Theorem 3.1). First, while CNFs can parameterize VAE encoders and decoders, doing so requires expensive likelihood evaluation every training step; SFA avoids this entirely. Second, 3- (Higgins et al., 2017) introduces regularization parameter to balance the reconstruction and KL terms, trading off density estimation quality for p( [X (z) against p(z) x). SFA requires no such trade-off: the latent z captures meaningful structure in the conditional path Pt. (x, 2z,, which simultaneously stabilizes training and ensures accurate marginal reconstruction. Fig. 3b provides a visualization samples of 2, and 2; trajectory from MNIST dataset.
ELBO with an objective grounded in the consistency
For different graphical models (cf. Fig. 2), SCFM objective can be adapted to accommodate their specific structures. We illustrate through three examples in the following subsections, spanning continuous, finite mixture and Markov dynamic latent structure. We chose these three examples because they are (1) widely applicable across different domains, (2) representative of different dependency types (continuous, finite mixture, temporal), and (3) sufficient to demonstrate the framework's 's flexibility.
3.1 CONTINUOUS LATENT VARIABLE MODEL
Consider the graphical model in Fig. 2a, where R2 and E IRd, 0 < p < d.
std p(2),
2z ind p(æl2))
2,
giving rise to the posterior 2//2; ~ p-p2 )) We estimate both the unknown likelihood and posterior from observed data, 31,i ~ Pdata (z) = J (=)=)p(=)dd under SCFM objective. Following from Thcorem 3.1, the likelihood model is the conditional CNF generated by the conditional vector field ut (x(2 0):
dt (x()) = =
(=))=))) Φ0(x)
Zo.
Do Po(())
In this example, the risk function follows directly from Eq. 6. For practical implementation, the mixed conditional vector field (2, j (mx( ( as commonly used in VAE (Kingma & Welling, 2013). The approximation family Q for the posterior path {ge (z x)) bx))ej0,1 admits different choices, each defining valid absolutely continuous curve in (P2 (IR)). W2).
can be approximated with single sample E ~ q. (z|1r)
Conditional CNF. One can parameterize (x((x))) via a conditional vector field re(z f), defining the posterior path as the pushforward base distribution (z(z) through the flow ODE s [0, t):
d
(2) = (z) 0); (ol(z) zo: ds
Zo (z)).
This is maximally expressive. The path can traverse all of Pa (RP), but it requires backpropagation through adjoint ODE steps (Chen et al., 2018) to evaluate the inner expectation, adding computational cost and potential training instability.
Gaussian family. Alternatively, Q can be a parametric family indexed by and x : Q = ((t,x) + N (Ho x), Et,t, (t)), which parameterizes path in the Bures- wrres-sesersttein submanifold. Gradients follow directly from the reparameterization trick, offering computational efficiency and training stability.

---

3.2 LATENT FINITE MIXTURE MODEL
In this section, we consider the generative model in Fig. 2b, where the latent variable follows finite mixture distribution with K classes. This graphical model takes into account the latent class for each sample a. It gives rise to posteriors on the local class label εt. the contimuous latent 2, and the global class proportion
& [K], where p(E, k
Generative Model
Inference Model
ι p>(( 6i/T sid
Cat(TT)
f πT ~ Cat(p(C: r)), 2,|m;, & p(z; n. (t), x|2 p(T[2p)))
2: p(z(Et), 2z ind p(z)=,),
When Dar(a), the posterior for overall proportions p( (11) has closed form Dir(ā) with ag ag + Σ 1fc k). The local label ç, are of major interest for drawing inference on the latent class assignment and value. Next, we adapt SCFM for latent finite mixture model: both & and now integrated out in the inner expectation. Applying Theorem 3.1 gives rise to Eq. 8.
inf rEQ.EEO
Baippasta (a1)
(E (() (ar 0)] Ua (az, x11) ||
(8)
(x(()) 1..nff(011
The design of likelihood model follows similarly in Eq. 7, which is CNF conditioned only, The approximation family for p,,) could be chosen as a Gumbel-Softmax distribution with time-dependent parameters, alternatively constant across f reduce the complexity of the model. The approximation family for pe (z|xx E), should be chosen as conditional CNF or parametric distribution indexed by f, æ,, As is unconstrained, Gaussian approximation family can be posited similarly as in Section 3.1.
3.3 LATENT DYNAMIC SYSTEM
We consider discrete-timg time sequential generation with continuous latent states following the graphical model in Fig. 2c. For index [S], the generative model is formalized as conditionally independent observations x given latent state z"; the inference model is focused on the full posterior of latent which can be factorized into full conditionals at
trajectory z1 given the observation sequence :s1, each index given all previous history.
Generative Model
z z p(2[2"=1),
2/|= p(#= (z()
Inference Model
z [20)
|S| [st p(z |2
II +E[S]
p(= [---1] x (51).
Given observed sample sequences {z.
[5] |5 | the SCFM objective can be shown to have the form Eq. 9. Accompanying theoretical results and numerical derivations are detailed Theorem A.2, extending Theorem 3.1.
inf gEQ,0E0
Ex
patato me (am Unif
Σ IE (a [0r(<;]2]; ue (xf x #E[S]
Here the sum over s E [S] captures the sequential dependencies inherent in the model structure. Following the assumed conditional independence, parameterize the likelihood using conditional
CNF' for each æ² E [S] according to Eq. 7. To approximate the posterior Po |2(2151|2215 [1[)), we employ
parametric family with its parameters indexed by [0, 1], previous states ==s-1), and the full observation sequence !S] Details of implementations can be found in App. B.4.
4 EMPIRICAL STUDIES
In this section, evaluate the proposed SFA across range of tasks and data modalities: (a) conditional density estimation for Pinwheel dataset; (b) latent clustering on MNIST dataset; (c) gene expression modelling single-cell RNA-seq data; (d) sequence modelling with Pendulum trajectory
7

---

Table 2: Comparison of metries for MNIST dataset between VAE, latent FM and SFA. Evaluated on a held-out set of size 1000. The OOD dataset consists of first 10 classes of letters and the first 10 classes of digits in EMNIST. The clustering done in the latent space via k-means with given.
pa122) f
log p(=(2) t
Vendi f
SSIM
NMI (OOD)
ARI(OOD)
VAE
-453.645
-85.448
63.286
0.419
0.039(0.033)
0.017(0.012)
Vamp VAE
-584.845
155.820
1.140
0.866
0.006(0.006)
0.000(0.000)
Latent FM
8.380
0.980
0.488(0.392)
0.381
Latent FM (VAE)
910.925
-11.192
19.631
0.697
0.309(0.152)
0.205(0.073)
SFA
916.901
793.262
25.589
0.716
0.490(0.394)
0.356
w/Deterministic
858.385
10.189
0.732
0.501(0.333)
0.379(0.155)
w/CNF Posterior
-907.998
351.141
23.166
0.654
0.485(0.325)
0.355(0.118)
video dataset. We compare our method mainly to VAE counterparts, including SVAE (Johnson et al., 2016), VAm (Tomczak & Welling, 2018), 3-VAE (Higgins al., 2017). Gaussian Mixture VAE (GMVAF) (Lin et al., 2018) and I Latent Flow Matching with deterministic autoencoders (LatentFM) (Dao et al., 2023) and probabilistic autoencoders MLatentFM w/VAF). In the following, we use Mixture-SVAF SVA refer to the SVAE with Gaussian mixture latent, and equivalently GMVAE; use GLD-SVAE to refer to the Gaussian Latent dynamics SVAE.
For fair comparison, we restrict the prior to be fixed in SVAE and only focus modeling the conditional probabilities. All the metrics are evaluated based on samples held out from training. More experimental details and supporting visualizations are hosted in App. B. We draw comparisons in the quality of (1) generated posterior samples (or latent distribution samples for Latent FM) and (2) generated posterior predictive samples, and downstream tasks, such latent space clustering.
For (1), given a large training dataset, the posterior distribution should be able to provide an accurate representation of the latent distribution. Therefore, under simulation settings where latent ground truth is known, we are able to evaluate the discrepancy between learned latent representation and the truth, provided
qi (z (a)) with P
For (2), we conduct posterior predictive check to evaluate the discrepancy of the samples from Pdota (x) and Ppred x) = P1 (क) 2 [=z== d)).. To sample from the latter, we follow Pdata (æ), 2|2r (=((rr,, a|2 (arlz).
We evaluate the diversity of the generated samples using Vendi score (Friedman & Dieng, 2022); quality of image generation with SSIM (Wang et al., 2004); quality of latent clustering with ARI (Hubert & Arabie, 1985), NMI (Strehl & Ghosh, 2002) and probabilistic version softNMI (Eq. 16).
Summary of findings. In conditional density modeling (Pinwheel), SFA consistently outperforms LatentFM and VAE-based models, showing better data density reconstruction and better latent space modelling. To assess scalability, we apply SFA to a single-cell RNA-seq dataset, where it effectively models high-dimensional gene expression data and outperforms VAEs in reconstruction quality, On image data (MNIST), both SFA and its mixture extension (Mixture-SFA. learn meaningful latent representations, generate high-fidelity samples, and perform well on latent-space clustering tasks. Finally, we highlight SFA's versatility on sequential data using the pendulum video dataset, where it successfully captures the oiw-dinensional periodic structure the underlying physical system.
Comparing posterior family. When the latent is diwer-dimensional, smaller posterior model is sufficient relative to the model size needed for the likelihood. When learning multiple components jointly, simpler parametric approximation families are preferred over conditional CNFs for training stability. Computation- wise, SFA (2.4M parameters) requires 13.220 1.848 seconds per epoch, comparable to VAE's 12. 2.01 seconds. With CNF the posterior of SFA, we observe drastic increase in the runtime to 167 460 + 176.817 seconds. This is due to the additional time arising from sampling CNF during training, which requires solving ODE each gradient evaluation step.
4.1 MIXTURE MODELING: PINWHEEL DATA
We first illustrate the ability of SFA in learning conditional distributions using the toy example of the pinwheel dataset, with five clusters each having the shape of blade (Johnson al., 2016). The
8

---

Table 3: Subspace clustering on MNIST latent mixtures models Mixture- SVAE and Mixture-SFA
Evaluated on a hutd-out set of size 1000.
p<)222 t
log (g)2())))
SSIM 1
softNMI f
NMI 1
ARI
Mixture-SVAr -1133
0.667 725
0.634
0.698
0.161
0.072
Mixture-SEF -FA
-906
0.779
0.728
0.489
0.332
6
9
0
/
2
3
6
9
O
2
3
4
5
6
?
8
9
0
/
2
3
4
S
7
g
9
(a) Mixturc-SVAL
(b) Mixture-SFA
Figure 4: Comparison of Mixture-SVV -VAL and Mixture-SFA on MNIST dataset. The first row displays the posterior predictive æt pi (ı(a))) (z|æ1.) Juz and latent class assignment probability & (g)m1,)) for data ø1,6. The second row shows the latent representation learned for digits, where each point is sampled from 2, qi (z()),)) for test data sample 11,1-
class membership is not provided during training. The goal is to evaluate if the posterior is able to uncover the latent structure of the data, and whether the model is able to capture the observed data distribution. In Fig. 1, we visualize the generated data together with their representation coded in 1D colorbar. SFA able to reconstruct the support of the ground truth distribution, in addition to capturing a meaningful latent representation for the angular rotation. In contrast, both the Latent FM and VAE-estimated density do not have well-separated components. As shown in Table I, SFA based methods achieve similar density estimation quality FM and comparable to ground truth, while SVAE based methods fail to model the density accurately.
4.2 IMAGE MODELING: MNIST DATA
Next, consider MNIST dataset LeCun al. (2010). We aim to recover the probabilistic assignment of each image to the 10 classes, and learn a low dimensional feature representation at the same time. We first consider the graphical model with continuous latent, and compare SEA VAE and Latent FM on the latent space clustering task. In addition, to check the learned latent space captures desirable structures in the data, such as stroke and abstract shape. sample from the latent distribution encoding Out-of- offDissriibution (OOD) data in the EMNIST dataset Cohen et al. (2017). Results and comparisons are summarized in Table 2. Visualizations are presented in Fig. Owing to posterior collapse, VAE learns an unstructured latent space with close to random samples (high Vendi score). SFA achieves the best balance of diversity (Vendi=25.6) and clustering (NM1=0.490) among all methods, while Latent FM trades off diversity for reconstruction (Vendi=8.4 4 but SSIM=0.980), Performance using mixture graphical model are organized Table 3 and Fig. 4. MixtureSVVAE improves generation quality and latent class separation over VAE. However, Mixture-SFA still achieves better clustering quality.
4.3 GENE EXPRESSION MODELING: SINGLE-CELL RNA-SEQ DATA
The dataset obtained from Lotfollahi et al. (2023) includes PBMCs from eight patients with Lupus. The data consists 7 cell types, and treated and control with IFN-3 (Kang et al., 2018). The observed count is normalized and log(x + 1) transformed, then 5, 000 HVGs are selected. We apply continuous latent to learn the low-dim representation of the high-dimensionn shows that both SFA and Latent FM produce meaningful clusters of the cell type in the latent space. SEA matches Latent FM in clustering accuracy while producing substantially more diverse samples, indicated by the higher Vendi score and clustering scores (Table 4a).
differential expression data. Fig.

---

Table Comparison of metries across different datasets and methods. (a) Kang HVG dataset evaluated on a held-out set of size 500. The observation has dimension 5000, due to the size, the log likelihood for CNF cannot be directly computed by solving adjoint- ODE, therefore left out of the comparison. (b) Pendulum dataset over posterior samples of observed (RMSEg), and latent (RMS kig) Evaluated on a held-out set of size 300.
(a) HVG Single-Cell RNA-Seq
(b) Pendulum Trajectory
log 1o/p()(z) f
Vendi (a) t
NMI t
ARI 1
RMSE, +
RMSE; +
VAE
40.04
26.58
0.412
0.257
LatentFM
5.801
0.617
0.457
GLD-SVAE LDS-SFA
4.574 3.233
8.090
SFA
384.1
737.7
0.633
0.460
1.526
4.4
SEQUENTIAL MODELING: PENDULUM TRAJECTORY VIDEO DATA
We choose pendulum trajectory dataset for LDS example. The dynamics is driven by the pendulum physical system modeled as a damped harmonic oscillator. The latent trajectory 2 dimensional, consisting of angle and angular velocity. The observation is video with discrete time frames mapped from latent trajectory. We compare SVAE with SFA in Table 4b, where we measure the discrepancy between the generated and ground truth of both observed and latent dynamics using RMSE. LDSSSFA reduces latent RMSE by over compared to GLD-SVAE (1.526 vs 8.090), suggesting much better recovery of the underlying dynamics. Additional details on model implementations are in App. B.4.
DISCUSSION
In this work, we introduced structured flow autoencoders (SFA), a framework that integrates flow matching (FM) with latent probabilistic graphical models (PGMs) to achieve both high-fidelity generation and structured latent representation learning, leveraging the consistency of marginal probability paths (Theorem 3.1). At the core of SFA is the structured conditional flow matching (SCFM) objective, which extends flow matching by explicitly modeling latent variables, enabling joint estimation of the generative likelihood p(x =) and the posterior p(z|z)) ) A key contribution of this work is the generality of the SFA framework: SCFM provides principled recipe for augmenting any graphical model that specifies conditional independence structure with wew-matching-bare based likelihoods. We instantiated this for continuous, finite mixture, and dynamical latent structures, and demonstrated consistent improvements over VAE-based and latent flow matching baselines across synthetic density estimation, image clustering. high-dimensional gene expression modeling, and sequential video generation in both generation quality and representation utility.
Limitations and future work. SFA with parametric posteriors is computationally comparable to VAEs, but the CNF posterior variant remains expensive due to ODE solves at each gradient step. In practice, simpler families suffice when the latent dimension is much smaller than the observation dimension; principled selection strategies remain open question. In addition, our contribution is primarily methodological: formulating how to compose representation learning with flow matching via SCFM. Scaling complex image datasets poses separate, architectural challenge: designing UNet-based decoders that meaningfully condition on a simultancously learned stochastic latent, rather than bypassing it through skip connections. Standard conditioning mechanisms such as FiLM can be insufficient when the latent carries distributional uncertainty, risking the decoder ignoring the latent entirely. Addressing this architectural challenge a natural next step that complements the methodological framework introduced here.
ACKNOWLEDGMENTS
This work was supported in part by funding from the Office of Naval Research under grant N00014 14- 13-10, the National Science Foundation under grant No. 2310831, No. 2428059, No. 2435696, No. 2440954, a Michigan Institute for Data Science Propelling Original Data Science (PODS) grant, Two Sigma Investments LP, and LG Management Development Institute AI Research. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the sponsors.
10

---

REFERENCES
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, arXiv preprint arXiv:2107.03000 2021.
Mohammad Bashiri, Edgar Walker, Konstantin-Klemens Klemens Lurz, Akshay Jagadish, Taliah Muhammad, Zhiwei Ding, Zhuokun Ding, Andreas Tolias, and Fabian Sinz. A flow-basee latent state gener- ative model of neural population responses to natural images. Advances in Neural Information Processing Systems, 14:15001-15815, 2021.
Jianfei Chen, Cheng Lu, Biqi Chenli, Jun Zhu, and Tian Tian. Vflow: More expressive generative flows with variational data augmentation. In International Conference on Machine Learning, pp. 1660-1669. PMLR, 2020.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances neural information processing systems, 31, 2018.
Siyi Chen, Huijie Zhang, Minzhe Guo, Yifu Lu, Peng Wang, and Qing Qu. Exploring low-dimensional subspaces in diffusion models for controllable image editing. arXiv preprint arXiv:2409.02374, 2024.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and André van Schaik. Emnist: extension of mnist to handwritten letters. arXiv preprint
2017.
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran. Flow matching in latent space. arXiv preprint arXiv:2307.08698, 2023.
Aram Davtyan, Sepehr Sameni, and Paolo Favaro. Efficient video prediction via sparsely conditioned flow matching, In Proceedings of the IEEE/CVF International Conference Computer Vision, pp. 23263-232744 2023.
Dan Friedman and Adji Bousso Dieng. The vendi score: diversity evaluation metric for machine learning. arXiv preprint arXiv:2210.024110, 2022.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint arXiv:2407.15595, 2024.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367. 2018.
Irina Higgins, Loic Matthey, Arka Pal, Christopher P Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. ICLR (Poster), 3, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 1866840-68511 2020.
Lawrence Hubert and Phipps Arabie. Comparing partitions. Journal of classification, 2(1):193-218 1985.
Noboru Isobe, Masanori Koyama, Kohei Hayashi, and Kenji Fukumizu. Extended flow match- ing: a method of conditional generation with generalized continuity equation. arXiv preprint arXiv:2402.18839, 2024.
Matthew J Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Composing graphical models with neural networks for structured representations and fast inference. Advances in neural information processing systems, 29, 2016
11

---

Hyun Min Kang, Meena Subramaniam, Sasha Targ. Michelle Nguyen, Lenka Maliskova, Elizabeth McCarthy, Eunice Wan, Simon Wong, Lauren Bymees, Cristina Lanata, et al. Multiplexed droplet single-cell raauennnnn using natural genetic variation. Nature biotechnology. 36(1): 89-94, 2018.
Jaivardhan Kapoor, Auguste Schulz, Julius Vetter, Felix Pei, Richard Gao, and Jakob H Macke. Latent diffusion for neural spiking data. arXiv preprint arXiv:2407.08751, 2024.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. CoRR, abs/1312.611ll. 2013. URL https::tp//.. /tpp. www.steisccooooeeeeee.... org/ Corpus 21780090.
Yann LeCun, Corinna Cortes, and Christopher JC Burges, Mnist handwritten digit database. http: //yann. lecun. cww/ wwo// mnist/, 2010.
Gen Li and Yuling Yan. Adapting unknown low-dimensional structures in score-based diffusion models. arXiv preprint 248.: 2024.
Wu Lin, Nicolas Hubacher, and Mohammad Emtiyaz Khan. Variational message passing with structured inference networks. arXiv preprint
005505 2018.
Yaron Lipman, Ricky TQ Chen, Heli Ben- HamuH Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXhv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2000220090,00,, 2022.
Mohammad Lotfollahi, Anna Klimovskaia Susmelj. Carlo De Donno, Leon Hetzel, Yuge Ji, Ig- nacio I. Ibarra, Sanjay Srivatsan, Mohsen Naghipourfar, Riza M Daza, Beth Martin, Jay Shendure, Jose L McFaline -FiFalineFFiucraa, Pierre Boyeau, F Alexander Wolf, Nafissa Yakubova, Stephan Günnemann, Cole Trapnell, David Lopez-Paz, and Fabian Theis. Predicting cel- lular responses to complex perturbations in high-throughput screens. Molecular Systems Bi- htt.:/o....gr...1122///momo..m.
ology, 19(6)e11517, 2023. doi: h/www. www.peebbssssaa..mm erg/doi/ abs /10 / /ws.//s../////. 200211517. 2
URL https:
Sarthak Mittal, Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Dif- fusion based representation learning. International Conference on Machine Learning, pp. 24963-24982. PMLR, 2023.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint (0Mi::20220066:1: 2021.
Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn. Dif- fusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 10619-10629, 2022.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502. 2020.
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415-1428, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1:Xiv:2011.134555 b021b.
Alexander Strehl and Joydeep Ghosh. Cluster ensembles knowledge reuse framework for combin- ing multiple partitions. Journal of machine learning research, 3(Dec):583-617, 2002.
Jakub Tomezak and Max Welling. Vae with a vampprior. International conference on artificial intelligence and statisfies, pp. 1214-1223, PMLR, 2018.
12

---

Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835- 8856. URL https::t///....... opppreveeeeeeeeeeeeee eww.. net/ forum? d=CDnnnc 3AW. Expert Certification.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in neural information processing systems, 34:1 1287-1130. 2021.
Binxu Wang and John J Vastola. Diffusion models generate images like painters: an analytical theory of outline first, details later. arXiv preprint arXiv:2303.024900 2023.
Peng Wang. Huijie Zhang, Zekai Zhang, Siyi Chen, Yi Ma, and Qing Qu. Diffusion models learn low-dimensional distributions via subspace clustering. arXiv preprint arXiv:2409.02426, 2024.
Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing diffusion models. In International conference on machine learning, pp. 36336-36354, PMLR, 2023.
Zhou Wang, Alan Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from eπror visibility to structural similarity. IEEE transactions image processing. 13(4):600-612, 2004.
Minkai Xu, Alexander S Powers, Ron O Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. In International Conference Machine Learning. pp. 38592-38610 10. PMLR, 2023.
13

---

A THEORETICAL DETAILS
We formally present the probabilistic representation of solutions to the continuity equation when the vector field fails to be Lipschitz w.r.t In this case, the solution the characteristic ODE (flow ODE) is not unique. When using neural nets to parameterize the vector field, we want to verify that the solution of the ODE indeed induces solution the continuity equation.
Firstly, we denote the space of absolutely continuous curves 7 [0,
[0, 1] + P(RR)) as the path of probability indexed by L, ACP(O, 0, 1; Rd) as
IR with finite p energy, i.e. 1 (0, 1). Denote I' as the space of continuous map y [0, 1] R. Let e, (x, + y(t) as the evaluation map. Then define the curve of probability measure induced by the evaluation map as
1e exty, t E [0.
where by definition,
I ((x)dup(())
f ((((t))dn(z;))) r), vi E (R te [0,
R
Then finally recall that the continuity equation
0.44 V (utp)) =0 in Rd (0,1).
Theorem A.1 (Ambrosio et al. (2008) Theorem 8.2.1). Let μ [0, + ((TR)) be a narrowly continuous solution of the continuity equation for suitable Borel vector field vn (x) such that for some > 1,
[. (x))))d < +00.
i Then (a) there exists a probability measure n in R& x such that that concentrates on the set of pairs (x, 7) such that y E (O, is solution of the ODE +(t)) (x(())) for L--a.e. E [0, 1] with ( a.
and (b) μr 1e 1L [0,
ii Conversely, any satisfies (a) and fo Jgaxr (d)((()))d)(zz, 7) < +oc induces solution of the continuity equation via μº entn with μiu coiy.
The converse argument can be readily extended to the conditional vector field v(., 2) [0. x 4 Rdd for any v-a.e. z. Then the solution curve the characteristic ODE would be indexed by z, denoted as 7z.
A.1 PROOFS IN 3
In the following, we use the same notation the main paper for the proof details $ 3. Recall the marginal vector field defined via the posterior expectation (Eq. 5):
(x(x)))
Uy a
(z(()) (e(x)z')p(2")) dz dz (pp(:ta))(((((>>))
The correct notion of distance between velocity fields, consistent with the Wasserstein-2 geometry of Ambrosio et al. (2008) (Chapter 8), is the kinetic energy norm:
f (w)(())))) (2) (x) dx dt,
which appears as the action functional in the Seenmou--Brenie formula for W2.
We now make precise the sense in which the conditional vector fields {ve(-12)}= give rise consistent marginal dynamics, restating the theorem for completeness.
Theorem (Theorem 3.1). Given conditional vector field v, (x) (z) that generates the path {pr (a = of probability kernel for p(z) a.e. Then vx(x)) = (pp(( [v,(a] (œ]2)] is the marginal vector field that
generates the marginal probability path pe(x) = f px (x) 2)) p(2) dz over [0, 1] under regularity
conditions. Conversely, the marginal probability path p. (x) induced by v, (x) via the continuity equation coincides with (=(22)) p(2) dz.
14

---

Proof. The proof proceeds in two steps: (i) the mixed vector field solves the marginal continuity equation, and (ii) uniqueness of the induced probability path.
Regularity conditions. Assume:
(R1) vz(2(2) is Borel measurable in (x, () and continuous in with So [+(())))))))) (z(x))) ((()))) dæ dt < oo for p(z)-a.e. 2.
(R2) v2(x) (e and V
are continuous and bounded in (ax, t, z).
(R3) p: (z) and at (a( are continuous in z). and de (e (a:(2) is uniformly bounded for all i E [0, 1J and p(z)
These conditions ensure the Leibniz integral rule applies, justifying the exchange of d and V with J [P(2)) dz.
de
By assumption, for p(z)-a.e. 2, the conditional pair (vr (-12), (c(+12) satisfies
d
हi P.(x12)) - (v.( z)) (e z)).
(10)
Integrating against p(z) and applying Leibniz (justified by (R2)-(R3)):
/ d p())
dl Pr
/ V (1 p(z) dz
= Vz / px
(z| p(z) dz.
(11)
Now we rewrite the integrand by multiplying and dividing by (x():
/ u p(z) dz = / v1(x)) x) P: (z(z)2) p(z)
Pc(x) dz
(x(x)
(/ ve (e(l)) (c(x))
(p (x)
pr
(12)
where the second line uses Bayes' rule (=(x)) (x(2) (((z)=)p()))))))) p(z)/p((x) and the last line is the definition of (x)). Substituting Eq. 12 into Eq. 11:
p((x) = Va (v((x) (z)).
confirming that (Ve, Pt) solves the continuity equation.
For the converse, verify that pa( is the unique solution induced by er. Under (R1), the integrability condition required by Theorem 8.2. Ambrosio et al. (2008) holds:
[' (0(a))) (x()) dx dt < [/ [x(()/))] (=)2)) p(z) dz dx dt +00,
where the first inequality is Jensen's 's applied to 1 under (x(x). The superposition principle then guarantees a probabilistic representation p: = extn. Uniqueness of the solution to the continuity equation follows from Lemma 8. 1.4 of Ambrosio et al. (2008) under the imposed regularity et- Hence the marginal path induced by i must coincide with J p. 22) p(z) dz.
This establishes the full consistency: the conditional vector fields {e, (+1=)]= uniquely determine both the marginal probability path (z) and the marginal velocity field v, (x)) = Ep [ux(x) and these are consistent solution pair of the continuity equation.
15

---

Remark (Posterior collapse). The collapsed posterior qt(z x) p(z) is special case of any distributional family Qain, so i R(8, q) I inf, R(8, :optimizing over 9 can only improve upon the collapsed solution. Unlike the VAE objective, R(0, q) contains regularization term (such as (k)( that penalizes the posterior for deviating from the prior, so there is no force driving q toward collapse during optimization.
A.2 PROOFS IN $ 3.3
Recall the posterior arising from the latent dynamic model is
~ p(z x (5)) = II
p(=11210 11-11 11)).
(13)
+E[S]
We present the extension of Theorem 3.1 to the latent dynamic model in the following theorem. The proof idea relies verifying the joint continuity equation over the trajectory is satisfied and the corresponding structured conditional flow matching objective is well defined.
Theorem A.2. With conditional flow defined by Eq. 7, and posterior defined by Eq. 13, the FM objective is derived to be Eq. 9, which has the same gradient as the flow matching objective that matches Ut the marginal vector field ur.
LSCFM Eeg, Paata perpp (ee, (2es) Uniff(.11]
Σ .E(S]
E (= (a)s) [1a(0]]= ue (ar acf)
Proof. Assume regularity conditions that guarantee the exchange of integration and divergence, differentiation w.r.t. t.
We first show that f Evess] ue we(a,, lc:)) mss)dda is the marginalized vector field that gener- ates {pr (. In the following p(æ} (2)) = p(æ+) for simplicity of indexing.
d dt (p(s)
/
d
dt Dc (mls)
|atss | dar(s)
/
Σ
(ax xt)
II
pe p(ats) days))
Pe
SE|S)
J
(p)
II (z )p(
sE(S]
(14)
Σ /
ue (xx |a) js
sE[S]
ue ue (x)) pe po(erss)
Σ #E(S)
The second equality is by conditional independence of the transported samples for each E [S] and applying chain rule on the product (=(xx==)) [I (x.)..). This shows the marginal vector field additive in the time index s following the marginal vector field defined for each s E [S].
Now, we'd like to derive the conditional flow matching objective from the marginal flow matching, and show the two has the same gradient with respect to the NN parameterized marginal vector field 2t. The marginal VF for LDS takes the form
Ep, sn) [[(()))))))) (.(m))))1
(15)
16

---

It is then sufficient look at the cross term and the squared term U.. Firstly,
B (x(s)) (vr(x)), 4x(x))
= - (10). Σ
/
(1( )ee (ar ((ss))d pe (2(ss)))
*E(S]
= Σ / (w)(2)), un (xxxa1=)))) / pi (ats) (m)))dd d Pe (da(ss
sE|S]
(2x(,,,, (e Io (x (s ) (er (acpss (ay
sE(s]
(x)), Σ u [x)())))
=B ( [sp
#E[S]
for the quadratic term, it directly follows from iterated expectations
(p (x)] = (a (l (z))l
Therefore optimizing is with Eq. 15 is equivalent to optimizing the marginal flow matching objective
inf E (1 )o(ass))))
ve(x)
ue (xx .E(S]
Σ
Finally, to introduce the structured FM objective with latent dynamical system, we verify the marginal vector field e arisen from marginalizing (z.,)zz.) generates the probability path {pr (2)())))) The proof is similar to the previous ones, where we verify that the continuity equation is satisfied.
d di (a]ss)
/
d di ((
/
Pr (x. z..
dt
SE(S]
jf8
Σ (x.(
II pa
SE[S]
5#
V
p(z(s))
eE(S]
Σ (2( p(a(s)) #E|S|
notable that the conditional independence and Markov assumption gives rise to the filtering probability p(z, ((s) and p(=!)=(s)
In particular, the objective Eq. 9 depends on the entire sequence through the sum over [S]. due to the conditional independence structure of the likelihood. As it requires access to the full observed sequence at every step [S]. the training procedure is entirely offline.
17

---

B EXPERIMENT DETAILS
All experiments are conducted on MacBook Pro equipped with an Apple M2 Pro chip and 16 GB of memory. Code is available https: /ttt::///tttt/ wwwppttww...oeeeeeeeeeeeeeeeeeee www/edeneee../eeee/ooe/oeoeee.ee
B.1 PINWHEEL
The Pinwheel dataset is classic benchmark for density estimation task. SFA is demonstrated to be able to meaningfully capture latent distribution as well as observed distribution.
For VAE, we use Gaussian distribution with diagonal covariance for posterior and likelihood family. The parameters
Gaussian are parameterized by ML.ss and mapped from context vector. For SFA, we use MLP to parameterize the conditional vector fields U/2 and use Gaussian parametric family indexed by f, x for the posterior flow.
With latent mixture, we use constant time Gumbel-Softmax network to model the latent class probability for both Mixture-SVAB and Mixture-SFA. We use tanh activation function for all models applied to this dataset.
To compare, the samples and latent representations are generated based on the estimated conditional probabilities,
2i i.d.d. p(z), 22124 (1()))), 2 f.dd (1(2(x,).
B.2 SINGLE CELL RNA-SFQ
The Single Cell RNA-seq dataset Kang et al. (2018) consists of transformed count vector of size 5000, which presents challenges in modelling with CNF and likelihood based optimization. SFA directly tackles this complexity by learning meaningful latent representation while not requiring computation of log-likelihood during training.
We parameterize the CNF U. with MLP, and uses a 32-dim Gaussian approximation family for the posterior that varies across time and observation æ. The VAE model uses Gaussian encoder and decoder with NN parameterized parameters that is time independent. For the latent FM, we use MLP encoder and decoder, 32 dim latent space and CNF to learn the latent distribution p(z).
B.3 MNIST DATA
For Mixture-SVAE we use Gaussian distribution with diagonal covariance for posterior and likelihood family. MIPs are used to map context vectors to the means and covariances of the Gaussians. The latent class probability via a Gunbbel-Softma network, which uses MLP to map from context vector to logits, then apply Gumbel-Softmax trick for sampling.
For Mixture-SFA, we also use MLP to parameterize the conditional vector fields e and Gaussian posterior with parameter indexed by I, æ. It is notable that with larger differences in the dimensionality and scale, a linear map is used to firstly map the context vectors to vectors of the same size. Then
(a)
(b) Latent
(c) VAE
Figure 5: RNAseq dataset: Latent space visualization in 2D, projected with TSNE (perplexity=30).
18

---

(a) SFA
(b) Latent FM
(c) VAE
(d) SFA- Mixture
(c) SEA (OOD)
(f) Latent FM (OOD)
VAE (OOD)
(h) Mixture-S SVAE
Figure 6: MNIST dataset: Latent space visualization in 2D projected with TSNE (perplexity=50). (a)-(c),(e)(()) follows from the continuous latent graphical model, (d) and (b) employs latent mixture model. The OOD dataset consists of first 10 classes of letters and the first 10 classes of digits in EMNIST.
apply concatenation and feed to the main network. We also use Gumbel-Softmax network model the latent class probability. Alternatively, we can use 10 dimensional vector field to model the distribution logits.
For both model, we use softplus activation function and train until convergence. We observe that smaller network usually suffices for modelling the latent, which also increases training speed.
We compare the performance of the two methods from 2 perspectives.
Posterior Predictive: for every test sample æ; we first sample from the posterior 2:|0; ~ q( ()), then sample from the likelihood æi,newlzz p(x]2z))
2. Latent Space Representation: for every test sample æ, we sample from the posterior 5il;; ~ ((E)=)), then from 2;/x; g(2[xx) to obtain a latent representation of the observed data point, We visualize the class probability samples (&:) with histogram, and the sample {zi} with TSNE (perplexity 30) projected from 64 dimensional space onto 3 dimensional space.
SoftNMI To assess the quality of latent probabilistic cluster assignment for Mixture-SEF SFA and Mixture- SVAE, we use a soft Normalized Mutual Information (softNMI), which computes the discrepancy between one-hot label vector and probability vector based on entropy.
softNMT(p. q) = H(p) +H(g) H(),)) H()) + +H())
E [0, 1).
(16)
where (p)) is the entropy function the marginal, H(p, q) is the entropy on the joint. Higher score suggests higher correlation between the posterior class assignment probability and true class label.
From Table 2, we observe the stochastic latent helps to increase the diversity of the generation, while variants VAE has poor generation quality, the FM based models have better performance image generation quality. However, there's ss notable trade-off in reconstruction quality and generation diversity, where stochastic latent have higher Vendi score while suffer a slight decrease SSIM score, and vice versa for the deterministic latent. In addition, for downstream clustering task, SFA with deterministic latent has better NMI and ARI score for within distribution sample, yet SFA with stochastic latent has better NMI and ART for -ut-o--distributiin samples.
19

---

0
5
9
0
4
7
7
2
9
7
2
2
2
2
2
2
2
2
3
4
3
6
7
7
3
3
3
3
3
3
3
3
4
3
7
9
4
4
4
4
4
4
4
4
5
7
0
6
3
5
5
5
S
5
S
S
S
6
0
0
3
7
9
8
8
6
6
6
6
6
6
6
6
7
í
3
7
4
7
7
7
7
7
7
7
7
8
7
3
9
7
7
8
8
8
8
8
8
8
8
9
3
3
6
9
9
9
9
9
9
(a) VAE Posterior Predictive
(b) SFA Posterior Predictive
0
0
O
0
/
/
/
/
1
1
1
1
2
5
2
2
2
2
2
2
2
2
②
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
5
5
5
5
5
5
5
5
5
5
S
5
S
S
5
5
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
7
7
7
7
7
2
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
(c) Mixture SVAE Posterior Predictive
(d) Mixture- SFA Posterior Predictive
Figure 7: Comparison of Posterior Predictive Results: (a) VAE, (b) SFA, (c) Mixture-SVAM and (d)
MFA.
20

---

-6
(a) Posterior Samples (5)
(h) Posterior Samples (2)
Figure Mixture-SVAi Posterior Samples: (a) samples from latent variable & each column corresponds to different images, each corresponds to different class label; and (b) latent variable z, with TSNE projection from 64 dimensional to 3 dimensional space.
(a) Posterior Samples (E)
(b) Posterior Samples (z)
Figure 9: Mixture-SFA Results: (a) Posterior Samples ( and (b) Posterior Samples (z).
B.4 LATENT DYNAMICAL SYSTEM
To model the full conditional posterior
sequence model to separately encode the observation sequence [S] and full latent history z(s-1| to context vector. The encoder processes each frame of (S) through a shared MIP to produce per-frame embeddings. which are computed once and cached. each autoregressive step 8, a GRU accumulates previously generated latents - into a summary vector, which is used query in cross-attention layer over all S frame embeddings. This allows the model to dynamically weight which frames are most informative for producing (=2,|=|s~1) S1) adapting its attention pattern based on what latents have already been generated. The attended æ-context is concatenated with the z-context and a positional scale embedding. then fed through an MLP to output the mean and variance of a diagonal Gaussian for at s.
Therefore, the full posterior is modeled by
z1 2=12|8- 2 Goussian (μ (1e(2l0- ne (20 E (210-11,,,), h (2((11)))
21

---

aa) Ground Truth
(b) EVAE
(c) LFA-SFA
Figure 10: Comparison of Pendulum Dynamies: Ground Truth, LDS-SFA and GLD-SVAE.
The likelihood, with conditional independence assumption is modeled by Gaussian (p(z,). E(z,)) The likelihood modeled by CNF, where the vector field parameterized with MLP with SiL.U activation, the incorporation of time and latent representation z is via FiLM.
We implement the probabilistic encoder VAE variant similar the SFA counterpart, and outputs mean-field approximation. The probabilistic decoder (likelihood) also shares similar structure, however assumed to be independent Gaussian for each pixel.
C THE USE OF LARGE LANGUAGE MODELS (LLMS)
LLM used to refine the code of model architecture and used to polish the writing of the draft. LLM is not used to generate research ideas or writing to the extent that it could be regarded as contributor.
22


Reviews:
[{"review_id": "TNNMeJ3org", "reviewer": "Reviewer_AJKg", "summary": "The authors propose Structured Flow Autoencoders (SFA), a framework that augments Continuous Normalizing Flows (CNFs) with graphical model structures.  The paper demonstrates the versatility of SFA across multiple settings, including models with continuous and mixture latent variables as well as latent dynamical systems. Experimental results show that SFA achieves superior data fitting performance while preserving meaningful and structured latent representations.", "questions": "1. How does SFA perform in high-dimensional domains, such as image or video modeling?\n\n2.Can the proposed structured latent representations be extended to discrete or hybrid graphical models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is clearly written and well organized, making it easy to follow. The proposed model and its accompanying theoretical analysis are sound and logically developed.", "weaknesses": "Despite its clarity, the paper presents several limitations that the authors should address:\n\n(a) The proposed approach essentially represents a straightforward combination of Flow Matching and Variational Autoencoders (VAEs). As such, the conceptual novelty appears limited, and the contribution may be incremental.\n\n(b) The probabilistic structures considered are restricted to latent variable models and latent chain models. Given this narrow scope, the term structured probabilistic models may be overstated; the work would be more accurately described as focusing on latent models with specific dependencies.\n\n(c) The experimental evaluation lacks results on natural image datasets, which are standard benchmarks for assessing the scalability and expressiveness of flow-based models. Including such experiments—and comparing SFA directly with Flow Matching and VAE baselines—would significantly strengthen the empirical claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cbtO5I9WE8", "reviewer": "Reviewer_EKhV", "summary": "This paper introduces structured flow autoencoders, a method that combines the advantages of flow matching (generating high-resolution diverse data ) and VAEs (learning structured latents for the data).  The idea is to jointly learn the likelihood model p_t(.|z) and the posterior p_t(.|x) using a flow matching objective, where the marginal velocity vector v_t(x) is replaced by an expectation over conditional velocity vectors v_t(x | z) under p_t(z | x).  The parametrization of the posterior  p_t(z | x) can be done either with a CNF or with simple parametric families e.g. gaussians.  They introduce several possible generative processes for x | z, either  continuous latent model, latent mixture model, or latent dynamic system, and such a process is set before training. \nThey run experiments on pinwheel, mnist, rna-seq and pendulum video data, and show that their method shows improvements compared to the baselines both in terms of (i) the structure of the learned latent space and (ii) the likelihood p(x|z) and the posterior p(z | x).", "questions": "Hello authors, \nFirst, I'd like to start by acknowledging that I couldn't spend as much time as I wanted to on your paper to fully understand the theory / baselines, and that there are some parts for which the theory is a bit vague for me, so apologies in advance if some of the questions/claims are evident. \n\n- Figure 1 : I don’t fully understand this figure. You train without the color and you want the learned latent variable to be the color? \n- For MNIST / Pinwheel: What is the structure of sampled the latent variable in this example? is it  a predicted  (mixture of ) mean and variance, and I guess they are one-dimensional here? If not what are their dimensions?\n- What do the intermediate-time latents z_t  look like, e.g. for MNIST or for pinwheel data ? are they interpretable? \n- In the paragraph \"SVAE and SFA comparison\", you claim that SFA has an advantage by jointly learning p(x|z) and p(z|x) , whereas SVAE has a \"generation-latent learning trade-off\". I don't understand this, could you elaborate more on it and why it is a drawback compared to SFA? \n- Could you explain the advantage of your method compared to latent flow matching with structured autoencoders? What is the advantage of also modeling z_t with time, instead of just learning the encoder/decoder for clean data as in latentFM?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- very nice and original idea to learn the latent structure with a flow matching objective instead of a VAE and theoretically grounded with the conditional vector field equation,\n- Clarity and mathematical / scientifc rigor of of the paper,\n- Solid and diverse set of experiments and of evidence,  and diverse set of baselines , nice 2D visualisations.", "weaknesses": "- Regarding the novelty of the eq.5, used in proposition 3.1 : this conditional vector field equation was introduced before in https://arxiv.org/pdf/2302.00482v1 eq 9 (however they are not conditioning on latents, but on the unnoised sample x_1)\n\n-  Regarding the clarity of the toy experiments: Although the theory for the SFM equation is clear, I wasn't able to understand the nature of the learned latent variables for the pinwheel dataset and MNIST dataset", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CrKBVroc0y", "reviewer": "Reviewer_KF7s", "summary": "The authors develop a structured conditional flow matching objective which is used to train their Structured Flow Autoencoder architecture. Likelihoods are similar to conditional continuous normalizing flow frameworks. Their aim is to achieve similar quality of samples to standard flow matching while retaining latent structure. They do this by jointly learning the conditional continuous normalizing flow likelihood and the approximate posterior.", "questions": "The log likelihood in Table 2 and Table 3 for SFA has value 725.232 and 384.137, which looks surprisingly large. How is this calculated, and any idea why the density is so concentrated?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Strengths\nA pretty comprehensive suite of small-to-medium scale experiments, with a toy Pinwheel dataset, MNIST, RNA, and a video dataset.\n\nTheir method seems to have the fidelity of the flow-based models while keeping the latent structure intact as with the VAEs. \n\nThe structure seems to be well preserved when visualizing the MNIST digits with TSNE, wit SFA-Mixture, SFA, Latent-FM, and GMVAE giving similar groupings for the classes.", "weaknesses": "Weaknesses:\n\nThere isn't really one consistent SFA model for testing on the different problems, with SFA, Mixture-SFA, and LDS-SFA, etc. Is there a necessity for there to be modifications to be different for every data set?\n\nAlso some terms like LDS-SFA is not properly defined. I can infer that it is SFA on the LDS dataset, but hard to grasp what makes it different.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "101CO80nfA", "reviewer": "Reviewer_fJoA", "summary": "This paper introduces Structured Flow Autoencoders (SFA), a new family of generative models that combine the high-fidelity density estimation capabilities of flow matching with the structured latent representation learning of probabilistic graphical models. The key technical contribution is the Structured Conditional Flow Matching (SCFM) objective, which explicitly incorporates latent variables into the flow matching framework, allowing joint learning of both the conditional likelihood p(x∣z) and the approximate posterior q(z∣x). The paper demonstrates the versatility of SFA across several domains, showing improvements in both sample quality and latent structure interpretability.", "questions": "1. As shown in Table 1, FM achieves performance comparable to SFA. It remains unclear what concrete advantages SFA provides over FM.\n\n2. Is there a general procedure for determining the specific form of SFA when applied to different types of graphical models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides a clear theoretical formulation that unifies two previously distinct paradigms: probabilistic graphical models and flow matching.\n\n2. The Structured Conditional Flow Matching (SCFM) loss is a well-motivated extension of the flow matching objective, effectively bridging high-fidelity neural density estimation with structured latent representation learning.\n\n3. Experiments cover toy 2D densities, images, and single-cell RNA data, demonstrating the method’s flexibility across diverse modalities and latent structural forms.", "weaknesses": "1. The paper discusses several posterior approximation families (time-indexed Gaussians, conditional CNFs, and Gumbel-Softmax for categorical latents )but provides only limited empirical comparison. It remains unclear how sensitive the final performance and stability are to these design choices.\n\n2. Several parts of the pipeline still require solving ODEs (e.g., conditional CNFs), which can be computationally intensive. Providing more detailed analysis of runtime, memory usage, and overall computational complexity would strengthen the evaluation.\n\n3. Most experiments are conducted on MNIST, Pendulum, or relatively small-scale RNA-seq datasets. It is unclear how SFA would perform on larger-scale or more complex datasets (e.g., CIFAR-10 or ImageNet). Scaling SCFM to higher-dimensional inputs or more diverse domains may introduce additional computational and optimization challenges that are not explored in this work.\n\n4. Typos: Line 451 Table 4 --> Table 2", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)