Paper: Hierarchically branched diffusion models leverage dataset structure for class-conditional generation (ICLR-2024)

Reviewer comment (Experiments):
The paper is well-written, but some claims could be improved for clarification. For example, in Section 4 Page 6, it is unclear how the model performs with versus without fine-tuning the upstream branches that also diffuse over the newly added class in the continual learning setting without certain empirical results as support. Another example is that I found the observation in Section 5 Page 6 that letters with a larger feature value tended to transmute to letters with a larger feature value is hard to interpret from Figure 3b with only scatterplots of some feature values given.

Author response:
**Fine-tuning upstream branches (W1)**

In our experiments of continual learning in branched diffusion models, we found that fine-tuning upstream branches did not help (or hurt) the generative performance of downstream tasks. This is a very expected result, as our branch points are explicitly defined to be the point in diffusion time after which classes diffuse nearly identically. Hence, fine-tuning these upstream branches is not expected to significantly change the model’s performance. We will better clarify this point in the revised version.

**Transmutation of letters (W1)**

In Figure 3b, we showed the scatterplots for just two examples of features, but the histograms below (still in panel 3b) show that we found positive (often strongly positive) correlations for *all* features in the dataset. Of course, if analogous features were not being retained in the process of transmutation, we would expect these histograms to be tightly centered around 0.

Reviewer comment (Theory):
The efficiency of branched models over standard linear models when sampling multiple classes is clear from Table S9 when the number of classes is relatively small. Could the authors provide some insights when the number of classes grows very large with potential comparison regarding complexity analysis?

Author response:
**Efficiency of branched models with more classes (W2)**

The speedup of branched-diffusion models is only indirectly influenced by the number of classes, and instead is directly determined by the *hierarchy* and how much diffusion time can be shared. The amount of shared diffusion time, in turn, is a function of how similar different classes are (and the diffusion process). For datasets with a lot of very similar classes, we are able to share much more diffusion time, and therefore will enjoy more speedup. However, datasets with very distinct classes will have very late branch points, leading to lower efficiency. It is also worth mentioning that datasets with only very distinct classes are not particularly well-suited to the method of branched diffusion which we propose in this work. Instead, the main contribution of our work is intended for scientific applications, whose datasets are typically characterized by a large amount of internal structure (e.g. cell types, chemical classes, structural folds, etc.).

Reviewer comment (Presentation):
In Figure 4, the interpretation at immediate branch points between two classes is shown and aligned with intuition. I am curious about how the visualization is for the branch point from more upstream such as the branch point between class 0 and the immediate branch point between class 4 and 9.

Author response:
**Interpretation of branch points between less-related classes (W3)**

Branch points between less-related classes can certainly still be interpreted (and still yield valid interpolations/intermediates). However, between very unrelated classes, the interpretations are naturally less meaningful, as only very high-level (and relatively uninformative) features will be shared between these classes.

To demonstrate this, we updated our manuscript to show the hybrid resulting from interpreting the branch point between more dissimilar classes (Supplementary Figure S9). For example, on our tabular letters dataset, we quantitatively show that the branch-point hybrids between dissimilar classes still exhibit feature values which are interpolated between the two original classes, and often act as intermediates which exhibit properties of both. We also visually show the branch-point hybrids between dissimilar digits of our MNIST dataset. Our resulting hybrids show the common features between dissimilar digits, which are fairly high-level, including: 1) the digit is in the center of the image; and 2) there are areas which are relatively empty (coinciding with the “holes” of how many people draw digits like 6s and 9s). This is a direct consequence of the lower similarity between two unrelated data classes rather than a limitation of branched diffusion models.

Reviewer comment (Evaluation):
Why the FID between true and generated cells are the same for the branched and linear model in Figure S3 c)?

Author response:
**FID of RNA-seq dataset (Q1)**

For the RNA-seq dataset, below we reproduce the FID values for each class. Due to the high variation in FID values between the classes (likely because some cell types are inherently harder to generate than others), differences in these values were not obviously rendered in Supplementary Figure S3.

