Structure Learning with Adaptive Random Neighborhood Informed MCMC

In this paper, we introduce a novel MCMC sampler, PARNI-DAG, for a fully-Bayesian approach to the problem of structure learning under observational data. Under the assumption of causal sufficiency, the algorithm allows for approximate sampling directly from the posterior distribution on Directed Acyclic Graphs (DAGs). PARNI-DAG performs efficient sampling of DAGs via locally informed, adaptive random neighborhood proposal that results in better mixing properties. In addition, to ensure better scalability with the number of nodes, we couple PARNI-DAG with a pre-tuning procedure of the sampler's parameters that exploits a skeleton graph derived through some constraint-based or scoring-based algorithms. Thanks to these novel features, PARNI-DAG quickly converges to high-probability regions and is less likely to get stuck in local modes in the presence of high correlation between nodes in high-dimensional settings. After introducing the technical novelties in PARNI-DAG, we empirically demonstrate its mixing efficiency and accuracy in learning DAG structures on a variety of experiments.

Paper

Similar papers

Peer review

Reviewer PZBP7/10 · confidence 4/52023-07-01

Summary

By adding various elaborations to the state-of-the-art Markov chain Monte Carlo inference algorithm, this paper achieves an efficient and effective Bayesian network inference algorithm. Bayesian networks are one of the main tools of machine learning with a long history. In general, their learning is known to be a computationally hard problem, but efficient inference algorithms based on Markov chain Monte Carlo have been actively studied. Section 2 describes the problem setup of a linear functional Bayesian network with a normally distributed weight matrix as a concrete example. Section 3.1 introduces the baseline algorithm as a state-of-the-art inference algorithm for Bayesian networks, applying the method of Liang et al. [2022] for Bayesian variable selection problems to Bayesian networks. In Sections 3.2 through 3.4, the authors' further elaborations are carefully described step by step. Section 3.2 shows how to effectively adjust the neighborhood of the proposals, inspired by Liang et al. [2022] for Bayesian variable selection problems. Section 3.3 introduces an effective device for (nested) sequential sampling of neighborhoods to eliminate problems that can arise in the DAG estimation problem, but not in Bayesian variable selection. Section 3.4 presents a method for properly restricting the neighborhoods, which tend to be enormous.

Strengths

- This paper is a solid proposal for possible reasonable improvements in the Bayesian network inference problem, with broad coverage of the latest developments in the surrounding fields of Bayesian machine learning. - The text is very detailed so that a wide variety of readers (from beginners to experts) can follow the history and the latest developments in the field. - The code is provided in such a way that the proposed algorithm can be easily followed up by subsequent research, which is a very significant contribution to the field.

Weaknesses

The candidate weaknesses listed below are based on questions I had during my initial peer review. As my misconceptions are resolved, they may cease to be weaknesses. - The paper is somewhat unclear in its claims about the mixing time analysis (or empirical observation) of the proposed MCMC algorithm, although there are several mentions of it (e.g., Line 13, 288, 365).

Questions

Thank you for sharing this very interesting paper. This paper has a lot of devices to help the reader understand. I really enjoyed reading about many of the parts. On the other hands, for some parts, I also worry that perhaps I am underestimating the value of this paper due to my own lack of understanding. Therefore, I would like to present some arguments below to improve my own understanding. Authors do not need to respond to the inconsequential ones, but if there is anything that authors infer that my understanding is lacking, it would be very helpful if you could respond. - I find the author's argument against mixing times in the proposed algorithm somewhat unclear. In this paper, there are several mentions of mixing (e.g., Line 13, 288, 365). As the author states "PARNI-DAG quickly converges to high-probability regions (Line 11),” I interpret these claims as "the proposed algorithm achieves more rapid mixing times or shorter arrival time expectations to regions of high probability (note: these two claims are equivalent (Theorem 1.4 of [*]))”. [*] Yuval Peres and Perla Sousi. Mixing times are hitting times of large sets. Journal of Theoretical Probability, 28(2):488–519, 2015. However, I think that in general (especially when it depends on input data) analyzing or empirically observing MCMC mixing times is not an easy task. To my knowledge, there are not many examples of MCMC mixing times being analyzed for many problems in general. The exceptions are the following few problems: - Bayesian variable selection: Yun Yang, Martin J. Wainwright, and Michael I. Jordan. On the computational complexity of high- dimensional Bayesian variable selection. The Annals of Statistics, 44(6):2497 – 2532, 2016. - Bayesian community detection: Bumeng Zhuo and Chao Gao. Mixing time of Metropolis-Hastings for Bayesian community detection. Journal of Machine Learning Research, 22:10:1–10:89, 2021, - Structure learning with directed acyclic graph: [+] Quan Zhou and Hyunwoong Chang. Complexity analysis of Bayesian learning of high-dimensional DAG models and their equivalence classes, Annals of Statistics, 2023 (or arXiv:2101.04084). Fortunately, the DAG inference problems addressed in this paper seem to be able to guarantee polynomial-time mixing times for certain algorithms (Theorem 6 of the above literature [+]). In this light, what observations can be made about the mixing time of the proposed MCMC algorithm? Can existing results for mixing times be easily applied to the proposed algorithm? Or is theoretical analysis of the proposed algorithm for mixing time a future issue? Also, as we can observe from the experimental results, can the proposed algorithm speed up the mixing time by an order of magnitude that depends on the problem size? Or is it an improvement of constant orders? Fortunately, theoretical analysis of MCMC mixing time for DAGs seems to be in progress, so we hope that the author's mention of these issues (separating what is known from what is not known) will provide important insights for the reader.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

