Gaussian Ensemble Belief Propagation for Efficient Inference in High-Dimensional, Black-box Systems

Authors: Dan MacKinlay, Russell Tsuchida, Daniel Edward Pagendam, Petra Kuhnert
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: factor graph, data assimilation, ensemble kalman filter, gaussian belief propagation, geospatial, probabilistic graphical model, hierarchical model

Abstract:
Efficient inference in high-dimensional models is a central challenge in machine learning.
We introduce the Gaussian Ensemble Belief Propagation (GEnBP) algorithm, which combines the strengths of the Ensemble Kalman Filter (EnKF) and Gaussian Belief Propagation (GaBP) to address this challenge.
GEnBP updates ensembles of prior samples into posterior samples by passing low-rank local messages over the edges of a graphical model, enabling efficient handling of high-dimensional states, parameters, and complex, noisy, black-box generative processes.
By utilizing local message passing within a graphical model structure, GEnBP effectively manages complex dependency structures and remains computationally efficient even when the ensemble size is much smaller than the inference dimension --- a common scenario in spatiotemporal modeling, image processing, and physical model inversion.
We demonstrate that GEnBP can be applied to various problem structures, including data assimilation, system identification, and hierarchical models, and show through experiments that it outperforms existing belief propagation methods in terms of accuracy and computational efficiency.

Supporting code is available at https://github.com/danmackinlay/GEnBP}{github.com/danmackinlay/GEnBP

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730135196623

thread_type:
official_review

<thread_object>
creator: reviewer_qpy9

<content>
summary:
Efficient inference is challenging in graphical models with high-dimensional variables. 
The authors introduce an inference algorithm based on Gaussian Ensemble Belief Propagation (GEnBP) and the Ensemble Kalman Filter (EnKF), combining strengths of both algorithms to arrive at an algorithm which is competitive in accuracy, in scalability with respect to the dimensionality of the variables, and which can be applied to a wide range of inference problem structures.

strengths:
The authors present a novel methodology with well motivated approximations and algorithmic structure, where important connections and details are highlighted. The methodology contains contributions and explanations that are valuable also on their own. Graphical illustrations and notation is used in an effective manner to improve clarity, and the text is well structured.

weaknesses:
Although I find the text to be well structured in terms of its content, the text contains many language errors, forgotten punctuation, and incorrect figure references. I have listed some errors I encountered below which the authors may find helpful for improvement of the manuscript.  

List errors/suggestions style and language:
- Line 22: no spaces around “–”. 
- Line 48: why “seems” novel and not “is” novel, if that is the case? 
- Line 68: “table” is not capitalised.  This is also the case in many other places. Note that it should be capitalised since “Table 1” is a name.   
- Line 71: “section” not capitalised.  This is also the case in many other places. 
- Line 72: missing punctuation after “approximations”.
- Line 81-95: In my opinion, excessive use of bolds. Is it needed more than say a couple of times here, or even at all?
- Line 83: How about “a novel message-passing method for inference in high-dimensional graphical models”?
- Line 91: “Rank” is (incorrectly) capitalised even though the list is written as a sentence. 
- Line 93: Same as above but for “Compute”. 
- Line 99: How about “We will now introduce our notation and essential concepts”?
- Line 100: “appendix” is not capitalized. This is also the case in many other places. 
- Line 118: How about “We will assume that queries are always ancestral variables, i.e. that …”?
- Line 134: “equation” is not capitalized.  This is also the case in many other places. 
- Line 146: Missing word 'with'. “Edges connect each factor node j with each variable node…”? 
- Line 173: Unnecessary abbreviations, “we’ve” and “we’re”.
- Line 180: observed is misspelt.
- Line 184: “figure” not capitalised.  This is also the case in many other places. 
- Line 191: Missing punctuation before “The following”.
- Line 202: “definition” not capitalised. This is also the case in many other places. 
- Line 227: Double punctuation. 
- Line 261: Missing punctuation before “Marginalization”. 
- Line 294: Missing punctuation before “Throughout”. 
- Line 312: Missing punctuation before “Here”. 
- Line 314: Misspelt equivalently. 
- Line 323: Incorrect grammar, “See appendix I for a comparison with a naive attempt to do without the ensemble”.
- Line 325: Missing punctuation before “Then”. 
- Line 336: Missing punctuation after “high dimensions”. 
- Line 353: “algorithm” not capitalised. This is also the case in many other places. 
- Line 376: Incorrect grammar: “For a variable with K neighbors, in the worst case M = KN , when the cost becomes…” and missing punctuation. 
- Line 381: How about “In many practical applications it holds that N ≪ D, resulting in significant computational savings compared to GaBP which scales poorly with D.”. 
- Line 385: How about “better suited”?
- Line 404: Incorrect grammar “While GEnBP scales favorably with D, but unfavourably with respect to the node degree K, scaling as O(K 3 ) for some operations.”


The presentation of the experiments is unclear at times, please see Questions.

questions:
Regarding Section 4.1:
- Figure 2 is never referenced. Although I assume the first paragraph describes observations made from it? 
- Why is the log likelihood only shown for GaBP in the case of the lowest dimensionality setting used? 
- I struggle to reach the following interpretation based on Figure 2: “We see that the Laplace approximation … but its posterior likelihood, while similar to GaBP, is even less stable, and both are inferior to GEnBP”:
- - For dimensionalities greater than around 100, it seems GEnBP is underperforming Laplace in terms of log likelihood (while the results provides no comparison to GaBP)? 
- - Why is GaBP and Laplace interpreted as less stable than GEnBP based on this figure? 

Section 4.2:
- Figure 3 is from the text implied to contain results using Langevin MC, but it does not --- only Figure 4 does. 
- The results show (and the text says) that GaBP performs better in terms of log likelihood for high viscosity settings. What do you suspect is the explanation for this? 

Overall, I think the paper is good and makes valuable contributions. However, I think the manuscript needs an overhaul to fix the language errors, missing punctuation, and figure references before it is ready for publication. If this is done and the experiments section is clarified I intend to raise my evaluation score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730412958660

thread_type:
official_review

<thread_object>
creator: reviewer_mt1v

<content>
summary:
The paper presents a generalisation of the ensemble Kalman filter to make inference on arbitrary graphical models using belief propagation  (BP) with ensembles. By relying on ensembles rather than keeping track of the full moments, the method is able to scale to problems where the state dimension of the nodes are high dimensional, e.g. weather forecasting. The authors propose several numerical tricks that make this possible while retaining scalability. This is tested against a standard Gaussian BP on some examples, showcasing its computational efficiency and performance improvements, offered by its improved handling of nonlinear conditional dependence between nodes.

strengths:
The paper is generally well-written and well-motivated. The methodology presented, while seemingly a small step beyond standard BP, is not at all trivial to execute, requiring various tricks to make it work in practice. In particular, I find the efficient computation of the factor-to-variable message and the ensemble conformation step to be an interesting contribution that is essential to retain the ensemble-based representation of the beliefs. The experiments showcase interesting settings where this technique could be applied, such as latent forcing identification in fluid models. In addition, it shows clear benefits of the approach over vanilla BP, with orders of magnitude increase in computation speed and performance gains.

weaknesses:
While the paper is clearly written for the most part, there are some parts that I find need more explanation. See my questions below for details. I also find that the experiments can be improved with further baseline comparisons. Is GaBP really the only baseline that can solve the system identification problems in the experiments? One may probably also consider methods like the integrated nested Laplace approximation (INLA), or particle MC methods (e.g. the particle marginal Metropolis-Hastings algorithm).

questions:
- It would be great if the authors could please clarify why we need a particle representation of the beliefs when we can just make computations with the DLR representation? My understanding is that in EnKF, it is used to propagate the particles by the nonlinear dynamics model and update them using Matheron's rule, but what does this correspond to in the general BP framework (in terms of steps (5)-(7))?
- Do we need tricks like covariance localisation and inflation in this general BP setting? These are almost always used to make EnKF work robustly in high dimensions.
- In Figure 2, why are some results missing for GaBP and Laplace method in the log-likelihood comparison? Was it producing NaNs due to overconfident predictions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730528447542

thread_type:
official_review

<thread_object>
creator: reviewer_hdem

<content>
summary:
The paper introduces the GEnBP algorithm, a method that combines Ensemble Kalman Filter (EnKF) and Gaussian Belief Propagation (GaBP) to efficiently perform inference in high-dimensional systems. GEnBP updates ensembles of prior samples into posterior samples using a message-passing strategy on a graphical model, making it suitable for applications with complex dependencies and high-dimensional variables, such as in geospatial and physical models.

strengths:
The paper is well-structured, making it easy to follow the ideas presented. The proposed method can be considered to be novel in its approach to addressing the computational challenges of GaBP in high-dimensional systems.

weaknesses:
While GEnBP demonstrates strong performance in high-dimensional contexts, its complexity scales disproportionately with the node degree in graphical models. This presents a significant limitation in highly connected graphs, where methods such as Forney factorization have been proposed by the authors as partial solutions. However, the effectiveness of these methods remains uncertain. Can some simple numerical results or insightful analyses be added to showcase this potential solution?

