Paper: Generative Marginalization Models (ICLR-2024)

Reviewer comment (Theory):
My primary concern centres around the aspect of marginalization self-consistency. I'm sceptical about whether the soft constraint presented in equation 7 can effectively ensure self-consistent marginalization.

Author response:
- **Concerns about marginalization self-consistency**
    
    Please refer to **general response C1** for details**.** Thanks for your suggestions. We conducted extra experiments to validate self-consistency on a synthetic problem with known ground truth PMF. We also measured self-consistency on real-world problems (see C1.2).

Reviewer comment (Experiments):
Is it possible to validate the effectiveness of the marginal constraint in some toy examples, like training the model on a predefined Gaussian distribution with a known marginal density and directly testing it using the mean squared error?

Author response:
Dear Reviewer kkqm,

Thanks again for your comprehensive review and valuable feedback. In our previous response, we have conducted additional experiments (on both predefined distribution and real-world distribution) following your suggestions to address your concerns centered around “marginalization self-consistency”. We have also answered the questions raised in your review. In light of our response, we would appreciate if you would consider re-evaluating our submission. Please let us know if there are further questions we can address before the rebuttal period ends.

Reviewer comment (Reproducibility):
When it comes to inference, do you employ equation 6 for sampling, or do you utilize the conditional distribution \( p_\phi(x_{\sigma(d)} | x_{\sigma(<d)}) \)? What’s the difference between these two schemes?

Author response:
- **Sampling with conditionals $p_\phi$ v.s. marginals in eq.(6)**
    - Please refer to **general response C1.2 Experiment 2.** We conducted extra ablation studies of the two sampling approaches. We have also experimented sampling with different number of variables at one step using marginals.
    - In **general response C1.1**, we additionally provided a synthetic problem that purely learns marginals and sample using marginals.

Reviewer comment (Theory):
Point 2) in section 4.3 appears unclear. To train an order-agnostic AutoRegressive Model within the EBM framework, it seems plausible to employ the reinforce gradient estimation method outlined in equation 9. By replacing \(\log p_\phi (x)\) with the Monte Carlo estimation of \(\mathbb{E}_\sigma \log p_\phi(x|\sigma)\), one could potentially achieve this. Therefore, the statement "ARMs cannot be trained with the expected DKL objective over all orderings simultaneously" is somewhat unclear to me.

Author response:
- **Why not train AO-ARM with MC estimation of the gradient in energy-based training?**
    
    Please refer to **general response C3.**

Reviewer comment (Theory):
In section 4.3, point 3), I hold a different perspective regarding the efficiency of MAMs compared to ARMs, especially in high-dimensional scenarios. While it's true that in ARM-Full, you do require D feed-forward runs for gradient computation, in MAMs, you also necessitate Gibbs sampling to generate samples from the model. Even if you employ block-wise Gibbs sampling, it still demands multiple steps to guarantee MCMC chain convergence. Hence, I doubt that MAMs also face challenges in high-dimensional problems.

Author response:
- **Block-wise Gibbs sampling still demand multiple steps to guarantee MCMC chain convergence.**
    
    Please refer to **general response C4.**

Reviewer comment (Experiments):
Auto-regressive diffusion models have demonstrated strong performance on the CIFAR-10 dataset. Is it conceivable that MAM could also yield promising results on CIFAR-10?

Author response:
- **Promising results on CIFAR-10?**
    
    We followed your suggestion and conducted experiments. Indeed MaM can yield promising results on CIFAR-10. Please refer to **general response C5** for details.

Reviewer comment (Evaluation):
Could you offer additional examples to illustrate why we should be concerned with any-order AutoRegressive models? Do any-order ARMs exhibit superior performance compared to their order-specified counterparts?

Author response:
- **Why Any-Order auto-regressive models, v.s. fixed order?**
    - The test NLL achieved by AO-ARM are often just slightly worse than fixed-order ARM, since it is solving a harder problem. But they are interesting and useful due to the flexibility in sampling. For problems without natural ordering, such as molecules, proteins, materials, AO-ARMs allow full control of the generation ordering, and arbitrary conditional generation that fixed-order ARM cannot. See EvoDiff [1] for a recent example in controllable protein design. We have also provided an experiment on molecule string conditional generation (**Figure 7, 24, 25**). Even for tasks such as languages with an ordering, AO-ARMs allow us to do in-fillings and re-edits on existing text, which is useful for iteratively improving and more customizable than current language models.
    - Access to marginals further increase the flexibility in generation, such as sampling multiple variables at one step (see **general response C1.2 Experiment 2**). Marginals also allow for much better scalability in evaluating marginal likelihood, which is crucial for test time inference of $\log p$ (such as whether a new protein sequence is better in $\log p$ than its counterpart) and scalable training in energy-based setting.

