Paper: Federated contrastive GFlowNets (ICLR-2024)

Reviewer comment (Novelty):
It should be noted that an alternative approach to the problem in the Bayesian posterior setting is to train a single GFlowNet with a *stochastic* reward, as done in [Deleu et al., UAI 2022] and [Deleu et al., NeurIPS 2023].

Author response:
**A:** Thanks for pointing this out. We are unsure about the relevance of the two references to stochastic rewards — both focus on causal structure learning. If they were not a typo, could you elaborate on that?

Nonetheless, we could use Zhang et al.'s (ICML 2023) scheme for training GFlowNets in environments with stochastic rewards drawn from the distribution

$$
R(x) = R_{c}(x), \, c \sim q(\cdot | \{1, \dots, K\}), % \{R_{1}(x), \dots, R_{K}(x)\}),  .  
$$

in which $q$ is a categorical distribution over the clients $\{1, \dots, K\}$. By assigning a weight $q(c)^{-1}$ to the $c$th client's reward, $R_{c}(x)^{q(c)^{-1}}$, Proposition 1 of Zhang et al. ensures that the trained GFlowNet would sample an object $x$ with probability proportional to

$\exp \underset{c \sim q}{\mathbb{E}} \left[\log R\_{c}(x)^{q(c)^{-1}}\right]  = \underset{1 \le k \le K}{\prod} {R}_{{k}}(x) $

The downside is that this requires many communication steps between clients and server. This is precisely the bottleneck we are trying to avoid with FC-GFlowNets, imposing one *single communication step* between the clients and the server. We wrote a short version of this discussion in related works (Appendix D).

Reviewer comment (Experiments):
There are still comparisons to be made in the experiments, such as using different training objectives for the client models (TB or CB).

Author response:
**A:** Thanks for the suggestion. In principle, FC-GFlowNets are agnostic to how the local models were trained as long as they provide us with forward and backward policies. We have now included Table 2 (below) in Appendix C, comparing the performance of FC-GFlowNets when clients are trained with TB vs. CB. The metrics are all very similar.

|                          		 | Grid World         		 |       		 | Multisets         		 |      		 | Sequence          		 |       		 |
|-----------------------------------|-------------------------|-------------------------|-------------------------|-------------------------|-------------------------|--------------------|
|                          		 | L₁ ↓          		 | Top-800 ↑      		 | L₁ ↓          		 | Top-800 ↑      		 | L₁ ↓          		 | Top-800 ↑      		 |
| FC-GFlowNet (CB)         		 | 0.038          		 | -6.355         		 | 0.130          		 | 27.422         		 | 0.005          		 | -1.535         		 |
|                          		 | (± 0.016)      		 | (± 0.000)      		 | (± 0.004)      		 | (± 0.000)      		 | (± 0.002)      		 | (± 0.000)      		 |
| FC-GFlowNets (TB)        		 | 0.039          		 | -6.355         		 | 0.131          		 | 27.422         		 | 0.006          		 | -1.535         		 |
|                          		 | (± 0.006)      		 | (± 0.000)      		 | (± 0.018)      		 | (± 0.000)      		 | (± 0.005)      		 | (± 0.000)      		 |

Reviewer comment (Novelty):
The "contrastive loss" is claimed as original, but in fact it is not. I suggest that the authors revise the discussion on this in section 3.3 and in the claimed contributions.

Author response:
**A:**  We once again thank you for noticing the relationship between the contrastive loss (CL) and the log-partition variance loss (VL) of Zhang et al. (ICLR 2023) — which we did not notice.  In fact, CL equals twice VL *in expectation*. Nonetheless, as you pointed out, CL and VL use different estimators for the variance up to some positive multiplicative constant. Additionally, our balance condition (CB) is novel per se and essential to deriving the federated loss (Corollary 1), which is the basis of our method.  We have updated the introduction, section 3.2 (which introduces the CB), and related works (Appendix D) to acknowledge this work.

Reviewer comment (Experiments):
There is uncertainty about whether differences between CL and TB (Figure 6) are only due to insufficiently high learning rate on logZ for TB, or differently tuned learning rates for the policies in both algorithms.

Author response:
**A:** To verify whether our gains stemmed from a poor choice of learning rate, we have run additional experiments (in Figure 8), sliding it from $0.001$ to $10$. In all cases, we saw CL converging to better optima than TB. That being said, while there may be some learning rate for log Z leading to better convergence, finding it may be difficult. We have updated Figure 6 to include the results from (former) Figure 7.

