Path-dependent Discrete Amortized Inference
Tiago da Silva, Esmeralda S. Whitammer, Salem Lahlou
cs.LG
Submitted: 2026-08-09
Updated: 2026-08-11
Comments: ICML 2026 (Oral)
Code: https://github.com/ax-ml/jax
License: http://creativecommons.org/publicdomain/zero/1.0/
The gist: We consider the problem of sampling compositional and discrete objects from a given unnormalized posterior distribution.
Terminology
Abstract
We consider the problem of sampling compositional and discrete objects from a given unnormalized posterior distribution. Notably, recent studies have shown that this problem can be efficiently solved by learning a deterministic Markov Decision Process (MDP) that progressively builds each object in proportion to the posterior. In this work, however, we demonstrate that the Markovian assumption can both hamper signal propagation during training and catastrophically reduce the learned sampler's expressivity due to state aliasing. To address these issues, we propose lifting the MDP with a learnable latent dynamical system that allows the underlying policy to depend on the entire past trajectory---and not only on the current state. In view of this, we refer to the resulting method as path-dependent discrete amortized inference. Importantly, we provably extend existing learning algorithms for discrete amortized samplers to our setting. In experiments on standard benchmark problems, we also show that our approach often leads to faster learning convergence and improved state space exploration relatively to prior techniques.
Sources
- A Conceptual Introduction to Hamiltonian Monte Carlo
- Action abstractions for amortized sampling
- A Survey of Inverse Reinforcement Learning: Challenges, Methods and Progress
- Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation
- The Logic of Graph Neural Networks
- Embarrassingly Parallel GFlowNets
- Stein Variational Inference for Discrete Distributions
- Approximating Continuous Functions by ReLU Nets of Minimal Width
- Boosted GFlowNets: Improving Exploration via Sequential Learning
- Generative Flow Networks: a Markov Chain Perspective
- Amortizing intractable inference in large language models
- A Modern Self-Referential Weight Matrix That Learns to Modify Itself
- QMDP-Net: Deep Learning for Planning under Partial Observability
- Symmetry-Aware GFlowNets
- Learning to Scale Logits for Temperature-Conditional GFlowNets
- Local Search GFlowNets
- Semi-Supervised Classification with Graph Convolutional Networks
- BatchGFN: Generative Flow Networks for Batch Active Learning
- VarGrad: A Low-Variance Gradient Estimator for Variational Inference
- Revisiting Non-Acyclic GFlowNets in Discrete Environments
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks