Research map / Probabilistic and scientific modeling / Variational and generative sampling
Generative Marginalization Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces marginalization models (MAMs), which directly model all induced marginals of discrete distributions and train them via marginalization self-consistency, enabling single-pass marginal inference and scalable any-order energy-based training.
Problem
Any-order autoregressive models need up to D network passes for arbitrary marginal inference and scale poorly for energy-based training on high-dimensional discrete data, where the training objective requires full log-probabilities.
Approach
Explicitly parameterize the marginal log-probability of any variable subset with one network (fed an augmented vector with a blank symbol for marginalized variables) and train it jointly with a conditional network by penalizing squared violations of the sum rule in log space, sampled randomly over orderings and data points.
Key results
- On Binary MNIST, MaM matches AO-ARM NLL (0.146 bpd) with Pearson 0.99 and about 0.018 s marginal inference vs 132.4 s for AO-ARM (roughly four orders of magnitude speedup).
- On the D=900 Ising energy-based task, MaM reaches 0.83 bpd while ARM and GFlowNet run out of GPU memory.
- On text8, MaM matches AO-ARM NLL (1.48 bpc) with Pearson 0.945 and 0.005 s inference vs 41.40 s.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MaM rather than AO-ARM for marginal inference | Binary MNIST · Binary MNIST | Time (s): 132.4 ± 0.03 -> 0.018 ± 0.00 | Helped |
| Use MaM rather than AO-ARM for marginal inference | CIFAR-10 · CIFAR-10 | Time (s): 2401 ± 1 -> 0.495 ± 0.00 | Helped |
| Use MaM rather than AO-ARM for marginal inference | ImageNet32 · ImageNet32 | Time (s): 4995 ± 1 -> 1.243 ± 0.00 | Helped |
| Use MaM rather than AO-ARM for marginal inference | Text8 · Text8 | Time (s): 41.40 ± 0.01 -> 0.005 ± 0.00 | Helped |
| Replace AO-ARM with MaM | CIFAR-10 · CIFAR-10 | NLL (bpd): 2.88 -> 2.88 | No clear effect |
| Replace AO-ARM with MaM | ImageNet32 · ImageNet32 | NLL (bpd): 3.60 -> 3.60 | No clear effect |
| Replace AO-ARM with MaM | Binary MNIST · Binary MNIST | NLL (bpd): 0.146 -> 0.146 | No clear effect |
| Replace AO-ARM with MaM | Text8 · Text8 | NLL (bpc): 1.48 -> 1.48 | No clear effect |
| Use MaM-Any-Order rather than ARM-One-Order | Ising model · Ising model (D=100) | NLL (bpd): 0.79 -> 0.80 | Hurt |
| Use MaM-Any-Order rather than ARM-MC-One-Order | Ising model · Ising model (D=100) | NLL (bpd): 24.84 -> 0.80 | Helped |
| Use MaM-Any-Order rather than ARM-One-Order for the larger problem | Ising model · Ising model (D=900) | NLL (bpd) | Helped |
| Use two-stage rather than joint training (ablation) | MaM · MNIST-Binary | epochs to convergence: 80 epochs -> 20 epochs | Helped |
Surprising findings
- Joint training with a small self-consistency weight lambda performs best in energy-based training, while large lambda hurts generative performance.
- Marginals can be used for sampling with arbitrary block sizes and orderings, with sample quality similar to conditional sampling.
Limitations
- Marginals are only approximately normalized, so normalized conditionals are needed for sampling and self-consistency is not strictly enforced.
- Out-of-distribution marginal estimates can be unreliable because the network is not robust off the training manifold.
- Joint training is sensitive to the self-consistency weight lambda.
Experimental setup
Datasets: Binary MNIST CIFAR-10 ImageNet32 text8 MOSES Ising model samples checkerboard synthetic
Benchmarks: Binary MNIST CIFAR-10 ImageNet32 text8 MOSES Ising model D=100 Ising model D=900 checkerboard
Metrics: NLL (bpd/bpc) Pearson correlation of marginals marginal inference time KL divergence Spearman correlation
Baselines: AO-ARM ARM ARM-MC P-AO-ARM PC (EiNets) GFlowNet D3PM PixelCNN++ VDM Discrete Flow
Builds on: any-order autoregressive models REINFORCE GFlowNets persistent contrastive divergence
Closest papers by meaning
- Efficient Marginalization of Discrete and Structured Latent Variables via Sparsity (NeurIPS 2020 · similarity 0.89)
- Bayesian Model Selection, the Marginal Likelihood, and Generalization (ICML 2022 · similarity 0.88)
- A Likelihood Based Approach to Distribution Regression Using Conditional Deep Generative Models (ICML 2025 · similarity 0.87)
- Algorithm- and Data-Dependent Generalization Bounds for Diffusion Models (NeurIPS 2025 · similarity 0.86)
- Stabilize the Latent Space for Image Autoregressive Modeling: A Unified Perspective (NeurIPS 2024 · similarity 0.86)
- Multimarginal Generative Modeling with Stochastic Interpolants (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.