A Neural Network Approach for Efficiently Answering Most Probable Explanation Queries in Probabilistic Models

Authors: shivvrat  arya
tahrima  rahman
vibhav  gogate
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 00:54:13

Abstract: We propose a novel neural networks based approach to efficiently answer arbitrary Most Probable Explanation (MPE) queries—a well-known NP-hard task—in large probabilistic models such as 
Bayesian and Markov networks, probabilistic circuits, and neural auto-regressive models. By arbitrary MPE queries, we mean that there is no predefined partition of variables into evidence and non-evidence variables. The key idea is to distill all MPE queries over a given probabilistic model into a neural network and then use the latter for answering queries, eliminating the need for time-consuming inference algorithms that operate directly on the probabilistic model. We improve upon this idea by incorporating inference-time optimization with self-supervised loss to iteratively improve the solutions and employ a teacher-student framework that provides a better initial network, which in turn, helps reduce the number of inference-time optimization steps. The teacher network utilizes a self-supervised loss function optimized for getting the exact MPE solution, while the student network learns from the teacher's near-optimal outputs through supervised loss. We demonstrate the efficacy and scalability of our approach on various datasets and a broad class of probabilistic models, showcasing its practical effectiveness.

Main paper: 
 introduction: 
 Probabilistic representations such as Probabilistic Circuits (PCs) [8] , graphical models [26] such as Bayesian Networks (BNs) and Markov Networks (MNs), and Neural Autoregressive Models (NAMs) [50] are widely used to model large, multi-dimensional probability distributions. However, they face a significant challenge: as the complexity of these distributions increases, solving practically relevant NP-hard inference tasks such as finding the Most Probable Explanation (MPE) via exact inference techniques [39, 40] becomes increasingly difficult and time-consuming. In particular, although various exact and approximate solvers exist for the MPE task in PCs, BNs and MNs, exact solvers are often too slow for practical use, and approximate solvers tend to lack the necessary accuracy, particularly in autoregressive models that currently rely on slow hill-climbing/beam search methods.In recent work, Arya et al. [4] proposed a method to overcome the limitations of existing approximate methods by using neural networks (NNs) to solve the MPE task in PCs. 1 Their method draws inspiration from the learning to optimize literature [12, 15, 29, 42, 55] . Given a PC and a predefined partition of variables into query and evidence sets, the core idea is to train a NN that takes an assignment to the evidence variables as input and outputs the most likely assignment to the query variables w.r.t. the distribution defined by the PC. Arya et al. suggest using either supervised or self-supervised learning techniques to train the NN; the former requires access to exact inference schemes, while the latter does not and is therefore more practical.In this paper, we address a more general and complex version of the MPE task than the one considered by Arya et al. Specifically, we assume that there is no predefined partition of the variables into evidence and query sets, which we refer to as the any-MPE task. The complexity of the any-MPE task arises from the exponential increase in the number of input configurations, compounded by the exponential number of possible divisions of variables into evidence and query sets. Furthermore, our method applies to a broad class of probabilistic models, including BNs, MNs and NAMs, whereas Arya et al.'s method is limited to PCs. In addition, Arya et al.'s method does not fully exploit the capabilities of self-supervision, and the benefits of combining supervised and self-supervised loss functions.This paper presents a novel approach that uses a NN for solving the any-MPE task in a broad class of probabilistic models (PMs) and achieves technical advancements in three key aspects:1. Efficient MPE Inference via Encoding Scheme and Loss Function: We introduce a new encoding scheme that tailors the NN architecture to the specific structure of the input PM. This scheme not only delineates the input and output nodes for the NN but also establishes a methodology for setting input values and extracting the MPE solution from the NN's outputs. Furthermore, we propose a tractable, and differentiable self-supervised loss function, enabling efficient training. 
inference time optimization with itself:: 
 We introduce a novel inference technique called Inference Time Self Supervised Training (ITSELF). This technique iteratively refines the MPE solution during the inference process itself. It utilizes gradient descent (back-propagation) to update the NN's parameters using our proposed self-supervised loss, leading to continual (anytime) improvement towards near-optimal solutions. ITSELF fully utilizes the power of our self-supervised loss, as it does not require labeled data or an external MPE solver.Without loss of generality, we use binary variables which take values from the set {0, 1}. We denote a random variable by an uppercase letter (e.g., X), and a value assigned to it by the corresponding lowercase letter (e.g., x). We denote a set of random variables by a bold uppercase letter (e.g., X) and an assignment of values to all variables in the set by the corresponding bold lowercase letter (e.g., x).Throughout the paper when we use the term probabilistic models (PMs), we are referring to a broad class of probabilistic models in which computing the likelihood2 of an assignment to all variables in the model can be done in polynomial (preferably linear) time in the size of the model. This class includes, among others, Bayesian and Markov networks collectively called Probabilistic Graphical Models (PGMs) [26] , smooth and decomposable Probabilistic Circuits (PCs) [8] , and Neural Autoregressive Models (NAMs) such as NADE [50] and MADE [17] .We are interested in solving the most probable explanation (MPE) task in PMs, namely the task of finding the most likely assignment to all unobserved (non-evidence) variables given observations (evidence). Formally, let M denote a probabilistic model defined over a set of variables X that represents the distribution p M (x). We categorize the variables X into evidence E ⊆ X and query Q ⊆ X groups, ensuring that E ∩ Q = ∅ and E ∪ Q = X. Then, given an assignment e to the set of evidence variables E, the MPE task can be formulated as:EQUATIONIt is known that the MPE task is NP-hard in general and even hard to approximate [9, 11, 36, 41, 44] .Motivation: The goal of this paper is to develop a method that trains a NN for a given PM and, at test time, serves as an approximate MPE solver for any-MPE query posed over the PM. By any-MPE, we mean that the NN can take an assignment to an arbitrary subset of variables (evidence) as input and output the most likely assignments to the remaining (query) variables. Recently, Arya et al. [4] proposed a NN-based solution for solving the MPE task in PCs under the constraint that the partition of the variables into evidence and query sets is known before training the NN. This constraint is highly restrictive because, for generative models, it is unlikely that such a partition of variables is known in advance. In such cases, one would typically train a discriminative model rather than a generative one. Unlike Arya et al.'s method, our approach yields an any-MPE solver. Additionally, Arya et al.'s approach has several limitations in that it does not fully exploit the benefits of self-supervision during inference time and requires the use of relatively large NNs to achieve good performance in practice. Our proposed approach, described next, addresses these limitations. 
