Generative Modeling of Discrete Latent Structures via Dynamic Policy Gradients

Abstract: Many scientific problems require inferring unobserved mechanistic latent states from indirect observations. While classical approaches, including expectation maximization, do not scale to combinatorially large spaces, deep learning approaches such as variational autoencoders typically form artificial latent states rather than reconstructing the mechanistic ground-truth states. Here, we introduce GReinSS, a policy learning framework that uses dynamically rescaled rewards to learn latent state distributions that maximize the observed data likelihood. We show that GReinSS accurately reconstructs simulated latent sets and latent graphs, outperforming alternative policy learning and generative modeling baselines. Additionally, GReinSS reconstructs isoforms from real short-read RNA sequencing data that better match isoforms detected by orthogonal long-read sequencing than the standard RSEM algorithm. Overall, GReinSS is a principled and practically effective approach for generative modeling and inference of combinatorial latent states from indirect observations.

• The paper presents GReinSS, a framework that adapts policy gradients to optimize the marginal likelihood for robust inference in large discrete spaces.
• It outperforms baseline methods in simulated graph, set inference tasks, and RNA isoform quantification with significantly higher F1 scores.
• Dynamic reward normalization prevents policy collapse, ensuring balanced optimization across observations and scalable probabilistic modeling.

Generative Modeling of Discrete Latent Structures via Dynamic Policy Gradients

Introduction and Problem Context

The paper presents Generative Reinforcement Learning of Structured States (GReinSS), a framework for generative modeling and probabilistic inference over discrete latent variable spaces. The motivation is the ubiquity of scientific problems that necessitate inferring mechanistic, combinatorial latent structures from indirect observations—a setting where standard unsupervised deep learning approaches (e.g., VAEs) are inadequate, as they typically learn artificial latent spaces and fail to reconstruct mechanistic ground-truth states. Classical EM-based methods, though principled, do not scale to exponentially large discrete spaces. GReinSS addresses this gap by leveraging policy gradients with dynamically rescaled rewards to optimize the marginal likelihood of observed data with respect to a generative distribution over structured latent states.

Formal Approach and Methodology

Central to GReinSS is the maximization of the observation likelihood $$\Pr(X_{1:N} \mid \theta) = \prod_{i=1}^N \sum_{S} \Pr(X_i \mid S)\Pr(S \mid \theta)$$. Unlike standard RL or variational inference, GReinSS explicitly constructs a reward for RL-style policy optimization that is dynamically adapted to optimize the marginal likelihood of indirect observations. This reward, for a sampled trajectory $\tau$ terminating in state $S(\tau)$, is given by $$r(\tau) = \sum_{i=1}^{N} \frac{\Pr(X_i \mid \tau)}{\Pr(X_i \mid \theta)}$$, guaranteeing an unbiased estimator of the likelihood gradient (Theorem 3.1). This dynamic normalization is critical—it prevents policy collapse to maximally rewarded individual trajectories, instead inducing learning of distributions that balance the likelihood across observations.

The policy over trajectories is implemented via a neural generative model parameterizing $\Pr(S|\theta)$. The method accommodates off-policy sampling; optimal off-policy proposals are analytically derived and realized via heuristic biasing toward high-likelihood regions (Theorem 3.3). The framework generalizes and unifies various latent-variable approaches: it reduces to standard generative modeling when all observations are states; to policy gradient RL if observational likelihoods are uniform; and subsumes approximate EM-like schemes for intractable latent structure inference.

Experimental Results

Simulated Graph and Set Inference

In simulated graph inference, GReinSS shows higher fidelity in reconstructing latent graphs than all baselines, including GEM-VAE, GEM-autoregression, discrete diffusion, GFlowNets, and local search. Notably, in low-information regimes (e.g., with only 10 random walks per latent state), GReinSS achieves a median F1 score of 0.891, while all other methods fall below 0.55. For simulated latent set inference, GReinSS also consistently outperforms baselines across all noise levels and set dimensions. The margin persists as the universe size grows—GReinSS remains robust (median F1 = 0.938 for universe size 1000), while non-policy learning methods deteriorate catastrophically.

RNA Isoform Quantification

A salient applied test case is transcript isoform discovery from short-read RNA-seq. Here, GReinSS outperforms the EM-based RSEM—prevalent in large-scale resources such as GTEx—when validated against long-read sequencing ground-truth. In 46.6% of genes, GReinSS reduces the isoform and proportion estimation error by at least 0.05 over RSEM; RSEM only exceeds GReinSS at this threshold on 9.4% of genes. Detailed analyses demonstrate that standard RL or naive policy gradients fail due to distributional collapse, confirming the necessity of the dynamic reward.

Theoretical and Practical Implications

GReinSS represents a fundamental advance in the application of policy gradient RL to likelihood-based inference in exponentially large combinatorial spaces. Its dynamic reward engineering enables likelihood maximization over distributions of discrete structures—contrasting the expected-return paradigm of standard RL or surrogate/ELBO optimization of variational methods. The approach is broadly flexible: it accommodates arbitrary discrete generative parameterizations (autoregressive, diffusion, neural combinatorial models), straightforwardly integrates off-policy correction, and is extensible to multi-environment or weighted observational likelihoods.

The strong empirical performance, especially in challenging indirect inference settings (as in RNA-seq), demonstrates the utility of probabilistic policy-based inference beyond domains traditionally addressed by EM. GReinSS also exemplifies a tight theoretical connection between policy gradient RL and statistical estimation for structured latent variable models.

Limitations and Future Directions

While the method admits scalability via minibatching and off-policy proposals, the performance is sensitive to the adequacy of the proposal distribution in high-dimensional or low-observational-information regimes. Analytical optimality of the off-policy sampling is proven, but practical implementations use heuristic samplers; future work could consider off-policy networks or amortized inference for proposal learning. Furthermore, the observation likelihood $\Pr(X \mid S)$ is assumed known; extending this to settings with partially parameterized or learned observation models would enhance applicability. Integration with Q-learning or actor-critic structures could further generalize the approach to complex interventional or online tasks.

Conclusion

GReinSS establishes a rigorous, scalable framework for generative modeling and inference over discrete latent structures from indirect observations, outperforming EM-based and deep generative baselines in both simulated and biological domains. By mapping marginal likelihood optimization to an RL policy gradient with dynamic reward normalization, GReinSS provides a new algorithmic foundation for probabilistic inference in large structured spaces, with broad theoretical and practical implications for discrete generative modeling, computational biology, and beyond.