As discussed in the above Questions, I am not sure the (theoretical) guarantee of mixing time of the proposed MCMC algorithm.

Reviewer 1zPb4/10 · confidence 2/52023-07-06

Summary

This paper proposes PARNI-DAG, a new MCMC method for sampling Directed Acyclic Graphs (DAGs) that can be used for the problem of structure learning under observational data. PARNI-DAG builds on top of PARNI, and similarly uses locally informed, adaptive random neighborhood with an efficient point-wise implementation, but introduces additional improvements include: pre-tune sampler parameters using a skeleton graph derived from other methods, augmenting the search neighborhood with an edge-reversal move, and neighborhood thinning to improve computational efficiency. Experiments on some toy datasets demonstrate advantage over existing baselines.

Strengths

The proposed method seems to be an effective adaptation of PARNI to do Bayesian learning on DAGs, and outperfoms a few (classical) baseline methods on a few toy benchmarks.

Weaknesses

I am not an expert in this field, but it seems to me this paper is largely adapting the existing PARNI method for Bayesian variable selection for DAG learning, and combining various techniques that are already present in the current literature on top of PARNI. It would be helpful if the authors can make a more compelling case for the contributions of the paper over existing work. For example, in L62-L72 contribution, it seems up until L69 it is just describing what PARNI already has. The procedure for pre-tune sampler parameters and do warm start also seem like a trivial application of exsiting method. In Section 3 titled "The novel PARNI-DAG" proposal, the entire Section 3.1 seems to be just a recap of wht PARNI already has. Section 3.2 is mostly a recap of Kuipers et al. [2022]. Section 3.3 introduces the reversal neighborhood but this is also done in for example partition MCMC. Section 3.4 also seems like a simple adaptation of the thinning procedure that is already present in PARNI. In fact a large part of Section 3 seems to belong to background review rather than a description of a novel method. It would be helpful if the authors can reorganze this way and also clearly state what exactly is the novel contribution of PARNI-DAG.

Questions

See weaknesses

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

1 poor

Limitations

The authors did not discuss limitations

Reviewer xD8f5/10 · confidence 4/52023-07-07

Summary

This paper presents a Markov chain Monte Carlo (MCMC) sampler adapted from previous PARNI sampler, called PARNI-DAG, which is designed for Bayesian structure learning under observational data. The authors assume causal sufficiency and target the posterior distribution on Directed Acyclic Graphs (DAGs). The proposed PARNI-DAG mainly relies on the PARNI sampler and modify it to suite the purpose of structure learning, including 1. warm-start of neighbourhood sampling probability 2. reversal neighbourhood step and 3. adaptive neighbourhood skipping probability. The main contributions of the paper are 1. A warm-start procedure of the sampler's parameters that exploits skeleton graphs derived through constraint-based or scoring-based algorithms, ensuring better scalability and mixing property with the number of nodes. 2. A reverse step to avoid getting trapped in the local mode. 3. An adaptive skipping probability such that not all intermediate neighbourhood sampling steps are executed. Empirically, the author demonstrates the advantage of PARNI-DAG using real-world examples, which shows advantages compared to the previous MCMC methods.

Strengths

## Originality The originality of the paper lies in the development of the PARNI-DAG algorithm, which combines the Point-wise Adaptive Random Neighborhood Informed (PARNI) proposal with new features specifically designed for structure learning in the space of Directed Acyclic Graphs (DAGs). The proposed algorithm addresses the challenges of mixing and convergence in high-dimensional settings, which is not adequately addressed by existing MCMC methods for structure learning. This work does not provides a completely new sampling algorithm, instead, modify the existing approaches to suite the purpose of structure learning. ## Clarity The paper is logically structured with clear presentation of the modifications made to the PARNI proposal. The authors have made efforts to provide intuitive explanations and motivate their choices in the development of the PARNI-DAG algorithm. The paper is easy to follow, and the appendices provide additional details on the derivations and calculations. ## Significance The significance of the paper lies in its potential impact on the field of structure learning and causal discovery. The PARNI-DAG algorithm aims to address the challenges of mixing and convergence in high-dimensional settings. The algorithm's improved performance over existing MCMC methods for structure learning should make it a reasonable contribution to the field but there are some limitations, which I will elaborate in the following.

Weaknesses

## Empirical experiments: While the experimental results demonstrate the advantages of PARNI-DAG over existing MCMC methods, the experiments primarily focus on the comparison with MCMC based approach. However, in the literature review, the author also mentioned several structure learning approach. Although the main claim of this paper is the improvement over existing MCMC, it is still beneficial to include a comparison to state-of-the-art Bayesian structure learning approach like [1,2,3]. [1] Cundy, C., Grover, A., & Ermon, S. (2021). Bcd nets: Scalable variational approaches for bayesian causal discovery. Advances in Neural Information Processing Systems, 34, 7095-7110. [2] Lorch, L., Rothfuss, J., Schölkopf, B., & Krause, A. (2021). Dibs: Differentiable bayesian structure learning. Advances in Neural Information Processing Systems, 34, 24111-24123. [3] Geffner, T., Antoran, J., Foster, A., Gong, W., Ma, C., Kiciman, E., ... & Zhang, C. (2022). Deep end-to-end causal inference. arXiv preprint arXiv:2202.02195. ## Limitation with linear model The proposed PARNI-DAG mainly targets at the linear model with Gaussian assumptions because the marginal probability are needed. For general nonlinear model, such integration is not tractable, which rendering the following PARNI-DAG proposal invalid. However, nonlinearity is everywhere in the real-world applications. To make the paper stronger, the author should discuss the implication of using linear model or demonstrates that linearity assumption does not harm the performances too much. This can be achieved by comparing some of the previous mentioned baselines [1,2,3] ## Computational complexity Although the proposed PARNI-DAG method uses many modifications to reduce the computation cost, this approach relies on the MH step to correct the bias. It is known that MH step scales linearly with the number of datapoints, which can be a huge computational bottleneck for large dataset. I suggests the author should explicitly discuss the computational complexity or limitations, and potential approach to remove this constraints.

Questions

All the questions are mentioned in the weakness section.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The author does not explicitly write a limitation section, I have made suggestions to include one like computational challenges and limitations of using linear model.

Reviewer xD8f2023-08-16

Thanks for the authors' feedback regarding the additional experiments. They addressed my concerns and promised to add new limitation section. So I will keep my current evaluation.

Authorsrebuttal2023-08-17

Response to Reviewer xD8f

Dear Reviewer xD8f, Thank you for your reply. We remain available in case any further clarification is needed. Kindest Regards, \ The Authors

Reviewer PZBP2023-08-17

I appreciate the author's very detailed and helpful responses.

I appreciate the author's very detailed and helpful responses. The detailed additional explanation for the MCMC mixing time exceeded my expectations. All my concerns have been addressed. Thank you very much. I am sure that these explanations will satisfy the more expert readers, although I believe that the original paper also entertained a diverse audience ranging from newcomers to the field to experts in Bayesian networks with a long history of research. I would keep the score unchanged from my initial positive impression. Thank you for sharing a solid and important paper.

Authorsrebuttal2023-08-17

Response to Reviewer PZBP

Dear Reviewer PZBP, We are very happy the response has clarified any outstanding concerns about the work. We are once again thankful for the positive comments, but mostly for the interesting discussion these have sparked about mixing time bounds of PARNI-like proposals. This will turn out to be very useful also for related ongoing (and future) research on the topic. We remain available in case any additional query arises! Kindest Regards, \ The Authors

Reviewer 1zPb2023-08-19

Thanks for the response

I thank the authors for the response. However after reading the response my confusion still remains. For example the authors mentioned that `Applying a vanilla PARNI proposal to a DAG learning problem results in sub-optimal performance and particularly poses serious concerns about the sampler’s scalability` but it's not immediately clear to me what they mean by `sub-optimal performance` and what are the `serious concerns`. And with all the additional explanations this work still seems like a simple adaptation of existing techniques. But as I mentioned in the initial review, I am not an expert in this area. I will keep my score as is since the rebuttal does not seem convincing to me (and in may places seems to be making claims without any explanations/support). But it is possible that I do not fully grasp the challenges the authors have to overcome in the relevant adaptions, and I would be happy to hear feedback from other reviews.

Authorsrebuttal2023-08-21

Response to Reviewer 1zPb

Dear Reviewer 1zPb, Thank you for your reply. We would just like to clarify the specific points raised in the response that might be source of confusion. - By ``serious concerns`` in relation to vanilla PARNI sampler’s scalability, we mean both its **higher computational costs** and its **slower mixing**. Issues concerning the higher computational costs of vanilla PARNI are addressed via the changes introduced in Section 3.2 (warm-up PEPs estimates) and 3.4 (directly controlling neighbourhood size). As an example relative to the changes implemented in section 3.2 alone, vanilla PARNI, where the PEPs are computed adaptively via Rao-Blackwellised estimates, scales linearly in the number of edges and quadratically in the number of nodes. PARNI-DAG instead scales strictly less than that, with computational savings’ attributable to the use of warm-up PEPs estimates that depend on the starting search space obtained from the PC algorithm. Section 3.3 instead deals with slow mixing issues. - By ``sub-optimal performance`` we mean that vanilla PARNI yields a significantly lower accuracy (in addition to longer computational time) in DAG learning tasks compared to PARNI-DAG (and closer to ADR’s one). This is naturally attributable to its slower mixing properties mentioned in the point above. Kindest Regards, \ The Authors

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC