Paper: Delta-AI: Local objectives for amortized inference in sparse graphical models (ICLR-2024)

Reviewer comment (Theory):
The paper hinges on the assumption that the factors of the distribution $\phi_{k}$ are known. This seems like a stringent condition, and it is worth exploring how the framework and algorithm proposed in the paper can be extended to scenarios where these factors are unknown and need to be learned.

Author response:
That is right, we assume the *factor structure* (the set of factors and their sets of arguments) are known. Their parameters are not necessarily known and can be learned with the aid of the amortized sampler (cf. the end of Section 2.1). We provide experiments in Section 5 where the task is to learn the parameters of the graphical model associated with a data-generation process, and demonstrate the benefit of using Δ-AI's partial inference capabilities to more efficiently learn factor parameters.

Learning the undirected graphical model structure itself is quite a different problem and is outside the scope of this paper (and we are not aware of any amortized inference algorithm that jointly learns the factor structure, the parameters, and the amortized sampler). Indeed, it is typical for related work to study the problem of inference and learning in fixed graphical models, for example, using GFlowNets [Zhang et al., ICML 2022] or other algorithms taking advantage of sparsity (e.g., [Heap et al., "Massively parallel reweighted wake-sleep", UAI 2023]).

We see two approaches possible to sampling the factor structure:
- The first is to use a continuous relaxation of the space of "presence" variables for every possible factor (indicating whether a factor is present or not in the graphical model) and to fit the presence variables jointly with the factor parameters using an approximate method such as Gumbel-softmax. This would, however, not result in an unbiased sampler of the posterior over factor graph structures.
- An approach more in line with the spirit of this paper would be to use a GFlowNet to sample the factor structure, perhaps by starting with an empty graph and inserting the factors one by one. This treats the factor structure as a latent variable and amortizes the sampling of the Bayesian posterior over this structure. We see this as a promising direction for future work, as GFlowNets have been successfully used to sample the structure of **directed** graphical models (Bayesian networks) ["Bayesian structure learning with generative flow networks", UAI 2022] and to sample continuous parameters associated to their factors ["Joint Bayesian inference of graphical structure and parameters", NeurIPS 2023].

Reviewer comment (Writing):
The paper asserts that it accomplishes simultaneous amortization over multiple DAG structures. However, it would be beneficial to provide a more detailed explanation of how this simultaneous amortization is achieved.

Author response:
We attempted to explain this in Section 3, but welcome any concrete suggestions on how to make this clearer. Would it help to incorporate into the text a description along the following lines?

"Figure 2 shows an example of multiple DAG structures (I-maps) for the same undirected graphical model. The conditionals present in the DAGs are different: for example, the conditional $p(v_1\mid v_2,v_3)$ appears in the two DAGs in the second row, but not in those in the first row. The proposed algorithm learns a model $q$ that matches the conditionals in the target distribution $p$. If $q$ has a structure that allows taking varying subsets of variables as input -- such as a masked autoencoder -- then it can be trained to match the conditionals appearing in multiple DAG structures simultaneously, and the resulting sampler can then be used for sampling in any of these DAGs."

Reviewer comment (Writing):
Not having direct background / experience with GFlowNets, I found myself wondering why have two representations of the model (the factorized Markov network p and the inference network q) - is the assumption that \( p \) is provided (not just the structure / factors (sets of nodes) but also the exact parameters) - or are the parameters of both learnt?

Author response:
We assume the *factor structure* of $p$ (what the factors and their sets of arguments are) is known. We can solve different problems using our proposed algorithm, depending on whether the *parameters* of the factors in $p$ are known.
- If the parameters of the factors are known, then $q$ can be trained using Δ-AI to match the conditional distributions of $p$. This question is interesting because sampling in undirected graphical models is generally intractable (see the experiments in Section 4).
- If the parameters of the factors are unknown, then the parameters of $p$ can be updated using a maximum-likelihood gradient using samples from $q$ (see the end of Section 2.1 and experiments in Section 5), while $q$ can be trained using Δ-AI to match the conditional distributions of a $p$ that evolves over the course of training.

We remark that these two settings do not rely on understanding of GFlowNets: both sampling in PGMs and training of factorized undirected graphical models are well-established problems outside the context of GFlowNets and the background is given in Section 2.1.

Reviewer comment (Writing):
The paper is based on taking advantage of local credit assignment -> local losses this part gets kind of obsfuscated as the paper goes along. (but this is a simple fix by rewriting a bit of portions in sections 3)

Author response:
Thank you for your comment. We believe you mean that local credit assignment is not revisited after the paragraph "Local credit assignment" in Section 3. We will address this point later in Section 3 in the revised manuscript.

Reviewer comment (Experiments):
How will the proposed inference algorithm scale to natural images?

Author response:
Scaling to probabilistic graphical models (PGMs) in the context of natural image generation is a challenging task due to the sheer complexity of the high-dimensional graphical model. This is why the vast majority of existing literature constrain their study of inference methods in probabilistic graphical models to small dimensional settings. In fact, GFlowNets (which are used as a baseline in our parameter learning experiments) have, to our knowledge, never been used to learn an amortized sampler for a PGM applied to computer vision datasets larger than MNIST (see [Zhang et al., ICML 2022] and [Hu et al., ICML 2023] for two examples where discrete MNIST is the most complex task that is explored).

Importantly, contrary to competing approaches, Δ-AI is expected to scale better to modeling natural images using a sparse PGM because of its partial inference/local credit assignment capabilities. Because our contribution is Δ-AI, an inference method, and we assume the graph structure is given, we chose to model MNIST images because MNIST images are expected to have strong local dependencies and a good fit can be obtained using a simple pyramid-shaped graph structure (see Figure 6). This graph structure may not provide the best fit for natural images, and a structure-learning algorithm would probably need to be used to obtain a graph structure before applying Δ-AI to train an amortized sampler on the corresponding graphical model. We will add discussion points in the revised manuscript about promissing future directions that could take advantage of the Δ-AI mechanism to address this.