two-phase pre-training with teacher-student architecture:: 
 To address challenges associated with self-supervised learning and ITSELF, we propose a two-phase pre-training strategy that leverages a teacher-student architecture. Self-supervised learning can suffer from overfitting and requires careful regularization. Additionally, ITSELF, especially with random initializations, might necessitate a substantial number of gradient updates to converge on optimal solutions. Our approach addresses these issues using the following methodology: (i) The teacher network first overfits the training data using ITSELF and (ii) The student network is then trained using supervised loss functions (e.g., binary cross-entropy) by treating the teacher network's output as pseudo-labels. This supervised training phase improves and regularizes the parameter learning process of the student network. It also provides a robust starting point for ITSELF, significantly reducing the required optimization steps and leading to substantial performance gains.Finally, we conduct a detailed experimental comparison of our method with existing approaches on several types of PMs such as PCs, PGMs and NAMs. Our results demonstrate that our method surpasses state-of-the-art approximate inference techniques in terms of both accuracy and speed.In this section, we develop a neural network (NN) based approach for solving the any-MPE task. Specifically, given a PM, we develop an input encoding (see Section 3.1) that determines the number of input nodes of the NN and sets their values for the given MPE query. Additionally, we develop an output encoding scheme that specifies the number of NN output nodes required for the given PM and enables the recovery of the MPE solution from the outputs. For training the NN, we introduce a tractable and differentiable self-supervised loss function (see Section 3.2), whose global minima aligns with the MPE solutions to efficiently learn the parameters of the NN given unlabeled data.Since NNs require fixed-sized inputs and outputs, we introduce input and output encodings that generate fixed-length input and output vectors for each PM from a given MPE problem instance MPE(Q, e). To encode the input, for each variable X i ∈ X, we associate two input nodes in the NN, denoted by Xi and Xi . Thus for a PM having n (namely, |X| = n) variables, the corresponding NN has 2n input nodes. Given a query MPE(Q, e), we set the values of the input nodes as follows: (1) If X i ∈ E and X i = 0 is in e, then we set Xi = 0 and Xi = 1; (2) If X i ∈ E and X i = 1 is in e, then we set Xi = 1 and Xi = 0; and (3) If X i ∈ Q then we set Xi = 0 and Xi = 0. (The assignment Xi = 1 and Xi = 1 is not used.) It is easy to see that the input encoding described above yields an injective mapping between the set of all possible MPE queries over the given PM and the set {0, 1} 2n . This means that each unique MPE query (Q, e) will yield a unique 0-1 input vector of size 2n.The output of the neural network comprises of n nodes with sigmoid activation, where each output node is associated with a variable X i ∈ X. We ignore the outputs corresponding to the evidence variables and define a loss function over the outputs corresponding to the query variables in the set Q. The MPE solution can be reconstructed from the output nodes of the NN by thresholding the output nodes corresponding to the query variables appropriately (e.g., if the value of the output node is greater than 0.5, then the query variable is assigned the value 1; otherwise it is assigned to 0).Since the output nodes of our proposed NN use sigmoid activation, each output is continuous and lies in the range [0, 1]. Given an MPE query MPE(Q, e), let q c ∈ [0, 1] |Q| denote the (continuous) Most Probable Explanation (MPE) assignment predicted by the NN. In MPE inference, given e, we want to find an assignment q such that log p M (q, e) is maximized, namely,log p M (q, e) is minimized. Thus, a natural loss function that we can use islog p M (q, e). Unfortunately, the NN outputs a continuous vector q c and as a result p M (q c , e) is not defined.Next, we describe how to solve the above problem by leveraging the following property of the class of PMs that we consider in this paper-specifically BNs, MNs, PCs and NAMs. In these PMs, the function ℓ(q, e) =log p M (e, q), which is a function from {0, 1} n → R is either a multi-linear polynomial or a neural network, and can be computed in linear time in the size of the PM. To facilitate the use of continuous outputs, we define a loss function ℓ c (q c , e) : [0, 1] n → R such that ℓ c coincides with ℓ on {0, 1} n . For PGMs and PCs, ℓ is a multi-linear function and ℓ c is obtained by substituting each occurrence of a discrete variable q i ∈ q with the corresponding continuous variable q c i ∈ q c where q c i ∈ [0, 1]. In NAMs, ℓ is a NN and we can perform a similar substitution-we substitute each binary input q i in the NN with a continuous variable q c i ∈ [0, 1]. This substitution transforms the discrete NN into a continuous function while preserving its functional form.An important property of ℓ c is that it can be evaluated and differentiated in polynomial time. Moreover, when ℓ is defined by either a neural network (in NAMs) or a multilinear function (in BNs, MNs and PCs), the minimum value of ℓ c over the domain [0, 1] n is less than or equal to the minimum value of the original function ℓ over the discrete domain {0, 1} n . Formally, Proposition 1. Let l(q, e) : {0, 1} n → R be either a neural network or a multilinear function, and let l c (q c , e) : [0, 1] n → R be its continuous extension obtained by substituting each binary input q i with a continuous variableq c i ∈ [0, 1]. Then, min q c ∈[0,1] n ℓ c (q c , e) ≤ min q∈{0,1} n ℓ(q, e)Following Arya et al. [4] , we propose to improve the quality of the loss function by tightening the lower bound given in proposition 1 with an entropy-based penalty (ℓ E ), governed by α > 0.EQUATIONThis penalty encourages discrete solutions by preferring q c j values close to 0 or 1, where α modulates the trade-off. Setting α to 0 yields the continuous approximation; conversely, an α value of ∞ results exclusively in discrete outcomes. From proposition 1 and by using the theory of Lagrange multipliers, we can show that for any α > 0, the use of the entropy penalty yields a tighter lower bound: Proposition 2.min q c ∈[0,1] n ℓ c (q c , e) ≤ min q c ∈[0,1] n ℓ c (q c , e) + ℓ E (q c , α) ≤ min q∈{0,1} n ℓ(q, e)How to use the Loss Function: Given a PM defined over n variables, we can use the self-supervised loss function ℓ c (q c , e) + ℓ E (q c , α) (treating α as a hyper-parameter) to train any neural network (NN) architecture that has 2n input nodes and n output nodes. This trained NN can then be used to answer any arbitrary MPE query posed over the PM. The training data for the neural network consists of assignments (evidence e) to a subset of the variables. Each training example can be generated using the following three-step process. We first sample a full assignment x to all variables in the PM using techniques like Gibbs sampling or perfect sampling for tractable distributions such as PCs and BNs. Second, we choose an integer k uniformly at random from the range {1, . . . , n} and designate k randomly selected variables as evidence variables E, and the remaining nk as query variables Q. Finally, we project the full assignment x on E. The primary advantage of using the self-supervised loss function is that it eliminates the need for access to a dedicated MPE solver to provide supervision during training; gradient-based training of the neural network provides the necessary supervision. At a high level, assuming that the NN is over-parameterized, if we use the self-supervised loss and repeatedly run (stochastic) gradient updates over the NN for a given dataset, theoretical results [2, 13] as well as prior experimental work [46, 56] suggest that the parameters of the NN will converge to a point near the global minimum of the self-supervised loss function. This means that through gradient updates, the network will find a near-optimal MPE assignment for each training example. This strategy of performing gradient updates over the NN can also be used during inference (test) time to iteratively improve the MPE solution, thereby maximizing the benefits of self-supervision.Specifically, at test time, given a test dataset (or example), we initialize the NN either randomly or using a pre-trained model and then run gradient-based updates over the NN iteratively until convergence. The gradient is computed w.r.t. the self-supervised loss function ℓ c (q c , e) + ℓ E (q c , α).We call the resulting algorithm ITSELF (Inference Time Optimization using SELF-Supervised Loss), as detailed in Figure 1 . The performance of ITSELF typically improves with each iteration until the loss converges.Our proposed method, ITSELF, is closely related to test-time training approaches which are widely used to solve problems in deep learning [1, 10, 19, 30-32, 38, 49, 51, 57] . Our method differs from these previous approaches in that the global minima of our proposed self-supervised loss correspond to the MPE solutions, provided that the penalty α is sufficiently large. 