|  | Branched | Label-guided |
| --- | --- | --- |
| CD16+ NK | 377.254 | 376.300 |
| Cl. Mono. | 163.470 | 166.414 |
| Late Eryth. | 150.605 | 153.134 |
| Macroph. | 151.118 | 156.442 |
| Megakar. | 64.968 | 65.506 |
| Mem. B | 1604.477 | 1610.171 |
| NK | 84.447 | 85.517 |
| Plasmabl. | 317.073 | 317.584 |
| Tem/Eff. H. T | 174.465 | 175.692 |

Reviewer comment (Experiments):
Could the authors give some insights on why only one baseline of the label-guided (linear) models (Ho et al., 2021) is used for comparison?

Author response:
**Choice of baseline (Q2)**

Our proposed method of branched diffusion is a fundamentally novel way to perform class-conditional generation using diffusion. We thus sought to compare the benefits of branched diffusion to the current state-of-the-art methods of class-conditional diffusion. In the literature, there are effectively two ways to perform class-conditional diffusion: classifier-guided diffusion and classifier-free diffusion (the latter of which we called label-guided diffusion). Classifier-guided diffusion requires training external classifiers on all noise levels, and is not commonly used today for this reason. Label-guided diffusion, on the other hand, is universally relied upon today for conditional diffusion, and state-of-the-art results in diffusion arise practically entirely from this method. Hence, we opted to perform comparisons relative to this technique (the only state-of-the-art method used today).

Reviewer comment (Evaluation):
Some of the experiments are not, at the moment, particularly convincing of the usefulness of the effects that they are showcasing. For the analogy-based generation with the RNA-seq data set, some marker genes were indeed changed, but do we have any other way to evalute the success of these generations? Could we formalize a clear objective on what does the conditional generation aim to do in the first place here? A similar issue exists for the molecule data set: Indeed regenerating does allow to generate cycled molecules from acyclic ones, but it is not clear what if any properties of the original molecule are retained this way. Just looking at the results, it seems possible that the generated molecules are a random mixture of the desired property and some atoms and bonds from the original molecules.

Author response:
Thank you for taking the time to review and provide valuable feedback! We have answered all the questions and comments in the following.

**Formalizing and quantifying the goals of transmutation (W1, Q1)**

In transmutation (i.e. analogy-based conditional generation), we start with an object $x_0$ of class $c_0$, and transmute it to another object $x_1$ of class $c_1$ (usually $c_0 \neq c_1$). The goal of transmutation is two-fold: 1) **efficacy**: $x_0$ should have features which define it to be in class $c_0$, and $x_1$ should have features which define it to be in class $c_1$; and 2) **analogy**: features which do not distinguish/define an object’s class to be in $c_0$ or $c_1$ should be retained. This makes analogy-based conditional generation substantially different from standard conditional generation, as we are interested in sampling from $q(c_1, x_0\in c_0)$ instead of $q(c_1)$. In addition to several qualitative results, we systematically quantified both objectives in our experiments, demonstrating both efficacy and analogy in multiple real-world datasets.

For the RNA-seq example, we showed **efficacy** by quantifying the marker genes before and after transmutation. The class of a cell (i.e. cell type) is a complex concept, and marker genes are by far the most widely accepted method for determining cell type. As such, we showed key marker genes which demonstrate that a cell was successfully converted from one type to another. To demonstrate **analogy**, we quantified the correlation of non-marker genes, namely genes responsible for COVID-related inflammation (which do not define the cell type). These inflammation genes were explicitly identified by Lee et. al. (2020) (the publication which is the source of the RNA-seq data) as being key upregulated genes in COVID-infected cells (regardless of cell type). The correlation here is quantified over a random sample of cells (each point in the correlation is a single cell). A high correlation indicates that if the starting cell $x_0$ had a high expression of the gene, then the transmuted cell $x_1$ *also* had high expression of that gene (and *vice versa*). For these inflammation genes, this is precisely the desired property, as it indicates that the model is transmuting COVID-infected cells of one type into COVID-infected cells of another type, and healthy cells from one type into healthy cells from another type (i.e. the non-cell-type-defining features are retained).

