Hierarchically branched diffusion models leverage dataset structure for class-conditional generation

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

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.

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.).

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.

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.

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).

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.

Thank you for clarifying the last few questions!

Fine-tuning a linear model

Linear (i.e. traditional, label-guided) diffusion models are generally single-task models which are not trained with multiple output heads (Ho et. al. 2021). Instead, class identity is specified to the model as a label embedding which is fed in as an input. It is certainly true that one of the major reasons why branched diffusion models are much easier to extend to new data is the ability to fine-tune a small number of parameters for only a small part of diffusion time, while guaranteeing that no previously learned tasks/classes are affected. Indeed, one of the contributions of branched diffusion models—in addition to recognizing that there are natural branch points along the diffusion process—is to turn the neural network into a multi-task network which optimally separates diffusion time into distinct branches and tasks to accomplish this separation of parameters.

In contrast, training a traditional diffusion model to be multi-task (with one output task per data class) is generally not a method which is used today. A lot of the diffusion process between classes (particularly at later times) would be inefficiently learned by multiple tasks, and so a single-task model with label embeddings is the ubiquitously relied-upon approach for class-conditional diffusion today. A branched model solves this problem of inefficient learning by leveraging a hierarchy of class similarities to define branches of shared diffusion. That said, we agree that this method of multi-task diffusion in an otherwise linear model would be far more amenable to continual learning than the ubiquitous single-task linear diffusion models, although it would still be less suitable than branched diffusion models due to the need to train on the entire diffusion timeline for each class (as you suggested).

Additionally, training a multi-task linear diffusion model would not allow for the many other benefits of branched diffusion models, such as transmutation, interpretability of branch points, or efficient generation.

FID on feature values

That's right! 2-Wasserstein distance is perhaps a clearer term to use here for similarity comparisons on general vectors. We will update the language in our manuscript accordingly.

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.

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

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.

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.

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.

We sincerely appreciate the time taken to take our new results and response into consideration and reflection! Here, we will address these remaining questions.

W2

Our selection of $\epsilon$ was done by first examining the similarity between pairs of classes throughout the diffusion timeline from $t = 0$ to $t = T$. Given our rough goal of having branch points not too close to $0$ or $T$, we considered the similarity between pairs of classes at $t = \frac{T}{2}$. We then simply selected our value of $\epsilon$ to be the average similarity over pairs of identical classes (i.e. the average similarity between random partitions of data from the same class), minus the average similarity over pairs of different classes, rounded to one significant digit for simplicity. This difference measures the expected gap in similarity between identical versus distinct classes at the midpoint of forward diffusion. We found that this procedure was sufficient to ensure that the branch points were not all too close to $t = 0$ or $t = T$.

Of course, our results in Supplementary Figure S6 also show that the branch points (and downstream performance) are fairly robust to the selection of $\epsilon$ anyways. We agree that this procedural detail could be very helpful in justifying our selection of $\epsilon$ (even if performance is not particularly sensitive to this hyperparameter), and we will update our manuscript accordingly.

Q3

The decision to start with a model of fewer classes was motivated by two reasons: 1) sheer convenience of faster training time and data loading (particularly during development and testing); and 2) we needed to hold out never-before-seen classes to introduce for our continual-learning experiments (i.e. if we had trained a model on all available classes, then we would not have new classes to add for our experiments).

Our expectation that there is no increase in difficulty when starting with a larger model versus a smaller model is firmly rooted in the following observations: 1) when we add a new data class to a branched diffusion model, the weights in the existing architecture are frozen so that catastrophic forgetting is impossible; and 2) learning to generate a new class effectively reduces to fine-tuning a standard diffusion model, where the early layers of the neural network are pre-trained, and the late diffusion times are transferred from similar data. Note that this process of freezing early layers and only training on early diffusion times (for a single class) is also what allows continual learning in a branched model to be much faster than in a linear model.

In order to support our expectation that there is no increase in difficulty for larger models, we note that our experiments in Supplementary Figure S7 showed that the continual-learning benefit of branched diffusion models remained equally effective and performative as more classes were introduced and the size of the model increased. Furthermore, we also showed continual-learning results on branched diffusion models trained on our much larger dataset of single-cell RNA-seq (Figure 2).

Finally, to directly answer the final question, training a 10-class model does not take a prohibitively large amount of time, and this was never a limitation for us. For example, to generate Supplementary Figures S5—S6, we trained many branched diffusion models (on the full 10-class dataset) for 150 epochs, which was the same as the linear model.