Streaming Bayes GFlowNets

Summary

Generative flow nets are extended to streaming inference by using the prediction of the previous network as the prior for the current one. Two training methods are proposed: streaming balance which uses squared difference of target and learned log posteriors, and VI which uses KL divergence. Analytic results for both methods bound the error in terms of updating error and error from the previous step. Three experiments show the method can accurately track the true posterior across updating steps.

Strengths

The methods are elegant and sound and the experiment results are strong. This is an important alternative to streaming VI that I'm excited to try out myself.

Weaknesses

Not a major weakness, but the extension to the streaming case (this paper's main contribution) is fairly straightforward. It was clear from the beginning that the predicted distribution from the previous step would provide the prior for the next, with the obvious adaptations of the balance and VI objectives. Both $p_B$ terms in the streaming balance condition (3) appear to be unnecessary. Prop 1 will hold without them. We just need a joint distribution over $(x,\tau)$ (equivalently, a distribution over $\tau$) with the desired marginal on $x$, and $f(\mathcal{D}_{t+1}|x) p^t_F(\tau)$ already has that property. Compare to the definition of the target distribution $p(\tau)$ in (5). Sec 3.2 just gives the objective and doesn’t explain how $p_F$ and $Z$ are learned (likewise with the argmax lines in Algo 1). Is it just SGD on $\mathcal{L}_{SB}$? The main text says nothing about how to construct the graph ($\mathcal{S}$ and $\boldsymbol{A}$). Details are given in the appendix -- in all cases the allowable paths build the desired objects from simpler ones -- but a sentence in the main text would help.

Questions

(2): why do we need to sample from $p_B$ and importance weight, rather than sampling directly from $p_F$? (5),170,178: $\pi_{t+1}$ should be $f(\mathcal{D}_{t+1}|x)$, yes? The derivations all look correct if so. 181: what is $\delta$? $Z_{(t+1)}$ at 205 and $\hat{Z}_{t+1}$ are the same? More of a suggestion: could the method be extended to cases where the target objects grow at each time step? For example, phylogenetic trees would grow as new species are observed. It seems like a natural application of flow nets since you can expand the graph by appending a new terminal layer.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

The set generation task is degenerate because the objective is additive. A model only needs to track the posterior on individual $x$s.

Summary

The authors provide a framework that allows training of GFlowNet models with streaming data by checkpointing a previously trained GFlowNet model and proposing the streaming balance condition. The problem as well as the approach is well motivated and theoretically sound and this work would be a valuable contribution to the research community. The authors conduct experiments on synthetic tasks and phylogenetic inference, and show improved scalability. Additionally, there is some theory that describes and bounds the approximation error based on errors on the checkpointed GFlowNet model and current estimation problem. Overall, I think the work is quite interesting but it would be nice to consider another experiment since currently most of the experiments are based on synthetic data.

Strengths

- The authors provide useful theoretical contributions that bound the approximation errors into two terms: error from solving the current distribution matching problem and error coming from suboptimal solution to the previous distribution matching problem. - The problem of training sampling methods that can incorporate streaming data is quite important and relevant to the field, especially for Bayesian Inference methodologies. - The authors do provide a good set of experiments, and show the benefits of their proposed approach over the standard training of GFlowNets.

Weaknesses

- I found Definition 2 to be quite confusing to read and it is still not clear to me how the equivalence holds. Could the authors clarify what the KL divergence is between, why is it valid and how does it circumvent the problem of parameterizing a partition function? - The original setup proposed by the authors follows from a trivial extension of the GFlowNets framework which begs to question the novelty of the work. - It is unclear how their formulation solves the problem of permutation invariance / iid treatment of observations. In particular, given two sets $D_1$ and $D_2$, we know that the posterior distribution $p(x | D_1, D_2)$ can be equivalently written as $$ p(x | D_1, D_2) = \frac{p(D_1, D_2 | x) p(x)}{p(D_1, D_2)} = \frac{p(D_1 | x) p(D_2 | x) p(x)}{p(D_1, D_2)} = \frac{p(D_1 | x) p(x|D_2)}{p(D_1 | D_2)} = \frac{p(D_2 | x) p(x|D_1)}{p(D_2 | D_1)} $$ Then, in particular, how is it maintained that given any split of the data $D$ into $D_1$ and $D_2$, the final GFN trained leads to the same solution. If the GFN learns the right solution, then it will follow simply, but it would be nice if the model could satisfy this permutation invariance inherently. Could the authors provide some ablations and control experiments to show to what extent is this constraint satisfied over $t$? - The authors should add an experiment on discovering causal graphs conditioned on observational samples, which would strengthen their paper considerably. - The authors should consider the following related work, which are either connected to GFlowNets, sampling from an unnormalized density, or modeling Bayesian posterior inference: Cranmer, Kyle, Johann Brehmer, and Gilles Louppe. "The frontier of simulation-based inference." Proceedings of the National Academy of Sciences 117.48 (2020): 30055-30062. Mittal, Sarthak, et al. "Exploring Exchangeable Dataset Amortization for Bayesian Posterior Inference." ICML 2023 Workshop on Structured Probabilistic Inference {\&} Generative Modeling. 2023. Zhang, Qinsheng, and Yongxin Chen. "Path integral sampler: a stochastic control approach for sampling." arXiv preprint arXiv:2111.15141 (2021). Richter, Lorenz, et al. "VarGrad: a low-variance gradient estimator for variational inference." Advances in Neural Information Processing Systems 33 (2020): 13481-13492. Sendera, Marcin, et al. "On diffusion models for amortized inference: Benchmarking and improving stochastic control and sampling." arXiv preprint arXiv:2402.05098 (2024). Berner, Julius, Lorenz Richter, and Karen Ullrich. "An optimal control perspective on diffusion-based generative modeling." arXiv preprint arXiv:2211.01364 (2022). Akhound-Sadegh, Tara, et al. "Iterated denoising energy matching for sampling from Boltzmann densities." arXiv preprint arXiv:2402.06121 (2024). Richter, Lorenz, Julius Berner, and Guan-Horng Liu. "Improved sampling via learned diffusions." arXiv preprint arXiv:2307.01198 (2023).

Questions

- It is not clear what do the authors mean by $\mathbf{A}$ in their notation at the start of Section 2. - The authors seem to have a typo in the equation under Line 142, essentially the right hand equation should have $p_F^{t+1}$. - Can the authors provide a visualization of the Set generation task? If I understand correctly, the authors first sample $d$ points randomly and the corresponding $f_i(j)$ randomly as well where $i = 1, ..., K+1$ and $j = 1, ..., d$. Then they define $R_i$ over each set, and a tempered version of $R_i$ with temperature $\alpha$. Is that it?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have adequately talked about the limitations and impacts of their work.

Summary

This paper introduces Streaming Bayes GFlowNets (SB-GFlowNets), which enables streaming Bayesian inference over discrete parameter spaces, relying on the expressive power of GFlowNets as amortized samplers over discrete compositional objects. The process is akin to Bayesian streaming, where the posterior updates with each new data stream and the posterior from the batch before acts as a prior for the current batch. The authors use a GFlowNet to fit the initial posterior and update it training only on new data stream batches. To do so, the authors propose two different solutions: enforcing a streaming balance condition at each streaming step, or relying on direct divergence-based updates (letting go of the estimation of the normalizing constant). For the latter, the authors show how to obtain low-variance estimators for the gradient of the proposed KL divergence loss. Finally, the authors provide a theoretical performance analysis of their two proposed variants along with a diverse set of experiments to showcase the performance in each case, and show the gain in training time in the streaming setting compared to retraining for all the data at once at each step.

Strengths

- The authors provide a complete study of all the important components within their work (experiments, runtime comparison to full retraining and practical performance enhancement etc). - The idea is interesting across many applications where data streams are processed continuously. - The paper is overall well written and easy to follow.

Weaknesses

- I'm not sure where equation (9) comes from. Equating \$\mathcal{L}\_{SB}(\tau)\$ to 0 would result in \\[ p\_F^{(t+1)}(\tau) = \frac{Z\_t}{Z\_{t+1}} \cdot \frac{p\_B^{(t+1)}(\tau \vert x)}{p\_B^{(t)}(\tau \vert x)} \cdot p\_F^{(t)}(\tau) \cdot f(\mathcal{D}\_{t+1} \vert x) \\] Even by assuming \$(p\_F^t, p\_B^t, Z\_t)\$ satisfy TB \$(\mathcal{L}\_{TB} \rightarrow 0)\$, I'm not sure how that would directly yield equation (10). The proofs for proposition 2 also seem to start from (9). Can you elaborate on that? - To what extent is the assumption that new data is independent of past data under the posterior predictive? This question is especially important in cases where you try to sample from the marginal likelihood over some object (for instance for structure learning), where you lose that independence under the posterior given only the structure (and not parameters for example). I suspect this might be the reason for performance degradation across streaming steps in the Bayesian phylogenetic inference case, given that Felsenstein's algorithm assigns a marginal likelihood. **Minor Comments** - Line 30, $\{y_i\}_{i=1}^2$ should be defined/introduced better. - Line 134, space after comma. - In line 142, in the second TB condition, it should be $p_F^{(t+1)}(\tau)$ instead of $p_F^{(t)}(\tau)$. - $p_F^{(t+1)}$ should be evaluated at $\tau$ in equation (10).

Questions

See weaknesses.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

- The authors adequately addressed limitations throughout the paper (for instance posterior error propagation etc) and through a separate section in the conclusion.

Summary

This paper introduces Streaming Bayes GFlowNets, a method for performing approximate Bayesian inference over discrete parameter spaces in streaming data settings. SB-GFlowNets allow efficient updating of posterior approximations as new data arrives. The method reduces training time compared to retraining GFlowNets from scratch and maintains comparable accuracy. In conclusion, SB-GFlowNets enable streaming variational inference for discrete parameters, opening up new applications of Bayesian methods to large-scale streaming data problems.

Strengths

- SB-GFlowNets significantly reduce training time compared to retraining GFlowNets from scratch for each new data batch. This is particularly valuable for large-scale streaming data problems. - The paper provides a theoretical analysis of how errors propagate through posterior updates - The method is demonstrated to work well on various tasks (set generation, Bayesian linear preference learning, and online phylogenetic inference)

Weaknesses

- As mentioned in the paper's limitations section, inappropriate approximations to earlier posteriors may propagate through time, potentially leading to increasingly inaccurate models. This is a form of catastrophic forgetting - Due to the error accumulation issue, there may sometimes need to retrain the model from an earlier checkpoint or using the full posterior - As the paper mentioned, for very sparse target distributions, the performance can degrade, particularly when using the KL-based training scheme

Questions

- Do you think the accumulative error issue can be solved by cache previous gradient (borrowing tricks from continual learning)? - It would be great if you can just compare with some continual learning method as baselines. - In sparse target distributions, if not using KL-based training scheme, what are the possible schemes?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

The authors have adequately addressed the limitations

Summary

The paper introduces a method for Bayesian inference in streaming data scenarios, particularly for discrete parameter spaces, called Streaming Bayes GFlowNets (SB-GFlowNets). The main contributions are: * Bayesian Streaming Inference: SB-GFlowNets enable the continuous updating of posterior distributions as new data arrives, without needing to recompute from scratch. * Addressing Intractability: The method addresses the challenge of approximating intractable posteriors in discrete state spaces, which is a limitation of existing variational inference (VI) techniques. * GFlowNet Utilization: SB-GFlowNets leverage GFlowNets, a class of amortized samplers, to approximate the initial posterior and update it incrementally with new data. * Case Studies: The effectiveness of SB-GFlowNets is demonstrated in linear preference learning and phylogenetic inference, showing its ability to sample from unnormalized posteriors efficiently in a streaming context. * Performance: The method is significantly faster compared to repeatedly training a GFlowNet from scratch for the full posterior.

Strengths

Two technically novel training algorithms for streaming inference on GFlowNets based on modified balance condition and VI with a control variate. Presentation was reasonably clear and the experimental results were convincing. Removes limitation of GFLowNets, an increasingly important family of models, to only batched training.

Weaknesses

A weakness of the paper is that it may not sufficiently highlight the significance of streaming inference compared to batch inference. This may be apparent to Bayesian practitioners but not so to a general ML audience. Are there more compelling applications of streaming inference?

Questions

Sentence in lines 58-59 trails off.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes, limitations have been addressed.

Reviewer Response

Thanks to the authors for providing a detailed and compelling rebuttal. While most of my concerns have been addressed, I would additionally like to point out that the work would be considerably stronger if the authors could consider a Bayesian linear regression model, with increasing number of observations seen on the X-axis and the KL divergence with the true posterior, which is available in closed form, on the Y-axis. However, I do understand that the discussion period is coming to an end soon and the authors might not have enough time to run this experiment for the discussion. Since all my other concerns have been addressed, I have updated the score accordingly.

Response to Rebuttal

I thank the authors for their thoughtful response to my questions and suggestions. That combined with the discussion with other reviewers below, I am glad to increase my score.

Thanks for your service

Dear reviewers and conference chairs, We are thankful for your service and commitment to the peer-reviewing process. Also, we greatly appreciate the reviewers’ input, which helped to significantly strengthen and clarify our work. All experiments and discussions will be included in the revised manuscript. Best regards, The authors.