supervised knowledge transfer from itself: 
 A drawback of our self-supervised loss function is that, unlike supervised loss functions such as binary cross entropy, it is a non-convex function of the NN outputs 3 . As a result, it has a significantly larger number of local minima compared to the supervised loss function, but also a potentially exponential number of global minima, because an MPE problem can have multiple optimal solutions [35] , all of which have the same loss function value. Thus, optimizing and regularizing using the self-supervised loss is difficult compared to a supervised loss, especially when the number of training examples is large.Moreover, our experiments show that large datasets necessitate large, over-parameterized neural networks (NNs) to achieve near-optimal MPE solutions for all examples. However, when the training data is limited and the NN is sufficiently over-parameterized, our preliminary findings, along with theoretical and empirical results from prior studies [3, 6, 23, 27, 28] , suggest that the NN is more likely to approach the global optima. Specifically, with a reasonably sized NN and a small dataset, the algorithm ITSELF tends to yield near-optimal MPE solutions. A further challenge with ITSELF is that even for small datasets, achieving convergence from a random initialization requires numerous iterations of gradient descent, rendering the training process inefficient and slow. Make a forward-pass over T to get an MPE assignment q i for e i 10:Update the entry in DB for e i with q i if it has a lower loss value than the current entry 11: end for 12: Update the parameters of S using the mini-batch D ′ and labels from DB and a supervised loss 13: T ← S ▷ Initialize T with S for the next epoch 14: end forTo address these challenges (using small datasets with ITSELF; designing better initialization for it; and using non-convex loss functions for training), we propose a two-network teacher-student strategy [7, 16, 20-22, 24, 37, 47, 52-54] , where we have two networks with the same structure that are trained via mini-batch gradient updates. The teacher network is overfitted to the mini-batch using our self-supervised loss via the ITSELF algorithm, and the student network is subsequently trained with a supervised loss function such as binary cross entropy. By overfitting the teacher network via ITSELF on the mini-batch, we ensure that it finds near-optimal MPE assignments for all (unlabeled) examples in the mini-batch and eventually over the whole training dataset.The student network then learns from the teacher's outputs, using them as soft labels in a supervised learning framework. This transfer of knowledge mitigates the optimization difficulties associated with the non-convex self-supervised loss, allowing the student network to achieve faster convergence and better generalization with a more manageable model size. Additionally, this strategy reduces the need for severe over-parameterization and extensive training iterations for the teacher network because it is operating on a smaller dataset. It also helps achieve better initialization for ITSELF.Our proposed training procedure, which we call GU IDE, is detailed in Algorithm 1. The algorithm trains a two-network system comprising a teacher network (T ) and a student network (S) with the same structure. The goal is to train the student network using a combination of self-supervised and supervised learning strategies. The algorithm takes as input the training data D, along with the teacher and student networks, T and S, respectively and outputs a trained network S. A database (DB) is utilized to store the best MPE assignment and corresponding loss value for each example in D. The parameters of T and S, and the entries in DB, are randomly initialized at the start.In each epoch, a mini-batch D ′ is sampled from the training data D. The parameters of the teacher network T are then updated using the ITSELF algorithm (which uses a self-supervised loss), applied to the mini-batch D ′ (the mini-batch helps address large data issues associated with ITSELF). For each example e i in D ′ , we perform a forward-pass over T to obtain an MPE assignment q i . The database DB is subsequently updated with q i if it has a lower loss value than the current entry for e i .Following this, the parameters of the student network S are updated using the mini-batch D ′ , the labels from DB, and a supervised loss function (ℓ sup ) such as Binary Cross Entropy or L2 loss. Finally, the parameters of the teacher network T are reinitialized with the updated parameters of the student network S to prepare for the next epoch (addressing the initialization issue associated with ITSELF). Figure 2 illustrates a single training epoch of GUIDE. Thus, at a high level, Algorithm 1 leverages the strengths of both self-supervised and supervised learning to improve training efficiency and reduce the model complexity, yielding a student network S. Moreover, at test time, the student network can serve as an initialization for ITSELF. 
experiments: 
 This section evaluates the ITSELF method (see section 3.3), the GU IDE teacher-student training method (see section 4) and the method that uses only self-supervised training, which we call SSMP (see section 3.2). We benchmark these against various baselines, including neural network-based and traditional polynomial-time algorithms that directly operate on the probabilistic model. We begin by detailing our experimental framework, including competing methods, evaluation metrics, neural network architectures, and datasets.We used twenty binary datasets extensively used in tractable probabilistic models literature [5, 18, 34, 50] -referred to as TPM datasets-for evaluating PCs and NAMs. For the purpose of evaluating PGMs, we utilized high treewidth models from previous UAI inference competitions [14] .To train Sum Product Networks (SPNs), our choice of PCs, we employed the DeeProb-kit library [33] , with SPN sizes ranging from 46 to 9666 nodes. For NAMs, we trained Masked Autoencoder for Distribution Estimation (MADE) models using PyTorch, following the approach in Germain et al. [17] . For Markov Networks (MNs), a specific type of PGM, we applied Gibbs sampling to generate 8,000, 1,000, and 1,000 samples for the training, testing, and validation sets, respectively. The query ratio (qr), defined as the fraction of variables in the query set, was varied across the set {0.1, 0.3, 0.5, 0.7, 0.8, 0.9} for each probabilistic model (PM).PCs -We used three polynomial-time baseline methods from the probabilistic circuits and probabilistic graphical models literature as benchmarks [41, 45] .• MAX Approximation (MAX) [45] transforms sum nodes into max nodes. During the upward pass, max nodes output the highest weighted value from their children. The downward pass, starting from the root, selects the child with the highest value at each max node and includes all children of product nodes. • Maximum Likelihood Approximation (ML) [41] computes the marginal distribution p M (Q i |e) for each variable Q i ∈ Q, setting Q i to its most likely value. • Sequential Approximation (Seq) [41] iteratively assigns query variables according to an order o.At each step j, it selects the j-th query variable Q j in o and assigns to it a value q j such that p M (q j |e, y) is maximized, where y is an assignment of values to all query variables from 1 to j -1.We further evaluated the impact of initializing stochastic hill climbing searches using solutions from all baseline approaches and our proposed methods for MPE inference, conducting 60-second searches for each MPE problem in our experiments, as detailed in Park and Darwiche [41] .NAMs -As a baseline, we used the stochastic hill-climbing search (HC) algorithm. Following a procedure similar to that used for PCs, we conducted a 60-second hill-climbing search for each test example, with query variables initialized randomly and setting evidence variables according to the values in the given example.PGMs -We employed the distributed AND/OR Branch and Bound (AOBB) method [39] as a baseline, using the implementation outlined in Otten [40] . Since AOBB is an anytime algorithm, we set a 60-second time limit for inference per test example.Neural Baselines -Arya et al. [4] introduced Self-Supervised learning based MMAP solver for PCs (SSMP), training a neural network to handle queries on a fixed variable partition within PCs. We extend this approach to address the any-MPE task in PMs (see Section 3.2), using a single network to answer any-MPE queries as an additional neural baseline.Evaluation Criteria -We evaluated competing approaches based on log-likelihood (LL) scores, calculated as ln p M (e, q), and inference times for given evidence e and query output q. Higher log-likelihood scores indicate better performance, while shorter inference times are preferable. In the first strategy, we performed a single forward pass through the network to estimate the values of query variable, as specified by Arya et al. [4] . The second strategy utilized our novel testtime optimization-based ITSELF approach for inference. The ITSELF optimization terminates after 100 iterations or upon loss convergence for both PCs and PGMs. For NAMs, we increase the limit to 1,000 iterations while keeping the convergence criterion.We standardized network architectures for PMs across all experiments. For PCs, we used fully connected Neural Networks (NN) with three hidden layers (128, 256, 512 nodes). For NAMs and PGMs, a single hidden layer of 512 nodes was employed. All hidden layers featured ReLU activation, while the output layers used sigmoid functions with dropout for regularization [48] . We optimized all models using Adam [25] and implemented them in PyTorch [43] on an NVIDIA A40 GPU.Results for PCs: We compare methods-including three polynomial-time baselines, neural network-based SSMP, and our ITSELF and GU IDE methods-on 20 TPM datasets as shown in the contingency table in figure 3a (detailed results in the supplementary materials). We generated 120 test datasets for the MPE task using 20 PCs across 6 query ratios (qr). Each cell (i, j) in the table represents how often (out of 120) the method in row i outperformed the method in column j based on average log-likelihood scores. Any difference between 120 and the combined frequencies of cells (i, j) and (j, i) indicates cases where the compared methods achieved similar scores. We present similar contingency tables for Hill Climbing Search over PCs (Fig. 3b ), NAMs (Fig. 3c ), and PGMs (Fig. 3d ) to benchmark the proposed methods against the baselines.The contingency table for PC (Fig. 3a ) shows that methods incorporating ITSELF consistently outperform both polynomial-time and traditional neural baselines, as indicated by the dark blue cells in the corresponding rows. Notably, GU IDE + ITSELF is superior to all the other methods in almost two-thirds of the 120 cases, while SSMP + ITSELF is better than both SSMP and GU IDE. In contrast, the polynomial-time baseline MAX is better than both SSMP and GU IDE (as used in Arya et al. [4] ), highlighting ITSELF's significant role in boosting model performance for the complex any-MPE task.We compare MAX and GU IDE + ITSELF using a heatmap in Figure 4a . The y-axis presents datasets by variable count and the x-axis represents query ratio. Each cell displays the percentage difference in mean LL scores between the methods, calculated as %Diff. = 100×(ll nn -ll max )/|ll max |. The heatmap shows that GU IDE + ITSELF achieves performance comparable to MAX for small query sets. As the problem complexity increases with an increase in query set size, our method consistently outperforms MAX across all datasets, except for NLTCS and Tretail, as highlighted by the green cells. In the 12 cases where GU IDE + ITSELF underperforms, the performance gap remains minimal, as indicated by the limited number of red cells in the heatmap. Results for NAMs: The contingency table in Figure 3c presents our evaluation of several methods for NAMs, including HC and two neural network approaches, SSMP and GU IDE, each tested with two inference schemes. We evaluated these methods on 20 TPM datasets, creating 80 test sets for the MPE task using 20 MADEs across four query ratios (qr).The GU IDE + ITSELF approach demonstrates superior performance compared to both baseline methods and other neural inference schemes, aligning with observations from PC. While HC outperforms SSMP, both GU IDE and the combination of SSMP-based training with ITSELF-based inference surpass HC, highlighting their advantages over the baseline.The heatmaps in Figure 4b further highlight the superior performance of GU IDE + ITSELF for NAMs, particularly in larger datasets where it outperforms the HC baseline by over 50% in most cases, as indicated by the dark green cells. The combination of GU IDE-based learning with ITSELF-based inference consistently outperforms the baseline across most datasets, with exceptions only in the Mushrooms, Connect 4, and Retail. Overall, the GU IDE + ITSELF approach significantly enhances the quality of the MPE solutions in NAM models.The contingency table in 3d compares the performance of AOBB and four neural-network-based methods on PGMs across four high-treewidth networks. For this evaluation, we generated 16 test datasets for the MPE task using four PGMs across four query ratios (qr).Consistent with results from previous PMs, methods using IT-SELF for inference consistently outperform the baseline methods AOBB and SSMP across most scenarios. Both GU IDE and SSMP outperform AOBB in at least 50 percent of the tests. The supplementary material presents comparisons against exact solutions, conducted on less complex probabilistic models where ground truth computation remains tractable.This analysis aims to evaluate the performance of GU IDE against traditional neural network training methods used in SSMP across different PMs and inference schemes. Using traditional inference scheme (i.e., one forward pass through the network), GU IDE consistently outperforms SSMP, demonstrating its superiority in 60% of scenarios for PCs, more than 80% for NAM models, and 75% for PGM models. When employing ITSELF-based inference, GU IDE maintains this advantage, achieving higher quality solutions in more than 75%, 85%, and 80% of cases for PCs, NAMs, and PGMs, respectively. Therefore, models trained using GU IDE are consistently superior to those trained with SSMP for the any-MPE task.In this analysis, we compare the performance of the single-pass inference method to that of the proposed multi-pass inference method (ITSELF). ITSELF combined with SSMP training outperforms the other methods in over 85% cases for PC, and more than 75% for NAM and PGM models. When used on models trained with GU IDE, ITSELF demonstrates even better results, achieving superior performance in nearly 90% of PC cases and 75% for both NAMs and PGMs. Overall, GU IDE with ITSELF inference emerges as the most effective method across all experiments. Empirical evidence consistently demonstrates ITSELF's superiority over single-pass inference across PMs.The inference time analysis, detailed in the supplementary material, compares computational efficiency across methods using the natural logarithm of execution time in microseconds. Neural network-based approaches with traditional inference demonstrate the fastest performance across all PMs, as they only require a single forward pass to compute query variable values. For MADE, models trained with GU IDE and ITSELF are the next most efficient. In PGMs, GU IDE + ITSELF ranks third, followed by SSMP + ITSELF. For PCs, MAX is marginally faster than both GU IDE + ITSELF and SSMP + ITSELF, while ML and Seq have the longest computational times. In general, models trained with GU IDE achieve shorter inference times than those trained with the self-supervised loss (SSMP), as they require fewer ITSELF iterations due to more effective initial training.Our experiments demonstrate that GU IDE + ITSELF outperforms both polynomial-time and neural-based baselines across various PMs, as evidenced by higher log-likelihood scores. Notably, ITSELF demonstrates significant advantages over traditional single-pass inference in addressing the complex any-MPE query task within probabilistic models, emphasizing the importance of Inference Time Optimization. Furthermore, the superior performance of models trained with GU IDE compared to SSMP highlights the effectiveness of the dual network approach, which improves initial model quality and establishes an optimal starting point for ITSELF. 