Reviewer comment (Theory):
The proofs of all the results in section 3.3 have one-line reductions to existing results (the TB training theorem and the TB gradient analysis theorem). The current proofs are quite obfuscated.

Author response:
**A:** We tried to make the proofs as clear and as self-contained as possible. In particular, as our arguments in Section 3.2 consist primarily of standard and easy-to-follow calculations, we believe that articulating a comprehensive proof in the Appendix is beneficial for enhancing the results' clarity. It is unclear whether shortening the proofs would make them easier to follow.

Reviewer comment (Theory):
The bound in Theorem 2 is nice but may not be very useful in practice. If one of the models is missing a mode, which is quite possible with reverse-KL objectives like those used here, $\alpha_n$ will be very small, and the final bound is not useful.

Author response:
**A:** We did not intend Theorem 2 to be used for diagnosing pathological implementations of FC-GFlowNet. Theorem 2 only aims to illustrate that the global distribution may diverge considerably from the target distribution when a client's policies are inadequately trained. Particularly, if one of the clients fails to detect some modes, as you suggested, then the global distribution will also miss these modes. This effect, similarly observed in e.g. De Souza et al. (AISTATS 2022), is accurately represented by our upper bound on the Jeffrey's divergence.

We agree that developing good diagnostics for federated GFlownets is an important topic that we would like to pursue in future works. However, the design of tight divergence bounds and effective diagnostic methods is currently an unaddressed issue in the GFlowNets literature, even for centralized settings.

Reviewer comment (Novelty):
I believe should give more concrete example on why learning a GFlowNet samples data proportional to product of reward is important in federated setting. From reading the paper and experiments, I am not clear about the significance of the problem studied by this paper. In fact, it seems the problem tackled by this paper is an special case while it is not evident that this special case is important in federated setting.

Author response:
**A:** The most common problem in federated learning is learning some point estimate $\theta^{\star}$ that minimizes the sum of loss functions  $\ell_{1}, \ldots, \ell_K$ for each client $k=1\ldots K$. Typically, the loss is equivalent to the negative log-likelihood of $\theta$ on some local dataset $\mathcal{D}_k$. In the Bayesian realm, the corresponding problem is sampling from a posterior distribution $p(\theta | \mathcal{D})$ proportional to the products of the local likelihoods and some prior over theta. Each of these terms in the product can be seen as a local reward. Then, the product of local rewards is proportional to the posterior. That being said, FC-GFlowNets are a natural solution for Bayesian federated learning cases in which $\theta$ is a discrete-valued random variable.  More generally, FC-GFlowNets work for any reward that factorizes as products. Thus, it can be used for, e.g., multi-objective active learning. Importantly, there well-established interest of the ML community in Bayesian federated learning — see., e.g., Appendix D for more details. To clarify the significance of our problem, we have also extended paragraph 3 (Introduction) in the revised manuscript. Please see our first answer to reviewer wWFS for further discussion about the motivation.

Reviewer comment (Novelty):
It seems that the main contribution of this paper is its problem setting where learning a GFlowNet model that samples data proportional to product of rewards can be done in a federated fashion. However, given the problem, learning the global model aggregating local GFlowNet seems to be straightforward. Therefore, I think I believe the contribution of the paper in federated learning is not high.

Author response:
**A:** We are not aware of any simpler approach to sampling from the product of GFlownets. The naive strategy would be to train a GFlowNet to sample from the product of terminal distributions $\hat{R} = \prod_{k=1}^{K} p_k^\intercal(x)$ of the clients $k=1\ldots K$. Note, however, that $p_k^\intercal(x)$ requires summing over all terminal paths leading to $x$ — which is infeasible in practice. Furthermore, using a Monte Carlo estimate of $p_k^\intercal(x)$ plus some standard GFlowNet loss results in biased gradients, which are undesirable if we want to converge with stochastic gradient descent.

Furthermore, sampling from products of arbitrary distributions over discrete objects is not a trivial task. This is the reason, for instance, why most of the literature on parallel MCMC focuses solely on continuous cases. Importantly, while there is a good body of literature on Bayesian federated learning, to the best of our knowledge, FC-GFlowNets is the first method capable of handling discrete cases.

Reviewer comment (Novelty):
However, the motivation for introducing GFlowNets into Federated Learning (FL) appears to be lacking. Despite the existing research gap, what specific challenges or benefits does FL encounter that require the integration of GFlowNets?

Author response:
**A:** Thanks for the opportunity to motivate Federated GFlowNets further. As mentioned in the Introduction, we foresee at least two impactful applications of FC-GFlowNets (our proposal): multi-objective active learning and Federated Bayesian Inference.

In multi-objective active learning, we want to obtain samples with high values across a pool of utility functions. For instance, in drug discovery, a key challenge is finding molecules that satisfy multiple constraints (e.g., affinity, solubility, safety). In a setting where different experts work independently to obtain property-specific models using GFlowNets, this challenge can be cast as sampling from the product of these local GFlowNets. This corresponds to the scalarization approach for multi-objective optimization, but we are interested in sampling to ensure diversity instead of maximizing --- which was also one of the primary motivations behind GFlowNets.

Federated Bayesian Inference/Learning entails sampling from a posterior distribution $p(\theta|\mathcal{D})$ where $\mathcal{D}$ is a dataset partitioned into a number of clients $k=1\ldots K$, each with a subset $\mathcal{D}_k$. The common setting is that the posterior factorizes as

$p(\theta|\mathcal{D}) \propto p(\theta) \prod_{k} p(\mathcal{D}_k | \theta)$, where $p(\theta)$ is a prior and $p(\mathcal{D}_k | \theta)$ is the likelihood contribution of client $k$. Note that the likelihood is typically different for each $k$; more generally, they don't even need to share a functional form. There are a series of methods to solve this problem in a single communication round when $\theta$ is a real-valued random variable. FC-GFlowNets (our proposal) solves this problem for the discrete case — and is the first work to do so.

Thanks to your comment, we have extended paragraph 3 (Introduction) in the revised manuscript. We also highlighted that our related work section discusses the Federated Bayes setting in detail (Appendix D).

Reviewer comment (Writing):
Conversely, is there a scenario in which GFlowNets necessitate Federated Learning (FL)? In my view, the explanation provided in the third paragraph is somewhat lacking. It would be helpful to provide more practical examples to illustrate this.

Author response:
**A:** Thanks for the opportunity to motivate Federated GFlowNets further. As mentioned in the Introduction, we foresee at least two impactful applications of FC-GFlowNets (our proposal): multi-objective active learning and Federated Bayesian Inference.

In multi-objective active learning, we want to obtain samples with high values across a pool of utility functions. For instance, in drug discovery, a key challenge is finding molecules that satisfy multiple constraints (e.g., affinity, solubility, safety). In a setting where different experts work independently to obtain property-specific models using GFlowNets, this challenge can be cast as sampling from the product of these local GFlowNets. This corresponds to the scalarization approach for multi-objective optimization, but we are interested in sampling to ensure diversity instead of maximizing --- which was also one of the primary motivations behind GFlowNets.

Federated Bayesian Inference/Learning entails sampling from a posterior distribution $p(\theta|\mathcal{D})$ where $\mathcal{D}$ is a dataset partitioned into a number of clients $k=1\ldots K$, each with a subset $\mathcal{D}_k$. The common setting is that the posterior factorizes as

$p(\theta|\mathcal{D}) \propto p(\theta) \prod_{k} p(\mathcal{D}_k | \theta)$, where $p(\theta)$ is a prior and $p(\mathcal{D}_k | \theta)$ is the likelihood contribution of client $k$. Note that the likelihood is typically different for each $k$; more generally, they don't even need to share a functional form. There are a series of methods to solve this problem in a single communication round when $\theta$ is a real-valued random variable. FC-GFlowNets (our proposal) solves this problem for the discrete case — and is the first work to do so.

Thanks to your comment, we have extended paragraph 3 (Introduction) in the revised manuscript. We also highlighted that our related work section discusses the Federated Bayes setting in detail (Appendix D).

Reviewer comment (Experiments):
The experimental setup regarding the clients is not sufficiently clear. Are all clients identical, or does this setup involve non-iid (non-independent and identically distributed) problems?

Author response:
**A:** In our experiments, clients are never identical, i.e., there is no pair of local rewards $(R_i, R_j)$ that are equal. For instance, in the phylogeny experiments, each client's reward depends on different parts of the nucleotide sequences. In this case, no two clients have access to the same part of the sequence. For the grid world, each client has different reward beacons. Finally, the multiset and design-of-sequences experiment, clients value items differently. We have clarified that no clients have identical rewards in the experiment section.

Reviewer comment (Presentation):
It would enhance clarity if the authors could include a figure or pseudocode illustrating the algorithm.

Author response:
**A:** Thanks for the suggestion. We have included a pseudo-code for FC-GFlowNets in the Appendix.