Reviewer comment (Theory):
The paper presents false claims, and the whole modelling framework is not theoretically solid/supported.

Author response:
- **Marginalization Self-Consistency**
    - **Theoretical foundation:**
        
        We understand your concern regarding the theoretical solidity of MaMs. The foundation of MaMs is built on probabilistic principles, i.e. that the model should adhere to the marginalization self-consistency. In training, we make use of neural nets as powerful function approximator to learn marginals that are self-consistent.
    - **Investigating self-consistency empirically:** 

      We have also investigated in depth on how well the marginalization self-consistency is enforced in MaMs and how the soft-consistencies are very useful in practice. Please refer to **general response C1.1** for details.

Reviewer comment (Theory):
By evaluating what authors call "augumented vector representation" we are not really summing-out the unobserved variables, i.e. we are not performing marginalisation. Rather, we are computing a simple p(x) where x belongs to an augmented state space, and therefore there's no guarantee that this satisfy what authors call "marginalization self-consistency".

Author response:
- **Marginalization Self-Consistency**
    - **Theoretical foundation:**
        
        We understand your concern regarding the theoretical solidity of MaMs. The foundation of MaMs is built on probabilistic principles, i.e. that the model should adhere to the marginalization self-consistency. In training, we make use of neural nets as powerful function approximator to learn marginals that are self-consistent.
    - **Investigating self-consistency empirically:** 

      We have also investigated in depth on how well the marginalization self-consistency is enforced in MaMs and how the soft-consistencies are very useful in practice. Please refer to **general response C1.1** for details.

Reviewer comment (Theory):
As far as I understood, there's no guarantee that the model satisfies the marginalization self-consistency constraint (Eq. 5), and therefore the model compares to any other dealing with approximate marginal inference, such as [1] [2].

Author response:
- **Approximate v.s. exact**
    
    We appreciate the feedback that the method is limited by its approximate nature. We would like to point out that there is a constant trade-off of approximate and exact. By utilizing the approximation power of NNs, we have to give up some nice theoretical properties. In this case, the marginals are soft-consistent instead of being exactly consistent. We also empirically measured how well the self-consistencies are preserved and how the marginals learned can be useful in practice. MaMs unlock **three new capabilities** — 1) significantly faster marginal $\log p$ inference, 2) scalable training under energy-based setting, and 3) flexible sampling of multiple variables in one step using marginals, all due to directly modeling the marginals —  **beyond previous NN-based methods**, while maintaining strong generative modeling performance. We think our work proposed a new idea and provided evidence that learning approximate marginals are worth exploring, in parallel to how PCs with tractate likelihood properties are worth exploring as well. At the moment, they might be best suited for different tasks where marginal inference are required. But there is more hope that by further exploring those ideas, powerful generative models with fast exact likelihood can become a reality.

Reviewer comment (Theory):
Sampling is still sequential, as standard autoregressive models. This is weird, as a model allowing for proper marginal inference should not have sequential sampling.

Author response:
- **Sampling is still sequential**
    - MaMs additionally allow sampling multiple variables of different sizes in arbitrary order at each step, which is more flexible than ARM. Please refer to **general response C1.2**. We acknowledge that PCs allow one-shot conditional sampling due to the nice structures of the sum and product nodes. We have included this discussion in the updated manuscript.

Reviewer comment (Theory):
INTRO: why EBMs should be limited in fixed-order generative modeling? I do not agree with this claim.

Author response:
- **Energy-based setting**
    - **“Why AO-ARMs are limited in fixed-order modeling”:**
        
        Please refer to **general response C3**.

Reviewer comment (Novelty):
What MaMs provide that PCs don't?

Author response:
We would like to follow up on the concern of **comparing with probabilistic circuits**. 