conclusion and future work: 
 We introduced novel methods for answering Most Probable Explanation (MPE) queries in probabilistic models. Our approach employs self-supervised loss functions to represent MPE objectives, enabling tractable loss and gradient computations during neural network training. We also proposed a new inference time optimization technique, ITSELF, which iteratively improves the solution to the MPE problem via gradient updates. Additionally, we introduced a dual-network-based strategy that combines supervised and unsupervised training which we call GU IDE to provide better initialization for ITSELF and addressing various challenges associated with self-supervised training. Our method was tested on various benchmarks, including probabilistic circuits, neural autoregressive models, and probabilistic graphical models, using 20 binary datasets and high tree-width networks. It outperformed polytime baselines and other neural methods, substantially in some cases. Additionally, it improved the effectiveness of stochastic hill climbing (local) search strategies.Future work includes solving complex queries in probabilistic models with constraints; training neural networks with losses from multiple probabilistic models to embed their inference mechanisms; boosting performance by developing advanced encoding strategies for similar tasks; implementing sophisticated neural architectures tailored to probabilistic models; etc. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our paper improves the speed, accuracy and scalability of MPE inference algorithms. Since algorithms already exist for solving MPE tasks, we do not perceive any negative or positive societal impacts beyond what currently exists.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Table 1 summarizes the datasets and the probabilistic circuits trained on them. We use the same datasets for both PCs [8] and NAMs [17, 50] . The selection includes both smaller datasets, such as NLTCS and MSNBC, and larger datasets with over 1000 nodes.For Markov networks, we utilize high treewidth grid networks, specifically grid40x40.f2.wrap, grid40x40.f5.wrap, grid40x40.f10.wrap, and grid40x40.f15.wrap. Each model contains 4800 variables and 1600 factors. Our experimental framework was designed to ensure consistency and efficiency across all conducted experiments. For NAM's, we used MADE, training the model with two hidden layers of 512 and 1024 units, respectively, using the hyperparameters from Germain et al. [17] .For neural network-based solvers, the mini-batch size was set to 512 samples, and a learning rate decay strategy, reducing the rate by 0.9 upon loss plateauing, was implemented to improve training efficiency. Optimal hyperparameters were identified via extensive 5-fold cross-validation.In discrete loss scenarios, the hyperparameter α played a pivotal role. We systematically explored the optimal α value across the range 0.001, 0.01, 0.1, 1, 10, 100, 1000 for neural-based models, including ITSELF and GU IDE. Notably, higher α values better constrain outputs to binary, thereby facilitating near-optimal results.The current approach can be extended to support both multi-valued discrete and continuous variables, broadening its utility in diverse scenarios.For multi-valued discrete variables, the method can be adapted by implementing a multi-class, multioutput classification head. Each query variable is represented by a softmax output node, which provides soft evidence by generating probabilistic distributions across multiple discrete values.To incorporate continuous variables, we introduce a linear activation function in the output layer. The loss function, specifically the multi-linear representation of the PM, is modified to accommodate continuous neural network outputs. For example, in Probabilistic Circuits that use Gaussian distributions, continuous values can be directly integrated into the loss function, facilitating gradient-based backpropagation.These extensions primarily involve adjusting the network's output layer and refining the selfsupervised loss function represented by the PM. Notably, other elements of our approach, including the ITSELF and GUIDE procedures, remain unchanged.Our approach further extends to additional inference tasks over probabilistic models, including marginal MAP and constrained most probable explanation (CMPE) tasks. However, the scalability of this approach depends on the computational efficiency of evaluating the loss function for each inference task. When this evaluation becomes computationally infeasible, the proposed method-training a neural network to answer queries over probabilistic models-may itself become infeasible. For example, performing marginal MAP inference over NAMs and PGMs requires repeated evaluations of the loss function associated with the marginal MAP task and its gradient during training. This iterative process, essential for updating the neural network's parameters, can become prohibitively resource-intensive due to the high computational demands of evaluating the marginal MAP loss over these probabilistic models.This section evaluates the performance of models initialized through various techniques-random initialization, SSMP, and GU IDE. Each plot represents the loss for a distinct test example, with the x-axis denoting the number of ITSELF iterations and the y-axis showing the Negative Log Likelihood (NLL) scores. Lower NLL values signify better solutions. Through this empirical assessment, we compare the impact of different pre-training methods on model performance. We randomly initialize the parameters for the random model and perform 1000 iterations of ITSELF for inference. For the two pre-trained models (SSMP and GU IDE), we update the top N layers, where N is the number of layers corresponding to that loss curve, and fix the remaining bottom layers. We extract features by passing the input through these fixed layers and then train the parameters of the top N layers. We again perform 1000 iterations of ITSELF for inference. For NAMs and PGMs, we use neural networks with up to one hidden layer, while for PCs, we employ models with up to three hidden layers.From the plots for the three Probabilistic Modelss (PMs), we observe that models pre-trained using the proposed GU IDE training scheme generally have a better starting point for ITSELF, indicated by a lower loss, compared to all other models. Across a wide array of datasets, PGMs, and query percentages, the GU IDE method consistently converges to a lower or equivalent loss compared to other models. Remarkably, it sometimes achieves a loss value that is less than half of the nearest competing model. Furthermore, the losses for GU IDE are typically more stable than those of other initialization. In some scenarios, all models achieve a similar final loss, although models initialized with SSMP and those randomly initialized may experience oscillations in their loss values.Models pre-trained using the traditional self-supervised loss (SSMP) typically have better or similar starting points than randomly initialized models. However, models pre-trained using the SSMP method might converge to a worse loss than their GU IDE pre-trained counterparts.In most cases, convergence is rapid, even with a reduced learning rate of 10 -4 compared to the experiments shown in the main paper. Most methods converge within 200 to 300 iterations, although some may still oscillate during the later iterations of ITSELF. Dataset: Audio, Query Percentage: 0.9 Each cell displays the natural logarithm of the time, measured in microseconds, for each method and dataset. Lighter colors indicate lower values. Notably, inferences using SSMP and GU IDE require the shortest time, as these methods necessitate only a single forward pass through the neural network to obtain the values for the query variables.For MADE, the subsequent fastest method employs a model trained with GU IDE and conducts inference using ITSELF, outperforming the approach that uses SSMP for training. This advantage stems from the reduced number of ITSELF iterations required by GU IDE, benefiting from a more effectively trained model. In PGMs, a similar pattern emerges with GU IDE + ITSELF as the next fastest method, followed by SSMP + ITSELF. For PCs, MAX ranks as the next fastest, closely followed by the GU IDE + ITSELF and SSMP + ITSELF methods. Finally, the ML and Seq methods display the highest inference times.Thus, if you require a highly efficient method capable of performing inference in a fraction of a millisecond, GU IDE is the optimal choice. It outperforms the baseline for both MADE and PGMs. However, if higher log-likelihood scores are necessary, GU IDE + ITSELF would be suitable, as it generally surpasses the baselines in speed and performance across various scenarios.Table 2 presents the log-likelihood score gap between the neural network methods (SSMP, GU IDE, SSMP + ITSELF, GU IDE + ITSELF) and exact solutions. These exact solutions are obtained using AOBB, which provides near-optimal results for smaller datasets. For each approach M, the gap is calculated as the relative difference between the score of the near-optimal solution (determined by AOBB) and the score achieved by M. This approach is feasible due to the use of small datasets, allowing identification of exact solutions.The final column highlights the neural-based approach achieving the best performance for each dataset and query ratio combination. Notably, GU IDE and ITSELF consistently surpass other neural baselines across almost all dataset-query pairs. This analysis provides a comprehensive assessment of the proposed methods relative to exact solutions on small datasets, enabling a direct comparison of their effectiveness. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).38th Conference on Neural Information Processing Systems (NeurIPS 2024).Arya et al. [4] developed a NN-based method for solving the marginal maximum-a-posteriori (MMAP) task in PCs. In this paper, we focus on the MPE task, also sometimes referred to as the full MAP task, which is a special case of MMAP. Our method can be easily extended for solving the MMAP problem in PCs and tractable graphical models. For simplicity of exposition, we concentrate on the MPE task in this paper.or a value proportional to it such as the unnormalized probability in Markov networks.Note that we are referring to convexity with respect to the outputs, not the parameters of the NN. 


Decision: There is agreement among the reviewers that the paper addresses an interesting and important problem and proposes an entirely new approach to answering MPE queries. Some concerns have been adequately addressed during the rebuttal.

Reviews: review: paper_summary: paper_summary: This paper proposes a novel neural network-based method to solve the computationally challenging task of finding the Most Probable Explanation (MPE), which is known to be NP-hard. The proposed method involves an inference-time optimization process with a self-supervised loss function to iteratively improve the solutions. It also employs a teacher-student framework that provides a better initial network, which in turn helps reduce the number of inference-time optimization steps. Experiments demonstrate the efficacy and scalability of the proposed method compared to various baselines across different datasets.
strength_weakness: strength_weakness: The presentation of the main experimental results could be improved. The captions of Figure 1 and Figure 2 are quite simple and do not provide enough details about the information conveyed in the figures. It would be better to include more descriptive captions that explain the meaning of the rows, columns, and color schemes used in the visualizations.
questions: questions: 1. Could the authors provide more detailed explanations or visualizations to help the reader better understand the experimental setup and the interpretation of the results presented in Figures 1 and 2?
2. Are there any plans to further extend or generalize the proposed method to handle other types of probabilistic inference tasks beyond the MPE problem?
limitations: limitations: N.A
paper_summary: paper_summary: This paper proposes a neural network-based approach to approximate the most probable explanations (MPE) on probabilistic models. The basic idea is to train a neural network that takes as input an encoding of an assignment to an arbitrary subset of evidence variables and outputs an assignment to the remaining query variables, with the objective to maximize the likelihood given the predicted assignment. This is feasible on probabilistic models that support tractable likelihood and gradient computations such as PGMs, PCs, and NAMs. The authors present two key strategies to enhance the basic architecture. First, they employ a teacher-student architecture (GUIDE) that trains a student network by supervised learning with the teacher network’s MPE solutions as labels. In addition, they introduce inference time optimization to further refine the pre-trained models given specific MPE queries. The proposed methods are evaluated empirically on benchmark datasets and compared against baseline approximation methods.
strength_weakness: strength_weakness: - It appears that the inference time optimization is doing a lot of the heavy lifting compared to pretraining. On experiments on PCs, the simple baseline MAX, which is a popular method for PCs, seems to outperform both SSMP and GUIDE without inference time training. The same is true, albeit to a lesser degree, for MADE.