For the molecular example, we quantified **efficacy** simply by counting the number of molecules which satisfied the goal structural property. We quantified **analogy** by computing the similarity of functional groups (which, in turn, endow the chemical properties of the molecule) before and after transmutation. Overall, in this experiment, the objective of transmutation is to generate molecules that satisfy the new class (e.g. “cyclic” or “halogenated”), while also resembling the initial molecules. Therefore, our results showing that the generated molecules satisfy the desired target property and are some “mixture” of similar chemical motifs from the original molecule, are certainly an indication that transmutation is working properly. For any starting molecule, it is desirable to see a distribution of similar atoms and functional groups in the generated molecules (i.e. analogy), but still being valid molecules which satisfy the target property (i.e. efficacy).

Reviewer comment (Experiments):
Continuing on the analogy-based generation, I feel that it would be appropriate to do an ablation where we use a regular diffusion model, noise out the data partially, and regenerate with the changed label. Would this work equally well, or differently somehow?

Author response:
**Attempting transmutation on a linear diffusion model (W2)**

It is certainly possible to perform what was suggested here in a linear (i.e. regular) diffusion model. However, there are no explicitly defined branch points in a linear model, so it is not clear to what point forward diffusion should be performed before reverse diffusing to a different class (i.e. what “turn-back point” to use). As such, we would be forced to select a *post hoc* turn-back point which the model was not trained with. This leads to several pitfalls:

- Choosing a turn-back point which is too early will hurt efficacy, as the model will not have enough diffusion time to generate an object of the target class.
- Choosing a turn-back point which is too late will hurt analogy, as the model will lose many features which are non-class-defining.
- A turn-back point which retains both efficacy and analogy for one individual example from the dataset may not be sufficient for another example, because a linear diffusion model is not trained with any concept of branch points, and so there is no imposition of this structure between classes as a whole.

One of our core contributions is precisely the definition of formal branch points, which are the optimal time points to perform transmutation. Branched diffusion models are explicitly trained with these branch points, and so the model learns to generate shared characteristics in upstream branches, and class-specific characteristics in class-specific branches, for all examples.

Following the suggestion, to better show the advantages of branched diffusion models for transmutation compared to standard (linear) diffusion models, we have added a new ablation study (Supplementary Figure S8). Our results show that in the absence of explicitly defined branch points, attempting to perform transmutation in a linear model is significantly more difficult. Imposing a *post hoc* turn-back point that is too early causes the model to fail at efficacy: the target class is not generated at all. Imposing a turn-back point which is too late causes the model to suffer in analogy: non-class-specific features are no longer correlated before and after transmutation. Additionally, choosing certain turn-back points can cause some objects to be transmuted efficaciously, and others to fail to generate the target class entirely. These issues all severely complicate the attempt to perform transmutation in a linear model. Furthermore, even if we had prior knowledge of what a branch point should be, the transmuted results from the linear model at the “correct” branch point still are lower quality than in the branched model, likely because of crosstalk between the classes, which is not controlled for at all in the training of the linear model.

Thus, branched diffusion models are much more well-structured for tasks like transmutation (and of course, branched models are also uniquely more easily extendable in a continual-learning setting, their branch points are interpretable, and generation of multiple classes is significantly more efficient).

Reviewer comment (Evaluation):
The points about interpretability are also interesting, but it remains a bit unclear what could be the use of these average branching points.

Author response:
**Benefits of interpreting branch points (W3)**

In our experiments, we have shown how hybrid intermediates naturally capture shared characteristics between classes (Figure 4). In this paper, we focused on validating this hypothesis, as it demonstrates how the proposed model is learning to hierarchically generate objects. We will extend our discussion to better highlight the potential benefits of branch points on interpretability, which include:

- Discovery. Branch points can help identify shared higher-order features across two (or more) classes. While class similarities could also be discovered through other techniques, using branch points enables a model-based definition of these shared features, as they depend not only on the underlying data distribution but also on the trained model, including the defined diffusion process and the entire dataset on which the model was trained.
- Distribution-based interpretation of shared features. Branch-point intermediates are *distributions* of objects, where each individual feature has its own distribution of values given any two classes. This allows us to sample from the intermediate distributions, thereby not only allowing us to glean insight into shared features and concepts, but also assess their significance/confidence.
- Improved model understanding. Branch points may be used to diagnose potential issues in the generative model, providing a quantitative/qualitative way of checking generation intermediates, in addition to generated samples. On a related note, during the development of this method, we partially relied on interpreting branch points to verify the integrity of our discovered hierarchies.

 We plan to further investigate applications of directly interpreting branch-point intermediates in future work.

Reviewer comment (Evaluation):
What does the high correlation of expression between genes before and after transmutation mean here? That those particular genes often did not change? Is this the property that we want?

Author response:
For the RNA-seq example, we showed **efficacy** by quantifying the marker genes before and after transmutation. The class of a cell (i.e. cell type) is a complex concept, and marker genes are by far the most widely accepted method for determining cell type. As such, we showed key marker genes which demonstrate that a cell was successfully converted from one type to another. To demonstrate **analogy**, we quantified the correlation of non-marker genes, namely genes responsible for COVID-related inflammation (which do not define the cell type). These inflammation genes were explicitly identified by Lee et. al. (2020) (the publication which is the source of the RNA-seq data) as being key upregulated genes in COVID-infected cells (regardless of cell type). The correlation here is quantified over a random sample of cells (each point in the correlation is a single cell). A high correlation indicates that if the starting cell $x_0$ had a high expression of the gene, then the transmuted cell $x_1$ *also* had high expression of that gene (and *vice versa*). For these inflammation genes, this is precisely the desired property, as it indicates that the model is transmuting COVID-infected cells of one type into COVID-infected cells of another type, and healthy cells from one type into healthy cells from another type (i.e. the non-cell-type-defining features are retained).

Reviewer comment (Novelty):
Since adding uncorrelated noise probably does not result in clean hierarchical class structures in all cases (maybe in more complex image data sets, as pointed out in the paper), do you think it would be possible to induce such structure, e.g., by diffusing in some specifically designed latent spaces or otherwise designing the diffusion process itself to encourage it?

Author response:
**Inducing hierarchies by defining alternative diffusion processes (Q2)**

We found that even with a simple Gaussian diffusion process, we were able to recover meaningful hierarchies in our datasets. For example, although there is often some variation in the branching structures which can arise, our models remain robust to this variation, and these branching structures generally make intuitive sense (Supplementary Figure S5).

Defining alternative diffusion processes or leveraging uniquely structured latent spaces are certainly interesting directions which we hope to explore in future work. Depending on the scientific goal in mind, a unique diffusion process may be defined which targets a particular outcome for tasks like transmutation. For example, in the generation of molecules, a discrete diffusion process which operates on junction-tree representations may further encourage whole substructures and motifs to be retained in the transmutation process. On the other hand, a more fine-grained diffusion process which readily breaks/makes bonds might be better for a goal like scaffold hopping. Leveraging a specifically designed latent space (e.g. one which preserves short distances between molecules which have similar solubility) may result in a branched model which is better tailored to solubility-related tasks, as more diffusion time can be shared.

Reviewer comment (Experiments):
I wonder if the example where catastrophic forgetting is avoided in MNIST is possible to extend to multiple steps, to a slightly more realistic continual learning scenario?

Author response:
**Multi-step continual learning (Q3)**

The continual-learning experiments we showed can certainly be extended to multiple steps. The biggest benefit of branched diffusion models for continual learning is that when we introduce a new data class, the existing model remains completely frozen (i.e. the model’s generation of previous classes remains absolutely identical). As such, we have no doubts that since we achieved good generative results after introducing one class, our findings easily extend to the case where we introduce additional classes.

To further show how our method performs in larger continual-learning experiments, we present novel results performing multi-step continual-learning experiments where we extended our MNIST models to three never-before-seen classes one by one: 7s, 1s, then 2s (a mixture of classes which are similar to or distinct from existing classes in the model) (Supplementary Figure S7).