Additionally, a notable drawback is the absence of comparative analysis with other related works, both those employing and not employing EnKF, for efficient inference in high-dimensional, black-box systems. Key related works include:

  - Chen Y, Sanz-Alonso D, Willett R. "Autodifferentiable Ensemble Kalman Filters," *SIAM Journal on Mathematics of Data Science*, 2022; 4(2):801-833.

  - Chen Y, Sanz-Alonso D, Willett R. "Reduced-order Autodifferentiable Ensemble Kalman Filters," *Inverse Problems*, 2023; 39(12):124001.

  - Lin Z, Sun Y, Yin F, Thiéry A. "Ensemble Kalman Filtering-Aided Variational Inference for Gaussian Process State-Space Models." *arXiv preprint arXiv:2312.05910*, 2023.

  - Girin L, Leglaive S, Bie X, Diard J, Hueber T, Alameda-Pineda X. "Dynamical variational autoencoders: A comprehensive review." *arXiv preprint arXiv:2008.12595*, 2020.

Including a discussion or comparison on system identification and data assimilation from these (and other existing) works would enhance the comprehensiveness of the proposed method's positioning.

questions:
1. Figure 2 is not cited in the text, which may confuse readers
2. Consider put Figure 1 in the experimental section
3. The paper highlights potential applications in fields like geospatial prediction and physical model inversion, which may involve real-time data. How does GEnBP handle streaming data, and what would be required to extend it to real-time applications with time-varying dependencies?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730614566879

thread_type:
official_review

<thread_object>
creator: reviewer_dsra

<content>
summary:
The authors introduce a novel method named Gaussian Ensemble Belief Propagation (GEnBP). They combine EnKF and GaBP together, in order to improve the efficiency in high-dimensional and non-linear cases. The method uses low-rank message-passing algorithms.

strengths:
The authors try to use the strength of GaBP to solve (or I would like to say mitigate) the bad performance of EnKP in high-dim and non-linear problems. In the given examples, the new method works well. The computational complexity reduces largely for the dimensions of the data. It seems that the paper demonstrates superiority over existing methods in various benchmarks (but the figures are hard to read). I believe this method is potential for real-world applications for example in geospatial modeling.

weaknesses:
1. Most important question for me: the scalability of the method to extremely large models, like weather simulations, remains to be fully demonstrated. Computational complexity still depends heavily on the degree of nodes in the graphical model, which means you must restrict the scale of the model. However, models in climate research always be extremely huge. How to choose N, how to choose K? If too small, can the model catches the nonlinearity well?
2. EnKF is somehow popular in climate or geophysical models, because of its computational complexity. But GaBP, in my impression, is a little bit outdated. For the nonlinear cases, we may use variational inference, or particle filter. Have you compared with them?
3. The innovation: GaBP mostly for low-dim problem. I am not sure it can help to improve the EnKF itself. More importantly, please see Bickson et al (2008) for the similar idea.
4. Even for small model, for example, Laplace seems better after consideration of trade-off. And Fig 2(c) seems not complete.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732116075431

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response 1

comment:
Thank you for your detailed reviews and suggestions for comparing our Gaussian Ensemble Belief Propagation (GEnBP) method with various existing techniques. We value your insights, which help clarify our method's place within the broader research landscape.

There are  themes which have arisen generally in the reviews which we feel deserve a general response, which we would classify as “general questions about the scope of the claims.”

Several suggested comparisons, including autodifferentiable EnKF and dynamical variational autoencoders, are indeed valuable within system identification and state filtering. However, GEnBP targets a broader application set within the "belief propagation" methods, applicable to more general graphical models, as exemplified in Figure 8 of the appendix.

We acknowledge that our initial problem *is* a problem of system identification type, which we have selected because of its popularity. It would be unlikely for GEnBP to dominate in general in system identification over all possible methods designed specifically for that task.
For this case, demonstration that the GEnBP method, devised for arbitrary graphs, happens to be competitive with existing methods customised for system identification problems is, we argue, sufficient to establish our claim as made.

Unlocking efficient inference in high-dimensional graphical models is a high value goal for all the reasons that belief propagation is a valuable tool in the first place; e.g. we can can estimate intervention effects, impute missing values, distribute computation over many nodes and so on.

That said, we are sympathetic to the idea that if we wish to make a claim about a more general graphical model structure, there is a case for including example problems more general than the system identification problem. We have made the tactical choice not to prioritise that because the methodology section of the paper is already dense, and the appendices are already twice as long as the paper itself. We are happy to consider this for future work.

To align our manuscript with these clarifications, we propose the following revisions:

1. **Introduction and Discussion:** We will clarify GEnBP’s role within belief propagation, differentiating it from methods geared specifically towards system identification or state filtering.

2. **Scope Clarification:** We will explicitly note that GEnBP is not designed to universally surpass all other methods on all graphical structure but to complement the existing toolkit for handling high-dimensional inference problems.

3. **Emphasizing Novelty:** The integration of ensemble methods with belief propagation in complex models is a significant innovation of our work, and we will highlight this contribution more prominently.

We appreciate the opportunity to refine our work and believe these revisions will make the paper's objectives clearer.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732273557387

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response 2: Manuscript updated.

comment:
We are grateful to the reviewers for their thorough evaluations and insightful comments on our manuscript.
We have prepared a revision. Changes are also colorized in the PDF. This pushes it over the 10 page limit, but when the change markers are removed the paper is still 10 pages long.

### 1. Scalability and Computational Complexity *(dsRA, HdeM)*

- **Clarified Scalability Considerations**: In response to **dsRA** and **HdeM**'s concerns about the scalability of our method to extremely large models and highly connected graphs, we have:

  - **Expanded the discussion in Section 3.4** to detail the computational complexity of GEnBP, emphasizing that while it avoids the \(O(D^3)\) scaling of GaBP, its complexity depends on the graph structure and node degree.

  - **Added to the conclusion**: “Comparing overall computational cost of the algorithms is complicated and depends upon the graph structure. Our empirical results demonstrate the existence of problems where GEnBP is much more efficient than GaBP, but we have not characterized the full range of problems where this is the case.”

### 2. Comparison with Related Methods *(dsRA, HdeM, Mt1V)*

- **Expanded Literature Review**: To address the absence of comparative analysis noted by **dsRA**, **HdeM**, and **Mt1V**, we have:

  - **Added references** to related works, including:  Chen et al.'s autodifferentiable EnKF methods, Lin et al.'s Ensemble Kalman Filtering-Aided Variational Inference, Girin et al.'s review on dynamical variational autoencoders, Methods like Integrated Nested Laplace Approximation (INLA) and particle Monte Carlo methods.

  - **Discussed the positioning of our work** relative to these methods, highlighting the unique aspects of GEnBP in handling high-dimensional nodes within arbitrary graphical model structures.

  - **Clarified why system identification methods were not directly compared**, explaining limitations in high-dimensional settings and differences in application focus.

### 3. Clarification of Novelty and Contribution *(dsRA)*

-  In response to **dsRA**'s mention of Bickson et al. (2008), we have **Clarified the novelty of GEnBP**, explaining that while Bickson's work focuses on parallel computation in sparsely connected graphs with low-dimensional nodes, our method addresses high-dimensional systems with non-linearities and high-dimensional nodes, and in fact already cites Bickson.

### 4. Improved Figures and References* *(dsRA, QpY9, Mt1V)*

- **Updated Figure 2**: Responding to **dsRA**, **QpY9**, and **Mt1V**, we have prepared a new version of Figure 2, addressing incomplete data and enhancing readability.

- **Moved Figure 2 to the appendix** and updated reference to streamline the main text, as it seemed to distract readers from the core content.

- **Revised Section 4.1**: We clarified that the purpose of the experiments was to demonstrate the functionality of GEnBP on simple, visualizable problems rather than claiming superiority in all scenarios.

 - **Explained missing log-likelihood values**: We clarified that missing values in Figure 2 are due to non-positive-definite covariance estimates, making the log-likelihood undefined. Our revised plotting method exludes undefined values and so has fewer missing ranges on the plot.

- **Corrected results in Figure 3**: For **QpY9**, we re-examined the high-viscosity results and found that the problem becomes ill-posed in this regime. We have updated the figure and discussion accordingly.

### 5. Explanation of Methodological Details *(Mt1V)*

- we have  **Explained the role of particle representation** in GEnBP, highlighting that it serves a similar function to the linearization step in traditional GaBP and is essential for handling non-linearities in high-dimensional nodes with reference to App. I and  Figure 11

 - **Acknowledged the potential benefits** of techniques like covariance localization and inflation, commonly used in EnKF, for enhancing GEnBP's performance in the conclusion


### 7. Stylistic and Language Improvements *(QpY9)*

- **Revised the Manuscript for Clarity**:

  - **Corrected grammatical errors and improved punctuation**, addressing **QpY9**'s concerns.

  - **Ensured consistency in terminology and notation** throughout the paper.

### 8. Future Work and Limitations *(All Reviewers)*

- **Acknowledged Method Limitations**:

  - **Emphasized that GEnBP is not universally superior**, and its performance depends on the specific problem structure and graph topology.

- **Outlined Directions for Future Research**:

  - **Suggested integrating techniques** from both GaBP and EnKF literature, such as graph pruning, robust distributions, and the aforementioned covariance localization and inflation.

  - **Proposed exploring hybrids** of GEnBP with other methods like INLA or particle methods, as suggested by **Mt1V**.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732634932457

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response 3: new revision.

comment:
We have uploaded a new revision of the manuscript. Differences from the initial version are colorized. This one includes fixes for typos introduced in previous revisions, and an expanded explanation on the rationale for performing the specific task of system identification with a generic belief propagation method.

To paraphrase our own conclusion, we have shown that traditional belief propagation methods are leaving performance "on the table" by not exploiting the DLR structure and convenient conditioning produced by ensemble methods, and moreover, we suspect that this innovation is largely "pluggable" into the existing ecosystem of practices in this domain. On that basis, we believe that we are able to provide immediate quality-of-life improvements to researchers applying machine learning to spatiotemporal data, and we hope that this work is published to disseminate that improvement.

If we are blessed with luck, some bonus experiments may land in the paper before the deadline. However, we believe that the paper capable of standing on its own merits as-is. We welcome any final engagement with or requests from our reviewers on what might induce them to share this opinion.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732726150075

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response 4: New experiment on domain adapatation

comment:
We would like to thank the reviewers for their thoughtful comments and for highlighting areas where our manuscript could be improved. As noted before, we believe that our reviewers have, to varying degrees, assessed the merit of this paper based on whether it represents the state-of-the-art (SOTA) in the system identification problem. While we acknowledge that this is an interesting application, we contend that the broader significance lies in whether this method advances the frontiers of generic belief propagation problems that cannot be treated solely within the existing system identification literature.

To address this, the new Appendix subsection (B.3) provides extensive explanations of the scope of a specific non-system-identification type. Specifically, we apply GEnBP to the problem of showcasing its applicability to generic belief propagation scenarios. The results demonstrate that GEnBP effectively handles a latent domain-adaptation problem for a high-dimensional neural network with high-dimensional outputs, thereby illustrating its broader utility beyond system identification.

This addition is referenced in Section 4 (Experiments) as follows:

>Although GEnBP is applicable beyond the system identification setting, this benchmark is chosen for its simplicity and popularity. More sophisticated graphical model problems are outlined in Appendix C, and in Appendix B.3 we demonstrate the utility of GEnBP applied to a different graphical structure: few-shot domain adaptation of a system dynamics emulator on unobserved states.

In the new subsection, we demonstrate the applicability of GEnBP beyond system identification by applying it to a problem involving few-shot domain adaptation of a neural network emulator on unobserved states, which requires the unusual ability of the GEnBP method to scale to nodes with millions of dimensions: specifically, the millions of weights in a neural network. We set up a graphical model where a physical simulator predicts system states based on unknown parameters, and a Bayesian neural network (using the Fourier Neural Operator architecture) is trained to emulate this simulator using noisy observations without access to ground truth. By employing GEnBP to approximate the intractable posterior distribution of the neural network weights, we adapt the emulator to the new target problem. Our experiments on a 32×32 grid show that, after applying GEnBP to the first five of ten simulated timesteps, the adapted emulator closely follows the target distribution in subsequent timesteps, while the unadapted emulator diverges. This illustrates GEnBP's effectiveness in more complex, high-dimensional settings, highlighting its broader utility beyond traditional system identification tasks.

The content of Section B.2 is approximately two pages and adds two additional diagrams (Figures 8 and 10); we refer the reviewers to the manuscript for details. The figures are draft quality and will require minor adjustments for the camera-ready copy, but the data should be unambiguous.

We are confident that this provides a solid response to the question about whether this paper is purely focused on system identification, and answers in the negative: **this is a paper about belief propagation, and the method is capable of solving interesting and non-trivial belief propagation tasks outside of system identification**. As such, this addition directly addresses the major outstanding issue raised by the reviewers, to the best of our ability to discern. We welcome commentary on this development.

We appreciate the reviewers' insights and believe that these enhancements strengthen our manuscript by highlighting the versatility and broader impact of GEnBP in the field of belief propagation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors introduce a novel method named Gaussian Ensemble Belief Propagation (GEnBP). They combine EnKF and GaBP together, in order to improve the efficiency in high-dimensional and non-linear cases. The method uses low-rank message-passing algorithms.

The paper received four reviews, and all reviewers agreed that the paper was above threshold for acceptance, with two arguing for acceptance. The authors satisfied several reviewers'' concerns and comments, with some raising their scores as a result. Working in some of that new information, space permitting, would make for a stronger final version.