- Evaluation based on average joint likelihood p(q,e) may be biased towards the neural-network based approaches, compared to evaluating using average conditional likelihood p(q|e). Because the loss function is based on joint likelihood, it will bias the network training to perform better on evidence e with larger probability p(e) than those with smaller probability. Average joint log-likelihood could easily hide poor results on instances with low probability evidence.


- The significance of Propositions 1 & 2 is unclear. More useful in motivating the training objective might be showing that the chosen loss function upper-bounds the target loss (negative log-likelihood) that we want to minimize.
questions: questions: 1. Can the proposed approach be extended to continuous variables?

2. In each iteration of GUIDE, the teacher network is reinitialized with the student network. Could this result in the network being stuck at a local minimum?

3. How do the results compare with exact solutions (on model and problem instances that can be solved exactly)?

4. How are the results if evaluating using conditional likelihood p(q|e) instead of p(q,e)?

5. How many unique query instances (evidence variables and assignments) were used in the training data? (ie what proportion of all possible instances)

Other comments:
- I would suggest adding a diagram showing the components of the proposed method to improve clarity.
- Figure 2: typos in the percentage difference labels (e.g., >-10 and <10 should be >10 and <-10).
- Why was this particular input encoding chosen? For example, a natural alternative using an input node to indicate whether a variable $X_i \in \mathbf{E}$ and another input node to denote its value $x_i$.
limitations: limitations: The current approach is currently limited to binary variables.
paper_summary: paper_summary: This paper proposes an approximate MPE inference scheme for multiple ***discrete*** probabilistic models that guarantee efficient log-likelihood computations. The approach is based on training a carefully-designed neural network to answer these queries, building upon three key ideas:
1) self-supervised learning, in order to avoid labelling via expensive MPE oracles
2) ITSELF: iterative optimization of the NN parameters at prediction time, resulting in both anytime inference and continual learning
3) GUIDE: a teacher-student training scheme for mitigating overfitting