As mentioned in our response, we have included more discussion on the comparison between neural network based generative models v.s. probabilistic circuits. Our paper included PCs as a baseline in the MNIST experiment. For Text8 and Molecule Sets, we are not able to find implementation of PCs that are well-suited for the problem. As far as we understand, PCs have not been applied to energy-based training either. If you think our comparison and discussion on PCs are inadequate, please let us know. We are happy to incorporate your suggestions and address any concerns you might have. 
Thanks again for your valuable feedback.

Reviewer comment (Evaluation):
The experimental results are not so good. In many experiments, the evaluation with NLL is slightly worse than the baseline, which does not fully demonstrate the superiority of the MAM.

Author response:
- **“Experiment results are not good”**
    
    We believe this is a misunderstanding on how methods should be compared to each other. The main merit of MaMs is in unlocking new capabilities beyond previous methods, while maintaining strong generative modeling performance. The new capabilities include: 1) significantly faster marginal $\log p$ inference, 2) scalable training under energy-based setting, and 3) flexible sampling of multiple variables in one step using marginals, all due to directly modeling the marginals. 
    
    - **Comparable generative modeling performance**: In MLE setting, the fair comparison of test NLL will be MaM with AO-ARM-Single, since that is how data is generated (using a single random ordering). Since both methods have learnable conditionals, they achieve the same NLL. The AO-ARM-Ensemble model uses an ensemble of orderings for evaluating a better quality marginal likelihood, which is used only as a reference. Because in reality, AO-ARM-Ensemble cannot be used for generation. We are sorry for the confusion.
    - **Significantly faster marginal** $\log p$ **inference** (see last column of Table 1,2,3,4) as compared with all other methods while maintaining its generative modeling performance. Probabilistic circuits (PCs) are also fast in marignal evaluation but their generative modeling capability is not as good as NN-based approaches.
    - **Training with scalability on high-dimensional problems under EB setting:** For example, Ising-model with $900$ dimensions and molecule strings with $500$ dimensions that other approaches fail to scale to.
    - **Flexible sampling with marginals**: In **general response C1.2**, we provided additional experiments illustrating the power of marginal estimations in sampling with various block variable sizes that other SOTA methods such as AO-ARM cannot perform.

Reviewer comment (Experiments):
The effectiveness of two-stage training has not been adequately explained. At least, there should be an experimental comparison with the case without the two-stage training.

Author response:
- **“Two-stage training explanation”**
    
    Thanks for the suggestion. We are happy to include more explanation. 
    
    - The basic intuition is that the optimal conditionals can be obtained by doing Stage 1 just by maximizing the likelihood. And this is shown by proof of Proposition 1 in Appendix A.1. In practice, we find it most important to use the expectation over conditionals $\mathbb E_{ x \sim p_{\text {data}}} \mathbb E_{\sigma \sim \mathcal U \left(S_D\right)} \sum_{d=1}^D \log p_\phi\left(x_{\sigma(d)} \mid x_{\sigma(<d)}\right)$ for the $\mathbb E_{ x \sim p_{\text {data}}} \log p_\theta( x )$ part of the loss in eq.(8). Otherwise if marginal network is used for this part, we empirically find maximizing likelihood pushes $\log p_\theta( x)$ to high values that break the self-consistency. We propose two-stage optimization since it is theoretically principled and this is more GPU memory efficient (than optimizing $\phi,\theta$ jointly). If there is enough compute, it is possible to train Stage 1 objective and Stage 2 objective jointly, i.e. $\min_{ \theta,\phi} - \mathbb E_{ x \sim p_{\text {data }}} \mathbb E_{\sigma \sim \mathcal U \left(S_D\right)} \sum_{d=1}^D \log p_\phi\left(x_{ \sigma (d)} \mid x_{\sigma(<d)}\right) + \lambda \mathbb E_{x \sim q(x)} \mathbb E_{\sigma \sim \mathcal U\left(S_D\right)} \mathbb E_{d \sim \mathcal U(1, \cdots, D)}\left(\log \left[p_\theta\left(x_{\sigma(<d)}\right) p_\phi\left(x_{\sigma(d)} \mid x_{\sigma(<d)}\right)\right]-\log p_\theta\left(x_{\sigma(\leq d)}\right)\right)^2$
    - In **general response C1.1**, we also tested just training MaM marginals (no conditionals are learned) on a synthetic problem. In that simple case, training with a regularized objective is totally fine and results in self-consistent marginals.

Reviewer comment (Theory):
I do not understand why two-stage training is effective. Do the authors have any hypothesis on it?

Author response:
- **“Two-stage training explanation”**
    
    Thanks for the suggestion. We are happy to include more explanation. 
    
    - The basic intuition is that the optimal conditionals can be obtained by doing Stage 1 just by maximizing the likelihood. And this is shown by proof of Proposition 1 in Appendix A.1. In practice, we find it most important to use the expectation over conditionals $\mathbb E_{ x \sim p_{\text {data}}} \mathbb E_{\sigma \sim \mathcal U \left(S_D\right)} \sum_{d=1}^D \log p_\phi\left(x_{\sigma(d)} \mid x_{\sigma(<d)}\right)$ for the $\mathbb E_{ x \sim p_{\text {data}}} \log p_\theta( x )$ part of the loss in eq.(8). Otherwise if marginal network is used for this part, we empirically find maximizing likelihood pushes $\log p_\theta( x)$ to high values that break the self-consistency. We propose two-stage optimization since it is theoretically principled and this is more GPU memory efficient (than optimizing $\phi,\theta$ jointly). If there is enough compute, it is possible to train Stage 1 objective and Stage 2 objective jointly, i.e. $\min_{ \theta,\phi} - \mathbb E_{ x \sim p_{\text {data }}} \mathbb E_{\sigma \sim \mathcal U \left(S_D\right)} \sum_{d=1}^D \log p_\phi\left(x_{ \sigma (d)} \mid x_{\sigma(<d)}\right) + \lambda \mathbb E_{x \sim q(x)} \mathbb E_{\sigma \sim \mathcal U\left(S_D\right)} \mathbb E_{d \sim \mathcal U(1, \cdots, D)}\left(\log \left[p_\theta\left(x_{\sigma(<d)}\right) p_\phi\left(x_{\sigma(d)} \mid x_{\sigma(<d)}\right)\right]-\log p_\theta\left(x_{\sigma(\leq d)}\right)\right)^2$
    - In **general response C1.1**, we also tested just training MaM marginals (no conditionals are learned) on a synthetic problem. In that simple case, training with a regularized objective is totally fine and results in self-consistent marginals.

Reviewer comment (Evaluation):
What is the definition of the marginal inference time in Tables 2, 3, and 4?

Author response:
- **“Definition of marginal inference time”**
    
    We are happy to make clarifications. Marginal inference here means given a $x_\mathcal{S}$ (for example $x_\mathcal{S}= (0,1,1,1,1,1,?,?,?)$) evaluate the marginal probability under the learned generative model, i.e. $p_\theta(x_\mathcal{S})$. With MaM, we use the marginal network $\theta$ to evaluate $p_\theta(x_\mathcal{S})$ with a single NN forward pass. As with AO-ARM, it requires evaluating a sequence of conditionals (NN forward passes) to get the marginal of $p_\theta(x_\mathcal{S})$. Those marginals are compared to the marginal from the best likelihood estimation model (AO-ARM-E) by measuring Spearman’s and pearson correlations.

Reviewer comment (Experiments):
One issue with this paper is that it mainly uses small, simple datasets for its experiments. While the results look good with these small datasets, I'm not sure how well the model would work with larger, more complex real-world data. Real-world data is often much bigger and more complicated, presenting a lot of variables and intricacies. For instance, the number of possible orderings scales fractionally with data dimension. To really understand if this approach is useful in practice, it would be helpful to test it on bigger and more diverse datasets.

Author response:
- **“Small, simple datasets for evaluation”**
    - Datasets like MNIST Binary are relatively small. We study on them first since this is a new method and we want to understand it better. Text8 and Molecular Sets (MOSES) are relatively big and widely used for testing generative models, see [1][2] for example. Text8 has 100M characters and sequence is of length $D=250$. MOSES contain ~2M molecular structures, and the sequence length is $D=55$ (for SMILES) or $D=57$ (for SELFIES). We have also included an experiment CIFAR-10 **in general response C5**, which has much higher dimensions ($D=3 \times 32 \times 32, K=256$) than MNIST Binary ($D=28 \times 28, K=2$). MaMs are able to learn marginals effectively and use them for marginal inference.

Reviewer comment (Theory):
How can MAMs generate multiple variables simultaneously? It appears that MAMs are limited to generating variables sequentially, as indicated in Equation (5) and (6).

Author response:
- “**How can MaMs generate multiple variables simultaneously?**”: For generating multiple variables simultaneously, in **general response C1.2** we provide more details and conducted experiments to show this. The idea is in similar spirit to eq.(6): evaluate marginals on all possible outcomes over combinations of multiple variables and get the normalized sampling probability, i.e. $p_\theta( x_{s_i} | x_{s(<i)}) = \frac{p_\theta([ x_{s_i}, x_{s(<i)}])}{\sum\nolimits_{ x_{s_i}} p_\theta([ x_{s_i}, x_{s(<i)}])}$, where $s_i$ is the set of multiple variables to sample at this step.

Reviewer comment (Experiments):
How does the order in which variables are sampled impact the quality of the samples?

Author response:
- **“How does the order in which variables are sampled impact the quality of the samples?”**
    - We do not observe much difference, since MaMs are trained by sampling with arbitrary orderings. Same applies to AO-ARMs. For problems with natural ordering, such as language or maybe images, it is observed that training on a fixed ordering gives better results than training on any-orderings, see [1]. But that approach loses the flexibility of generating in any-order.

Reviewer comment (Experiments):
How does different $q$ (e.g., uniform v.s. $p_{\mathrm{data}}$) affect the learning of MAMs?

Author response:
- **“How does different q (e.g., uniform v.s. p_data) affect the learning of MAMs?”**
    - Please refer to **general response C2.**

Reviewer comment (Reproducibility):
I'm a bit confused about the sample generation process with MAMs. In Section 3, it mentions using the normalized conditional (Equation (6)) for sampling, but in Section 6.1, it's said that the conditional model $p_{\phi}$ is used to generate data. Can you clarify this process?

Author response:
- “**How is sampling/generation and evaluation done with MaM?**”:
        - We can use either the conditional or the marginal for generation and evaluation. The sampling and evaluation of test NLL in the original submission is using conditionals. We have additionally included sampling and evaluation using marginals with different number of variables at one step.  Please refer to **general response C1.2** for details.
        - The test NLL is evaluated with products of conditionals $p_\phi$ or $p_\theta$ in eq.(6),. The numbers from the output $p_\theta(x)$ on a complete $x$ are not used since they are not exactly normalized probabilities. In experiments (see **general response C1.2 Experiment 1**), we evaluate the quality of $p_\theta$, and show they are highly correlated with the actual $\log p$ from conditionals.
        - In **general response C1.1**, we additionally included a synthetic task that purely learns marginals and samples using marginals.

Reviewer comment (Experiments):
Given the two-stage training approach for MLE, it seems that the learned conditional model $p_{\phi}$ is independent of the marginalization model $p_{\theta}$. As mentioned on page 5, $p_{\theta}$ is described as distilling marginals from conditionals. How does the conditional model, as an AO-ARM model, compare to the AO-ARM baselines?

Author response:
- **How does the conditional model, as an AO-ARM model, compare to the AO-ARM baselines?**
    - They perform relatively the same since both are trained with the same objective $\mathbb E_{ x \sim p_{\text {data }}} \mathbb E_{\sigma \sim \mathcal U \left(S_D\right)} \sum_{d=1}^D \log p_\phi\left(x_{\sigma(d)} \mid {x}_{\sigma(<d)}\right)$, except MaM additionally learn marginals.

Reviewer comment (Evaluation):
In the experiments, how are the NLL (bpd) values estimated for MAM? Are these values the outputs of $p_{\theta}$ or the logarithmic products of $p_{\theta}$?

Author response:
- “**How is sampling/generation and evaluation done with MaM?**”:
        - We can use either the conditional or the marginal for generation and evaluation. The sampling and evaluation of test NLL in the original submission is using conditionals. We have additionally included sampling and evaluation using marginals with different number of variables at one step.  Please refer to **general response C1.2** for details.
        - The test NLL is evaluated with products of conditionals $p_\phi$ or $p_\theta$ in eq.(6),. The numbers from the output $p_\theta(x)$ on a complete $x$ are not used since they are not exactly normalized probabilities. In experiments (see **general response C1.2 Experiment 1**), we evaluate the quality of $p_\theta$, and show they are highly correlated with the actual $\log p$ from conditionals.
        - In **general response C1.1**, we additionally included a synthetic task that purely learns marginals and samples using marginals.