Technical Summary: Generative Modeling of Discrete Data Using Geometric Latent Subspaces

Problem Statement

Generative modeling of discrete (categorical) data presents unique challenges compared to continuous data, particularly regarding the representation of statistical dependencies and the geometry of the underlying probability spaces. Existing approaches often treat discrete data as points in an ambient Euclidean space or rely on fixed-size inputs with independence assumptions between classes. Furthermore, standard dimensionality reduction techniques like Principal Component Analysis (PCA) fail to respect the intrinsic manifold structure of discrete joint distributions, often resulting in reconstructions that do not belong to the discrete data space (e.g., non-integral values). The paper addresses the need for a framework that accurately encodes discrete data, compresses high-dimensional discrete distributions, and facilitates efficient generative modeling while respecting the geometric properties of statistical manifolds.

Methodology

1. Geometric Latent Subspaces (GPCA)

The authors propose Geometric PCA (GPCA), a framework that extends classical PCA to the exponential parameter space of product manifolds of categorical distributions.

• Parameterization: Discrete data points $x \in \mathcal{X}^n$ are modeled as fully factorized product distributions $s \in \mathcal{S}_c^n$ (where $\mathcal{S}_c$ is the probability simplex). These distributions are parametrized by natural parameters $\theta$ in the exponential family, where $s = \partial \psi(\theta)$ and $\psi$ is the log-partition function.
• Latent Subspace: A low-dimensional latent subspace $\mathcal{U} \subset T_0^n$ is defined as the span of basis vectors $V$. The natural parameters are constrained to this subspace: $\theta = V \odot z$, where $z \in \mathbb{R}^d$ are latent coordinates and $d \ll nc$.
• Induced Manifold: The mapping $\partial \psi$ acts as a decoder, mapping the linear latent subspace $\mathcal{U}$ to a nonlinear data manifold $\mathcal{M} = \partial \psi(\mathcal{U}) \subset \mathcal{S}_c^n$.
• Optimization: The GPCA parameters ($V, Z$) are obtained by minimizing an empirical objective based on the Bregman divergence $D_{\psi^*}$ generated by the convex conjugate of the log-partition function. This is equivalent to minimizing the negative log-likelihood (NLL) or, when data is smoothed into the simplex, the cross-entropy.

2. Riemannian Geometry and Isometries

A core contribution is the specific geometric structure imposed on the data space to enable consistent flow matching.

• The e-Metric: Instead of the Fisher-Rao metric, the authors equip the simplex $\mathcal{S}_c^n$ with the e-metric (induced by the exponential connection). In the $\theta$-coordinate system, this metric corresponds to the standard Euclidean scalar product.
• Isometric Embedding: The mapping $\partial \psi: \mathcal{U} \to \mathcal{M}$ is shown to be an isometric embedding when $\mathcal{M}$ is equipped with the induced e-metric. Crucially, this implies that geodesics in the latent space $\mathcal{U}$ (straight lines) map exactly to e-geodesics on the manifold $\mathcal{M}$.
• Geometric Regularization: To ensure that the Riemannian distance between data and reconstructions remains small (a requirement for theoretical guarantees on flow matching), the authors introduce a regularized loss function. This loss combines the standard cross-entropy (KL divergence) with an explicit penalty on the squared Riemannian distance ($d_e^2$), controlled by a parameter $\lambda$. This prevents reconstructions from approaching the boundary of the simplex where the geometry becomes singular.

3. Flow Matching on the Latent Subspace

The generative model is trained using Flow Matching (FM) restricted to the GPCA latent subspace.

• Vector Field Learning: The goal is to learn a time-dependent vector field $u_t$ that transports a base distribution $p_0$ to the data distribution $p$.
• Efficiency via Pullback: Due to the isometric property, the conditional flow matching objective can be formulated entirely in the low-dimensional natural parameter space $\mathcal{U}$. The objective simplifies to minimizing the squared Euclidean distance between the learned vector field and the straight-line interpolation of natural parameters ($\theta_1 - \theta_0$).
• Computational Advantage: This formulation allows training to occur in the low-dimensional latent space $Z$ rather than the high-dimensional ambient space, significantly reducing computational cost and memory usage while maintaining geometric consistency.

Key Contributions

1. Geometric PCA (GPCA): A novel dimensionality reduction method for discrete data that operates in the exponential parameter space, inducing a nonlinear manifold that respects the statistical structure of categorical distributions.
2. Isometric Flow Matching: The demonstration that using the e-metric allows for an isometric relationship between the latent subspace and the data manifold. This enables flow matching to be performed via straight-line interpolation in the latent space, making the training of generative models for discrete data computationally efficient.
3. Regularized Objective: The formulation of a geometry-aware loss that balances reconstruction accuracy (Hamming distance) with geometric consistency (Riemannian distance), ensuring the validity of the flow matching approximation.
4. Empirical Validation: Extensive experiments showing that low-dimensional latent representations suffice to accurately model high-dimensional discrete data (e.g., images, DNA sequences) and that the method achieves competitive generation quality.

Experimental Results

The authors evaluated GPCA on synthetic data, MNIST, FashionMNIST, Cityscapes, and DNA sequence datasets (PROMOTER, ENHANCER).

• Reconstruction: GPCA achieves near-zero Hamming distance (mislabeled nodes) on image datasets (MNIST, Cityscapes) with relatively small latent dimensions (e.g., $d=44$ for MNIST, $d=133$ for Cityscapes).
• Trade-offs: The experiments highlight a trade-off between minimizing cross-entropy (NLL) and minimizing Riemannian distance. While NLL minimization yields low Hamming distance, it does not guarantee small Riemannian distances. The proposed regularization ($\lambda > 0$) successfully controls this trade-off, reducing Riemannian reconstruction error while maintaining low Hamming error.
• Generation Quality: Generated samples for MNIST and FashionMNIST are visually coherent and novel (not mere copies of training data). On DNA datasets, the method shows competitive performance in Fréchet Biological Distance (FBD) compared to baselines like Fisher-Flow, though the authors note that compression performance varies by dataset.
• Efficiency: Training in the latent space $Z$ (e.g., $d=64$) is significantly faster and requires less memory than training in the full data space, even with larger network architectures in the latent space.

Significance and Claims

The paper claims that its primary significance lies in bridging continuous geometric modeling with discrete data generation. By leveraging the information geometry of exponential families (specifically the e-connection), the authors provide a framework where:

• Discrete data can be compressed and represented on a manifold that preserves statistical dependencies.
• The "straight line" property of geodesics in the latent space makes flow matching effective and computationally efficient, avoiding the complexities of Riemannian flow matching in high-dimensional ambient spaces.
• The approach offers a unified, dataset-independent framework for generative modeling once the geometric latent subspace is learned.

The authors modestly acknowledge limitations, specifically that the compressibility of the data is a prerequisite (observed weaker compression on certain DNA datasets) and that the selection of the regularization parameter $\lambda$ currently relies on empirical tuning. They suggest future work involves exploring consistency models within the GPCA subspace and further analyzing the isometry between latent and data domains.