The approach is evaluated on multiple binary benchmarks from the TPM and PGM literature.

**update after rebuttal**

The authors effectively addressed the concerns on the paper. Overall, I think that this is solid work and I hope that the authors consider adding their explanation on generalizing the approach to continuous variables and the runtime results in the main text.
strength_weakness: strength_weakness: - The limitations of the approach are not clearly stated.
- Some experimental details are unclear
- Runtime evaluation, a crucial aspect of any approximate algorithm, is deferred to the supplementary work
questions: questions: 1) "Without loss of generality, we use binary variables" How is this WLOG? Does the approach seems limited to discrete models. If that's the case, I would be more upfront about that.
2) Figure 1: what are the numbers in the contingency tables?
3) How often is the approach recovering the true MPE assignment? Why not devising an experiment where ground-truth solutions can be obtained?
4) What is the training time of the proposed approach?

---

### Minors:

Lines 139-143: is there a difference between the NN relaxation and the PGM/PC? It looks identical on a first glance.

Figure 2: Shouldn't the caption be "Top - PC, Bottom - ***NAM***"? Also in b), shouldn't it be "GUIDE + ITSELF ***vs.*** HC"?
limitations: limitations: The limitations of the approach are not clearly stated in the text.

Comments: title: Response to AC
Response to the authors
Please read and respond to the rebuttal
comment: Dear AC,

Sorry to hear that you found my review shallow. To clarify, Q1 refers to Weakness 1, Q2 and Q3 refer to Weaknesses 2 and 3, and the authors adeguately addressed those concerns in their response. I should have mentioned that.

Since the discussion time was (is) not over, I was simply taking my time to formulate further questions. 
I don't see the hurry in immediately raising the score, we have a 9 days long AC/reviewers discussion phase ahead.
We believe what you are suggesting is to evaluate the impact of instances with relatively low probability of evidence on the performance of our method. Your idea of using average conditional likelihood instead of just average log likelihood is an excellent one and will help address the influence of low-probability instances on our method’s performance. While this is definitely a step in the right direction, it might not fully capture the entire picture.

To build on this, and because your suggestion has inspired further thinking, we plan to use techniques like slice sampling or SampleSearch to specifically generate and evaluate these low-probability instances. This approach will allow us to thoroughly assess our model’s performance, particularly in handling rare events. We believe ITSELF and GUIDE will perform better on such instances because they use test time training.
Thank you for the detailed response. Most of my comments have been address, but I have a follow up regarding comment 2.

>Furthermore, using the conditional likelihood $P(q|e)$ as the loss function would be identical to using the likelihood since $\log P(e)$ does not contribute gradients to the loss.

While this is true when $e$ is a single fixed evidence, the data consists of multiple instances with different evidence assignments. In average joint log-likelihood, $\log P(q,e)$ for some evidence $e$ with high probability $P(e)$ could dominate $\log P(q',e')$ for another evidence $e'$ with much lower probability $P(e')$. Average conditional likelihood on the other hand weighs instances equally regardless of the probability of evidence of each instance. The dataset is constructed by sampling from the PM, which is likely to sample evidence assignments with higher probability more often, and average joint likelihood would weigh those instances even more. Thus, comparison by average joint log-likelihood could theoretically hide poor performance on low likely instances.

While the discussion period is closing soon, I would highly encourage the authors to consider also evaluating with average conditional likelihood in the revision.
Dear reviewer,

Your initial review was short and shallow. For instance, you wrote that some experimental details are unclear without giving examples or pointers. Now, after the authors have provided a rebuttal, answering all of the weaknesses you mentioned and your questions, you simply answer that you don't have any more questions, but without raising your score. Unfortunately, I will not be able to consider your review in my decision making unless you provide more substance and concrete reasons why the paper should not be accepted.
Thank you for clarifying some aspects of your work. I don't have any further questions at the present time.
Dear reviewers,

First of all, thank you for your service to the ML community. Writing high-quality reviews and engaging with authors' responses is essential to a healthy ML research community. 

It is essential that you read the rebuttals and provide a response and/or follow-up questions within the next few days. This will allow the authors sufficient time to react. While a detailed response addressing all points is not necessary, at a minimum, you should indicate you have read and considered the review and whether you will maintain or revise your score. Please also take the time to read the other reviews. Understanding your fellow reviewers' key arguments can help converge towards more similar ratings in cases of diverging scores.

I want to thank you again, and I look forward to following the discussions here.