As expected, our results here show the exact same trend as before: branched models are easily extended with an efficient fine-tuning step, where the performance of previous classes is untouched upon the introduction of a new data class. In label-guided (linear) models, fine-tuning on new data causes catastrophic forgetting. It is also worth noting that due to the reliance of label-guided models on class embeddings, typical neural-network architectures are unable to accommodate never-before-seen labels unless they are explicitly initialized to allow for extra class labels; in our experiments on continual learning, we needed to separately train a label-guided model on fewer classes, but with sufficient extra capacity (”embedding slots”) to accommodate future classes. In contrast, branched models have no such limitation, highlighting an additional advantage of branched models in continual learning.

Reviewer comment (Experiments):
A more thorough experimental validation for continual learning and analogy-based generation would be in place, so that the reader would have clear takeaways. For the analogy-based generation, some kind of formalization of what are we targeting with the conditional generation, would also help with showcasing the potential significance.

Author response:
Thank you for the feedback and for clarifying remaining concerns!

Firstly, we would like to note that we do conduct fairly extensive experiments over many different data types and datasets, spanning from **images** to **graphs** to **tabular data**. Two of our datasets were also **large, real-world scientific datasets**, and altogether we showed **global, quantitative** results for each of the benefits brought by our proposed method of branched diffusion (continual learning, transmutation, interpretability, and efficiency). As such, we believe that these experiments (in addition to the additional experiments we have shown throughout the rebuttal period) are sufficient to demonstrate the advantages and disadvantages of our method relative to the current state of the art.

**Transmutation**

Unfortunately, typical probability notation is not very well-suited to capture the process of transmutation. To be more clear, we defined $q(c_{1}, x_{0}\in c_{0})$ above to be the conditional distribution of objects for class $c_{1}$, but *also* conditioned on an object $x_{0}\in c_{0}$. If we were to draw an object from $x_{1}\sim q(c_{1}, x_{0}\in c_{0})$, we would like both efficacy and analogy to hold (i.e. $x_{1}$ belongs to class $c_{1}$ and features that do not define class identity are shared between $x_{1}$ and $x_{0}$).

Although we agree that it would be nice to have more formalization on the process of transmutation, it is unfortunately not possible to mathematically formalize both goals of efficacy and analogy in a straightforward way. In particular, it remains generally infeasible to clearly distinguish and define the class-defining features which should be modified in efficacious transmutation, and the instance-specific features which should be kept analogous. Although class-defining features and instance-specific features certainly exist, they are oftentimes *latent*, *high-level*, and *entangled* with each other. If the class-defining and instance-specific features were disentangled and clearly defined, then a general mathematical formalization would be more accessible, but unfortunately, the realization of these features are highly dependent on the specific application, including the dataset, the definition of classes, and the meaning of the underlying data. Unless we have a more simplistic dataset or we have adequate domain knowledge to perform this disentanglement and latent discovery, it remains infeasible to formally define transmutation at a feature level.

For this reason, we showed **globally quantitative** results on transmutation for *many datasets of different data types* (Figure 3). For example, we showed transmutation on a tabular dataset, where we could examine feature analogy by inspecting individual features. We also showed globally quantitative results demonstrating both efficacy and analogy on two real-world large scientific datasets: RNA-seq and drug-like molecules. For RNA-seq, to show efficacy and analogy we relied on knowledge of the domain: namely the presence of marker genes which define cell type and analogous COVID-related inflammation genes which were first discovered in Lee et. al. (2020). For the molecules, our quantification of efficacy and analogy relied on high-order structural properties and knowledge of functional groups.

Overall, however, it remains relatively infeasible to clearly mathematically formalize the low- and high-order features which we expect to be retained for analogy but modified for efficacy, particularly when we are attempting to remain general to most scientific applications. Importantly, note that transmutation can be thought of as a generalization of neural style transfer (NST) (i.e. generating images where the subject of the image is given the style of another image). Note that in the seminal works which define NST [1, 2], *there is also no mathematical formalization of the goals of NST*, due to the inability to define ahead of time the features which we expect to be modified or retained.

