Technical Summary: Generative Modeling of Discrete Latent Structures via Dynamic Policy Gradients

Problem Definition

The paper addresses the challenge of inferring unobserved, mechanistic latent states ($S^*$) from indirect observations ($X$) in scientific domains where the latent space is combinatorially large. Unlike general unsupervised learning (e.g., clustering or topic modeling) which learns artificial latent representations, this work focuses on reconstructing ground-truth mechanistic states such as chemical reaction pathways, transportation networks, or gene isoforms.

The core problem is formulated as a learning task (Problem 2.1) to find model parameters $\theta$ for a generative distribution $Pr(S|\theta)$ that maximizes the likelihood of observed data $X_{1:N}$, where $Pr(X|\theta) = \sum_{S} Pr(X|S)Pr(S|\theta)$. A secondary inference task (Problem 2.2) involves estimating the specific latent state $\hat{S}_i$ for each observation $X_i$.

Classical approaches like Expectation Maximization (EM) struggle here because the E-step requires computing expectations over exponentially large state spaces, which is often intractable. While Reinforcement Learning (RL) is naturally suited for sequential generation of structured states, standard policy gradients optimize expected rewards under stationary reward functions, not the marginal likelihood of indirect observations.

Methodology: GReinSS

The authors introduce GReinSS (Generative Reinforcement Learning of Structured States), a framework that adapts policy gradients to optimize the marginal likelihood of indirect observations.

Dynamic Reward Rescaling

The central innovation is a dynamically rescaled reward function. In standard RL, the policy gradient is $\mathbb{E}_{\tau}[r(\tau) \nabla_\theta \log Pr(\tau|\theta)]$. GReinSS defines the reward $r(\tau)$ for a trajectory $\tau$ (which yields a latent state $S(\tau)$) as:
$$r(\tau) = \sum_{i=1}^{N} \frac{Pr(X_i | \tau)}{Pr(X_i | \theta)}$$
where $Pr(X_i | \theta) = \mathbb{E}_{\tau}[Pr(X_i | \tau)]$.

Theorem 3.1 establishes that this specific reward function makes the policy gradient an unbiased estimator of the gradient of the data log-likelihood $\nabla_\theta \log Pr(X_{1:N} | \theta)$.

• Mechanism: The denominator $Pr(X_i | \theta)$ rescales the contribution of each observation. This prevents the policy from collapsing onto a single high-reward trajectory (which would maximize one observation's likelihood but fail others) and instead encourages the policy to learn a distribution $Pr(S|\theta)$ that balances the likelihood across all observations.
• Training Loop: The algorithm alternates between sampling trajectories, calculating dynamic rewards based on the current $\theta$, and updating $\theta$ via gradient ascent.

Off-Policy Learning

To address the sparsity of trajectories that yield non-zero likelihoods for specific observations, the paper employs off-policy sampling. Theorem 3.3 identifies the optimal variance-minimizing off-policy proposal distribution as $q(\tau | X_{1:N}, \theta) = \frac{1}{N} \sum_{i=1}^N Pr(\tau | X_i, \theta)$. While exact sampling from this distribution is often intractable, the authors utilize heuristics (e.g., biasing sampling based on observation features) combined with importance sampling corrections to maintain unbiased gradient estimates.

Mini-batching

The framework supports mini-batching (Corollary 3.2), where rewards are calculated over a subset of observations $B \subset [N]$, allowing the method to scale to large datasets.

Key Contributions

1. Formulation: The paper formulates the problem of learning distributions over discrete latent states from indirect observations as a policy optimization problem, bridging the gap between probabilistic modeling and reinforcement learning.
2. Theoretical Connection: It provides a theoretical proof (Theorem 3.1) that policy gradients with dynamically scaled rewards are equivalent to maximizing the marginal likelihood of indirect observations, a result not achieved by standard RL or Variational Inference (VI) formulations.
3. Generalization: Unlike previous applications of similar techniques (e.g., in cancer phylogeny or copy number variation), this work generalizes the procedure to arbitrary discrete structures (graphs, sets, sequences) and establishes connections to existing methods like Generalized Expectation Maximization (GEM) and GFlowNets under specific conditions.

Experimental Results

The authors evaluate GReinSS on simulated and real-world datasets, comparing it against local search, standard policy gradients, GFlowNets, and GEM-based implementations using Variational Autoencoders (VAE), autoregressive models, and discrete diffusion.

1. Simulated Process Graph Inference

• Task: Reconstruct directed graphs from random walk endpoints.
• Results: GReinSS consistently achieved the highest median F1 scores. Notably, with low information per observation ($k=10$ random walks), GReinSS achieved an F1 of 0.891, while all baselines (including naive policy gradients and GFlowNets) scored below 0.55. The naive policy gradient baseline often collapsed to predicting empty graphs, highlighting the necessity of the dynamic reward rescaling.

2. Simulated Subset Inference

• Task: Reconstruct subsets from noisy element measurements.
• Results: GReinSS outperformed all baselines across varying noise levels ($\sigma$) and universe sizes ($|U|$).
  • For small universes ($|U|=10$), most methods performed well.
  • For large universes ($|U|=1000$), baselines (VAE, diffusion, autoregression) collapsed (median F1 < 0.4), whereas GReinSS maintained a median F1 of 0.938.
  • The results demonstrate that GReinSS scales effectively to large combinatorial spaces where other generative models fail.

3. Real-World Application: RNA Isoform Reconstruction

• Task: Infer RNA isoform structures and proportions from short-read sequencing data (GTEx dataset), validated against long-read sequencing ground truth.
• Comparison: Compared against RSEM, the standard Expectation Maximization-based algorithm used in the GTEx project.
• Results: GReinSS significantly outperformed RSEM.
  • For the gene MBD2, RSEM predicted a major isoform with 18% proportion (vs. 96% ground truth) and incorrect minor isoforms, resulting in a prediction error of 0.5374. GReinSS predicted 92% for the major isoform with an error of 0.0067.
  • Across 14,390 genes, GReinSS outperformed RSEM by a margin of $\ge 0.05$ in 46.6% of cases, while RSEM outperformed GReinSS in only 9.4% of cases. The median difference in error was $-0.0405$ in favor of GReinSS.

Significance and Claims

The paper claims that GReinSS is a "principled and practically effective approach" for generative modeling in combinatorial spaces. Its primary significance lies in:

• Overcoming Intractability: It provides a solution for latent state inference in spaces where exact EM is infeasible due to exponential complexity.
• Superiority over Standard RL: It demonstrates that standard policy gradients fail in this context without dynamic reward rescaling, and that GFlowNets (which balance trajectory probabilities) are insufficient when the reward structure does not perfectly align with the observation likelihoods.
• Real-World Impact: The successful application to RNA splicing, outperforming a widely used standard (RSEM) on real biological data, validates the method's utility beyond synthetic benchmarks.

The authors note that while the method currently relies on heuristics for off-policy sampling and assumes a known $Pr(X|S)$, the framework offers a foundation for future extensions, including Q-learning, actor-critic frameworks, and learning the observation model itself.