We are also happy to hear that our novel ablation study shown in Supplementary Figure S8 was helpful. Importantly, we would like to clarify that *transmutation is a novel technique that we presented, which was **not** a part of traditional linear diffusion models*. Indeed, our main contribution is the discovery of branch points and applying them to a diffusion model, so that we can perform transmutation in the first place. As our results show, attempting to force a conventional linear model to perform transmutation is unnatural and difficult, and transmutation as a technique in linear diffusion models also did not exist prior to our paper. As such, beyond this ablation study showing why branch points are important for transmutation, we do not find it particularly meaningful to perform such an extensive evaluation (as we have done in our other experiments) against such an unnatural and difficult technique which understandably did not exist before.

Reviewer comment (Experiments):
Some of the experiments are not particularly compelling and serve more as examples of the potential of the technique rather than providing convincing evidence in support of the claims of the paper. In the image domain, analysis is limited to MNIST and the letter-recognition dataset; this leaves open the question as to how well the approach scales to other (less-structured) types of images and more challenging image classes, where the class hierarchy may not be so clear.

Author response:
**Application to scientific domains and datasets outside of images (W1)**

The main contribution of our work is intended for scientific applications, whose datasets are typically characterized by a large amount of internal structure (e.g. cell types, chemical classes, structural folds, etc.). As such, although we demonstrated our method’s soundness on well-known image benchmarks such as MNIST, we also focused much of our work on real-world large scientific datasets such as single-cell RNA-seq and drug-like molecules, where we showed the promise of branched diffusion to perform real scientific discovery by recovering known biology.

Notably, our proposed method is very much dependent on datasets with internal structure. Datasets which have do not have inherent structure (or only very weak structure) between classes are not suitable for branched diffusion, whose core hierarchical backbone reflects the intrinsic similarities between classes in the dataset. We will modify the language in our manuscript to better clarify this intended use case of our method.

Reviewer comment (Experiments):
The branching point definition involves a threshold. There does not appear to be a concrete recipe for the selection of this threshold. Based on the paper it seems to be left to the practitioner to determine when branching points are “too close” to 0 or T. There is no investigation of how the selection of this threshold impacts performance.

Author response:
**Defining thresholds and robustness to thresholds (W2)**

We found that in general, our selection of $\epsilon$ was not too difficult, as there is a range of possible values which lead to branch points that are fairly close to the time range of $t \in [0.1, 0.5]$. Additionally, we also found that variation in the branch points (resulting from randomness in the branch-point discovery procedure) still yielded similar results (Supplementary Figure S5).

To further demonstrate this, we present novel results on the robustness of our method to different values of $\epsilon$ in the updated version of our manuscript (Supplementary Figure S6). On the MNIST dataset, we swept through possible values of $\epsilon$ ranging from $10^{-5}$ to $10^{-1}$, uniformly distributed in logarithmic space (note that the $\epsilon$ we used in our main results was $5 \times 10^{-3}$). We computed branch points for each value of $\epsilon$, and trained a branched diffusion model on each. The two largest values of $\epsilon$ yielded class-specific branches of length 0, and therefore were discarded from this analysis. Supplementary Figure S6 shows that the branch points identified from our various values of $\epsilon$ are very similar to the distribution of branch points identified in Supplementary Figure S5 (variation due to randomness in sampling and in diffusion). This suggests that our branch-point discovery algorithm is rather robust to various values of $\epsilon$, provided obviously inappropriate values (such as those which yield length-0 terminal branches) are not considered. Furthermore, models trained with these branch points still yielded generative performance which was largely robust to $\epsilon$.

Reviewer comment (Experiments):
The major weakness of the paper is that most of the claims in the paper are supported by relatively limited experimentation. For example, the paper claims that the method achieves similar or better generative performance as the state-of-the-art, but does not clearly preface this with the clarification that the outperformance is observed only for two simple character-based image datasets and the similar performance is only established for one other dataset. The paper would be considerably more convincing if there were experiments on more challenging image datasets.

Author response:
**Application to scientific domains and datasets outside of images (W1)**

The main contribution of our work is intended for scientific applications, whose datasets are typically characterized by a large amount of internal structure (e.g. cell types, chemical classes, structural folds, etc.). As such, although we demonstrated our method’s soundness on well-known image benchmarks such as MNIST, we also focused much of our work on real-world large scientific datasets such as single-cell RNA-seq and drug-like molecules, where we showed the promise of branched diffusion to perform real scientific discovery by recovering known biology.

Notably, our proposed method is very much dependent on datasets with internal structure. Datasets which have do not have inherent structure (or only very weak structure) between classes are not suitable for branched diffusion, whose core hierarchical backbone reflects the intrinsic similarities between classes in the dataset. We will modify the language in our manuscript to better clarify this intended use case of our method.

Reviewer comment (Experiments):
“In general, the branched diffusion models achieved similar or better generative performance compared to the current state-of-the-art label-guided strategy. In many cases, the branched models outperformed the label-guided models, likely due to the multi-tasking architecture which can help limit inappropriate crosstalk between distinct classes.” – these sentences seem to be strong claims when the experiments are conducted on three datasets (two of which are similar). There seems to be no discernible outperformance for the single-cell RNA-seq data. The outperformance is really only for two character-based image datasets, so “In many cases” seems to be a stretch. Considerably more extensive experiments on a variety of datasets are required to support the general claim made in the paper. Alternatively it could be restricted to “For experiments performed on two character-based datasets and an RNA-seq dataset, ….” Can the authors clarify whether they consider the current experiments to be sufficient to demonstrate similar or better generative performance?

Author response:
**Comparison of generative performance and use of FID (W3, Q1, Q2)**

In our manuscript, we quantified the generative performance of branched vs linear (i.e. traditional) diffusion models over three datasets. For MNIST, we saw better performance in the branched model for 10/10 classes. For letters, our branched model outperformed in 24/26 classes. For RNA-seq, our branched model outperformed in 8/9 classes. Given these results, we consider it reasonable to claim “in general, the branched diffusion models achieved similar or better generative performance compared to the current state-of-the-art label-guided strategy”, and that “in many cases, the branched models outperformed the label-guided models” (quoted from Section 3.1). Note that our claims here are certainly within the context of our experiments, and we are not attempting to claim that branched diffusion models always offer improved generative performance over their linear counterparts in all situations. For improved clarity, we will modify these sentences as follows: “*In our experiments*, the branched diffusion models generally achieved similar or better generative performance…”

Specifically for the RNA-seq dataset, below we reproduce the FID values for each class. Due to the high variation in FID values between the classes (likely because some cell types are inherently harder to generate than others), differences in these values were not obviously rendered in Supplementary Figure S3.

|  | Branched | Label-guided |
| --- | --- | --- |
| CD16+ NK | 377.254 | 376.300 |
| Cl. Mono. | 163.470 | 166.414 |
| Late Eryth. | 150.605 | 153.134 |
| Macroph. | 151.118 | 156.442 |
| Megakar. | 64.968 | 65.506 |
| Mem. B | 1604.477 | 1610.171 |
| NK | 84.447 | 85.517 |
| Plasmabl. | 317.073 | 317.584 |
| Tem/Eff. H. T | 174.465 | 175.692 |

Finally, we agree that every metric for performance (including FID) has its advantages and disadvantages. In our analyses, we believe that FID is sufficient to capture the global trends in generative performance that we wish to understand. In addition to the analysis of FID, we also verified the generative performance of our branched diffusion models using other methods (Supplementary Figures S1—S2).

Importantly, however, we emphasize that the main contributions of our work lie *beyond* generative performance, and we simply wish to ensure that branched diffusion models are not *significantly worse* in performance compared to their linear counterparts. We believe that our analyses in this regard are sufficient to demonstrate this claim.

Reviewer comment (Experiments):
Why is the new class experiment limited to training on 3 classes? Is this to make the task easier? The new class experiment for the single-cell RNA-seq dataset seems to be similarly limited (just starting with two classes and adding a third). Is there a reason that the more obvious experiment of removing just one class and adding it back is avoided? What happens if the “1” class is already included (i.e. something that is much closer to the introduced task)?

Author response:
**Additional continual-learning experiments with larger models (Q3)**

The core benefit of branched diffusion models for continual learning is that the entire existing model can be frozen, and only a single new branch is added for a new class. Thus, in the exploration of this continual-learning benefit, we expect no difference in difficulty if we were to start with a larger existing model or a smaller one. Specifically for the continual-learning experiments, we opted to use smaller models simply for efficiency in training/loading. Regarding class extension when a highly similar class is already present (e.g. the model has already been trained to generate 1s and we are now introducing 7s for the first time), we would expect branched models to perform even better (or at least, certainly not worse) in these experiments, while we expect label-guided models to suffer from the same issues.

Additionally, we opted to show the result of adding a new data class which has never before been seen by the model (instead of removing an existing class and adding it back), as we believe this better simulates how these scientific models may be deployed in practice (i.e. new data coming in from experiments on new cell types or systems).

To further show how our method performs in larger continual-learning experiments, we present novel results performing multi-step continual-learning where we extended our MNIST models to three never-before-seen classes one by one: 7s, 1s, then 2s (a mixture of classes which are similar to or distinct from existing classes in the model) (Supplementary Figure S7).

As expected, our new results highlight the same trend as before: branched models are easily extended with an efficient fine-tuning step, where the performance of previous classes is untouched upon the introduction of a new data class. In label-guided (linear) models, fine-tuning on new data causes catastrophic forgetting. It is also worth noting that due to the reliance of label-guided models on class embeddings, typical neural-network architectures are unable to accommodate never-before-seen labels unless they are explicitly initialized to allow for extra class labels; in our experiments on continual learning, we needed to separately train a label-guided model on fewer classes, but with sufficient extra capacity (”embedding slots”) to accommodate future classes. In contrast, branched models have no such limitation, highlighting an additional advantage of branched models in continual learning.

Reviewer comment (Experiments):
“Additionally, this limitation on images may be avoided by diffusing in latent space.” – is there evidence for this claim?

Author response:
**Diffusion in latent space for images (Q5)**

To clarify, our intended meaning for this sentence was that we *expect* this issue for images to be avoided by diffusing in latent space (i.e. “may” here being an expression of possibility rather than assured potential), although we leave exploration of this for future work, particularly because our contribution is meant to be for general scientific data (including RNA-seq, molecular graphs, etc.) instead of being focused on images. We will clarify and reword this to be clear in the revised revision.

That said, our RNA-seq diffusion models were trained in latent space, so we already have some evidence that supports the use of latent-space representations for branched-diffusion-model training.

Reviewer comment (Experiments):
A concern arises regarding the scalability of the proposed method as the number of classes increases.

Author response:
Thank you for taking the time to review and provide valuable feedback!

To answer your question about datasets with more classes, we note that one of the datasets we presented had 26 classes, which we consider fairly typical in magnitude (or even larger in magnitude) compared to most scientific applications.

For example, training a model with cell types for classes would require only 10 - 15 classes for most tissue systems. For drug-discovery applications where the classes are target receptors, most families of receptors only have 5 - 10 members under investigation at any given point.

Reviewer comment (Experiments):
The experiments conducted appear to be limited to datasets with a small number of classes.

Author response:
Thank you for taking the time to review and provide valuable feedback!

To answer your question about datasets with more classes, we note that one of the datasets we presented had 26 classes, which we consider fairly typical in magnitude (or even larger in magnitude) compared to most scientific applications.

For example, training a model with cell types for classes would require only 10 - 15 classes for most tissue systems. For drug-discovery applications where the classes are target receptors, most families of receptors only have 5 - 10 members under investigation at any given point.

Reviewer comment (Experiments):
It would be beneficial if the authors could present results for datasets with a larger number of classes.

Author response:
Thank you for taking the time to review and provide valuable feedback!

To answer your question about datasets with more classes, we note that one of the datasets we presented had 26 classes, which we consider fairly typical in magnitude (or even larger in magnitude) compared to most scientific applications.

For example, training a model with cell types for classes would require only 10 - 15 classes for most tissue systems. For drug-discovery applications where the classes are target receptors, most families of receptors only have 5 - 10 members under investigation at any given point.