Hierarchically branched diffusion models leverage dataset structure for class-conditional generation
Authors: Alex M Tseng
Max W Shen
Tommaso Biancalani
Gabriele Scalia
Venue: Accepted by TMLR ()

Abstract: Diffusion models have attained state-of-the-art performance in generating realistic objects, including when conditioning generation on class labels. Current class-conditional diffusion models, however, implicitly model the diffusion process on all classes in a flat fashion, ignoring any known relationships between classes. Class-labeled datasets, including those common in scientific domains, are rife with internal structure. To take advantage of this structure, we propose hierarchically branched diffusion models as a novel framework for class-conditional generation. Branched diffusion models explicitly leverage the inherent relationships between distinct classes in the dataset to learn the underlying diffusion process in a hierarchical manner. We highlight several advantages of branched diffusion models over the current state-of-the-art methods for class-conditional diffusion. Firstly, they can be easily extended to novel classes in a continual-learning setting at scale. Secondly, they enable more sophisticated forms of conditional generation, such as analogy-based conditional generation (i.e. transmutation). Finally, they offer a novel interpretability into the class-conditional generation process. We extensively evaluate branched diffusion models on several benchmark and large real-world scientific datasets, spanning different data modalities (images, tabular data, and graphs). We particularly highlight the advantages of branched diffusion models on a single-cell RNA-seq dataset, where our branched model leverages the intrinsic hierarchical structure between human cell types.

Published in Transactions on Machine Learning Research (08/2024)
Hierarchically branched diffusion models leverage dataset structure for class-conditional generation
Alex M. Tseng
tseng. alen@gene com meen.man@gene.co com biancall@yene. com scalia gwbrieabbbrrreee.e. com
Max Shen
Tommaso Biancalani
Gabriele Scalia
Biology Research | Development
Cenentech
Reviewed on OpenReview: https: http::///pperevieet net/ permfo id- sGTfagRbes
Abstract
Diffusion models have allained state-of-the-arl performance in generating realistic objects, including when conditioning generation on class labels. Current class-conditional diffusion models, however, implicitly model the diffusion process on all classes in a flat fashion, ignoring any known relationships between classes. Class-labeled datasets, including those common in scientific domains, are rife with internal structure. To take advantage of this structure, we propose hierarchically branched diffusion models as novel framework for class-conditional generation. Branched diffusion models explicitly leverage the inherent relationships between distinct classes in the dataset to learn the underlying diffusson process in a hierarchical manner. We highlight several advantages of branched diffusion models over the current state-of-the-arl methods for class-conditional diffusion. Firstly, they can be easily extended to novel classes in a contimual-learning setting at scale. Secondly, they enable more sophisticated forms of conditional generation, such as analogy-based conditional generation (i.c. transmutation). Finally, they offer a novel interpretability into the class-conditioaal generation process. We extensively evaluate branched diffusion models several benchmark and large real-world scientific datasets, spanning different data modalities (images, tabular data, and graphs). We particularly highlight the advantages of branched diffusion models on a single-cell RNA-seq dataset, where our branched model leverages the intrinsic hierarchical structure between human cell types.
1 Introduction
Diffusion models have gained major popularity as method for generating data from complex data distributions -Sohl-Dickstein et al. 2015 Ho et al.] 2020 Song et al. 2021). Furthermore, they have also been successful in
performing conditional generation, where we wish to sample some object x conditioned on a label y. Recent
works in conditional diffusion have arguably received the most popular attention (Song al. 2021 Dhariwal
& Nichol 2021 Rombach et al., 2022 Ho et al. 2022), and have rapidly bocome a staple in generative AI.
limited in their treatment of labsled datasets. Conventional diffusion models learn the diffusion process flatly for each class, disregarding any known relationships structure between them. In reality, class-labeled datascts in many domains, such as those characterizing scientific applications, have an inherent structure between classes which can be thought of as hierarchical. For example, human cell types are organized hierarchically by nature: keratinocytes are very distinct from neurons, but the latter subdivide into excitatory and inhibitory neurons. Additionally, even when a datasct. has no pre-defined hierarchical label set (c.g. an explicit ontology), hierarchical structure can often be found, because some subsets of classes are invariably more similar than others.
Current diffusion models, however,

---

Published in Transactions on Machine Learning Research (08/2024)
In order to leverage this intrinsic structure within these datasets, we propose restructuring diffusion models to be hierorchicalyy branched, where the branching structure reflects the inherent relationships bet ween classes in the dataset (the underlying diffusion process remains unchanged). By modeling the diffusion process in a hicrarchical fashion, branched diffusion models enjoy several advantages which make them much more suitable in many applications, particularly scientific seltings. We apply branched diffusion to four different datasels, including two large real-world scientific datasets, spanning different data modalities (images, tabular data, and graphs) and showcase the following advantages over the current state-of-the-art method for conditional diffusion:
Branched models offer a novel way to perform class-conditional gencration via diffusion by organizing labels hierarchically.
They can be easily extended to generate new, never-before-seen data classes in a continual-earrning setting.
They can be used to perform more sophisticated forms of conditional generation, such analogy-based conditional generation (or "transmutation")
Diffusing across their branched structure offers interpretability into the relationship between classes from a generative perspective, such as elucidating shared high-level features.
2
Related Work
Most diffusion models today are defined by a pair of forward and reverse stochastic differential equations (SDEs). The forward equation injects random noise over continuous time to transform the initial data distribution pa( into a tractable prior distribution π(x) from to the time horizon t T:
dx (a, + g())w,,
(1)
where w a standard Wiener process (i.e. Brownian motion). fxx t) and g(t) are the drift and diffusion coefficients, respectively. For the variance-preserving SDE, for example, f(x, - f and g(t) /Bc for some noise schedule Be (yielding a prior (2x) which is an isotropic Gaussian).
In order to sample an object from po(x), we first tractably sample from π(x), and then follow the associated reverse SDE to recover a sample from Px()::
dx = -((,,t)-g(('(((r) () + g(()dw';
(2)
where w' is a standard Wiener process in the reverse direction. s(x, () = Va log(p<(x:) is the Stein score of x at diffusion time t. A neural network is trained to predict s0 t) x s(x, t), made possible by defining f(r, t) and (t) such that the true Stein score is tractably defined in closed form (such as with the variance-preserving SDE).
generation, the diffuusin model needs learn the conditional distribution of data for each class of the dataset. The current stato-of-the-art method for class-conditional diffusion proposed in Ho et al. (2021), termed "classifier-fre conditional generation". In this method, the reverse- auxiliary input, which guides the generation of objects
In order to perform
diffusion neural network is given the class label to specific classes:
dr = [f(,,t) - so f, c)] dt + g(t)dw',
(3)
where s0 1, c) is the trained neural network which approxinates the Stein score at x and time / for the conditional distribution of class c.
This method of conditional generation has achioved state-of-tho-art performance sample quality (Rombach et al. 2022; Ho al. 2022). and in contrast to the previous method of classifier guidance (Song al. 2021) )
it be applied to both continuous- and discrete-time diffusion models.
2

---

Published in Transactions on Machine Learning Research (08/2024)
Note that the term "hierarchical diffusion" is somewhat overloaded, as there exist other works which use this term, but describe vastly different methods than the one proposed here. For example, Qiang et al.] (2023) describes applying distinct diffusion processes iteratively to generate coarse-grained features before fine-grained features. Lu et al. (2023) proposes latent-space diffusion where latent embeddings of various granularities and sizes are passed to the diffusion denoising model. In contrast, the "hierarchy" in our work refers to a hierarchy of similarity between classes in the original dataset.
3 Hierarchically branched diffusion models
Suppose our dataset consists of a set of classes C. For example, let us consider MNIST handwritten digits, where C {0,1.... 9}. We wish to leverage the fact that some classes are inherently more similar than others (e.g. the 4s and 9s in MNIST are visually more similar to each other than they are to 0s). As noise is progressively added to data, there is some point in diffusion time at which any two samples from two different classes are so noisy that their original class effectively cannot be determined; we call this point in time a branch point. A branch point is a property of two classes (and the forward diffusion process), and and-importantly- more similar the two classos are, the carlier the branch point will be.
These branch points underpin the main distinction between branched diffusion model and a traditional linear one. We define the branch point between two classes as the earliest diffusion time point when objects of the two classes sufficiently similar in distribution, such that reverse diffusion after this point could be predicted by the same neural-network model (without specifying class identity). More formally, defined criterion for measuring when the distribution of two noisy classes is sufficiently similar (Equation D based off of the well-established metric of energy distance. Mathematical justification and more details can be found in Appendix [A] Note that branch-point discovery is performed once for the dataset at the beginning, and this proposed procedure's runtime is orders-of- -fgnn-it smaller than the time taken to train the model. Alternatively, since branch points reflect the inherent similarity between classes of the dataset, they may also be entirely defined by domain knowledge (e.g- an ontology describing known similarities between cell types, chemical classes of Brug--ike molecules).
Together, the branch points between all classes in C' naturally encode a hierarchy of class similarities (Figure [.) This hierarchy separates diffusion time from a single linear track into a branched structure, where each branch represents the diffusion of a subset of classes, and a subset of diffusson times. For |C| classes, there are 2|C1- - 1 branches. Each branch b, = (s,, 44, C.) is defined by a particular diffusion time interval [s,,t, t.) (where 0 S s; < < T) and a subset of classes C. S C (where C, * 0). The branches are constrained such that every class and time (c,t) C x [0, T) can be assigned to exactly one branch d; such that C2 and E (s). The branches form a rooted tree starting from t = T to t O. Late branches (large t) are shared across many different classes, as these classes diffuse nearly identically at later times. Early branches (small t) are unique to smaller subsets of classes. The earliest branches are responsible for gencrating only a single class.
Additionally, as opposed to a conventional ("linear"; or "non-hierarccicall") diffusion model which learns to reverse diffuse all classes and times using a single-task neural network, a branched diffusion model is implemented a multi-task neural network. Specifically, each output task predicts reverse diffusion for single branch (e.g. in an SDE-based diffusion framework (Song et al. 2021). each prediction head learns the Stein score for a specific branch) (Figure [p). The multi-task architecture allows the model to learn the reverse-diffusion process distinctly for each branch, while the shared parameters allow the network to learn shared representations across tasks without an explosion in model complexity:
dr ((())))))(()>s((m) (), p) dt + g(t)dw',
(4)
where be. is the branch index corresponding to the class c of object x, and t is the diffusion time. Note that the forward- and reuerse-diffusion processes are identical traditional diffusion models (Equations D However, in contrast to traditional linear models where the neural network sa(x, t, e) takes in the class as an input to learn class-conditional distributions (Equation the neural network for a branched model explicilly
3

---

Published in Transactions on Machine Learning Research (08/2024)
a)
b)
Branch
Reverse diffusiion (scoore) prediction
Branch bj
Branch li
4
Branch by
Shvred laperss
Branch b
9
Branch lo.
Qrli)]
golz)
Forward diffusion
(2tz
4
4
(are L. m)
( 4)
[
Traditional closs coss-caantiooall diffusion All classes and flussoo times predicted NN
by Each clees has NN which pretictss
Hlerarchikan chiccllyy branched diffuston off limes of di difierent classers mesimallly hierardhy, each prodicts classes diffluiion
diffusiinn times
Mathematically dentical diffusion processes
Figure Schematic of a branched diffusion model. a) Some classes are more similar than others from a diffusion perspective: after adding sufficient noise, similar classes become indistinguishable from each other after branch points. Together, the branch points (purple dots) between all classes define a hierarchical structure. Although there is still only single underlying diffusion process, the hierarchy separates diffusion time into branches, where each branch represents diffusion for a subset classes and a subset diffusion times. An example of one diffusion intermediate is highlighted in blue; this example is MNIST digit that is 4 9, and is branched diffuusio model realized as a multi-task neural network (NN) that predicts reverse diffuusin (one output task for each branch). The prediction path for the blue-highlighted MNIST digit in panel a) is also in blue. c) We show a progression of methods from traditional linear diffusion models to hierarchically branched models. Left: A traditional linear conditiodil diffusion model trains single NN for all classes and diffusion times. Both class label and time are NN imputs. Middle: Consider modification, where each class has its own NN, each of which is trained reverse diffuse the full diffusion timeline for only associated class. generate an object of specific class, take reverse-diffusion predictions from the associated NN. Right: At late diffusion times, different classes diffuse very similarly. A branched diffusion model shares as much diffusion time as possible, thereby leading to a hierarchy of times, where each branch is learned by output task of multi-task NN. To generate specific class, take reverse-diffusion predictions from the appropriate subset of branches which cover the full diffusion timeline from = T to f
intermediate diffusion time. b) A
learns the conditional distribution of each branch as separate output task (each branch has an associated output task with index be.c).
Training a branched diffusion model follows nearly the same procedure with standard linear model, except for each input, we only perform gradient descent on the associated branch, which corresponds to a model output task) (see Algorithm D. To sample an object of class c, we perform reverse diffusion starting from time T and follow the appropriate branch down (see Algorithm 2
To better understand branched diffusion models and how they are different from traditional linear models, consider the following progression of methods (Figure [p): 1) A traditional coas--onniiionall diffusion model, where all classes are treated in flat fashion. A single neural network learns reverse diffusion for all classes and diffusion times; both the time and class label are inputs to the model (Equation 2) Instead of a single neural network, Irain a separate network for each data class (the diffusion process is still the same for each class). In contrast with method 1), we reverse diffuse a particular class by choosing the appropriate network. This would allow for benefits like class extension in continual lcarning (Soction 5). but the cost. of inefficient parameterization and training time. Furthermore, fully separating diffusion timelines for each class into distinct neural networks ignores inherent class similarities, and result, this approach does not.
4

---

Published in Transactions on Machine Learning Research (08/2024)
Algorithm 1 Training branched diffusion model
Input: training set ((x!5) c( branches } repeat
Sample (Ia, c) from training data {(x((H)
Sample t ~ (0,
Forward diffuse zz ~ (2)2z))
Find branch b (si, Ci) s.t. 8, Si c C; Gradient descent P(o.. (Te, t)] (on output task i) until convergence
Algorithm 2 Sampling branched diffusion model
Input: class c, trained pa, branches {b,}
Sample 2 2T from (2x)
for t = T to 0 do
Find branch b = (s;, ti, Ci) s.t. 8, sistchy CE C; i + & [)] (take output task i)
end for Return &C
allow for benefits like transmutation (Section or interpretability of diffusion intermediates (Section 3) Owing to the inherent structure of the dataset, we can leverage the fact that noisy objects of different. classes indistinguishable after a certain diffusion time (i.e. branch points). Thus, we maximize sharing of diffusion time between classes via hierarchy, and train multi-task neural network where each task predicts reverse diffusion for single branch (Equation 4 We now generate a particular class by choosing the appropriate sel of branches (output tasks). This allows us to retain benefits like class extension, and gain benefits like transmutation and interpretability, with a much more efficient parameterization and reduced training time compared method 2).
Importantly, the underlying diffusion process and subsequent mathematical properties in branched diffussin model directly inherited from and identical to that of a conventional linear model. A branched model is characterized by the explicit definition of branch points which separale the responsibility of reverse diffusing different subsets of classes and times into separate branches, where each branch is predicted by different head of a multi-task neural network.
4 Class-conditional diffusion via branched diffusion models
Branched diffusion models are a completely novel way to perform class-conditional diffusion. Instead of relying external classifiers labels auxiliary neural-network inputs, branched diffusion model generates data of a specific class simply by reverse diffusing down the appropriate branches.
We demonstrate branched diffiusion models on several datasets of different data modalities: 1) MNIST handwritten-digit images (LeCun et al.);; 2) a tabular dataset of several features for the 26 English letters in various fonts (Frey Slate| 1991p: 9; 3) a real-world, large scientific dataset of single-cell RNA-seg, measuring the gene expression levels of many blood cell types in COVID-19 patients, influenza paticnts, and hcalthy donors [Lee et al. 2020); ; and 4) ZINC250K, large dataset of 250K real drug-like molecules (Irwin al., 2012 We trained continuous-time (i.e. SDE-based (Song et al. 2021)) ) branched diffusion models for all datasots. The branching structure was inferred by our branch-point discovery algorithm (Supplementary
Tables ST
S8 Appendix A.. We verified that our MNIST model generated high-quality digits (Supplementary
Figure $1
For our tabular-letter dataset, we followed the procedure in Kotelnikov et al. (2222 to verify that
the branched model generated realistic letters that are true to the training data (Supplementary Figure S2
We compared the generative performance of our branched diffusion models to the current state-of-tho-ar methods for conditional generation via diffuusio,, which are label-guided (linear) diffusion models (Ho et. al.
5

---

Published in Transactions on Machine Learning Research (08/2024)
Note that although Ho et al. 2021 called these "classifier-fres conditional diffusion models, we will refer to them as "label-guided" in this work, since branched diffusion models also allow for conditional generation without the use of any external classifier. We trained label-guided diffusion models on the same data using the analogous architecture. We computed the Fréchet inception distance (FID) for cach class, comparing branched diffusion models and their linear label-guided counterparts (Supplementary Figure S3 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 oulperformed the label-guided models, likely due to the multi-tasking architecture which can help limit inappropriate crosstalk between distinct classes. This establishes that branched diffusion models offer competitive performance in sample quality compared to the state-of-the-art methods for conditional diffusion.
2021
Although we will focus our later analyses on contimuous-time (i.e. Sasebased (Song et al. 2021 ) diffusion models, we also trained discrete-tůme branched model (i.e. based on DDPM (Ho et al. 2020 to generate MNIST classes (Supplementary Figure SSI.. This illustrates the ficxibility of branched diffusion models: as they are generally orthogonal to the underlying diffusion process, they perform equally well in both continuous- and discrete-timg diffusion settings.
5 Extending branched diffusion models to novel classes
The problem of incorporating new data into an existing model is a major challenge in the area of continual learning, as the emergence of new classes (which were not available during training) typically requires the whole model to be retrained (van de Ven & Tolias 2019). Conventional (linear) diffusion models are no exception, and there is a critical need to improve the extendability of these models in a continual-learning sctting. This requirement is typical of models trained on largo-scale, integrated scientific datasots, which grow steadily as data of new, never-before-seen classes is experimentally produced (Han al.) 2020: Almanzar et al. 2020 Lotfollahi et al. 2021). For example, large single-cell reference atlases comprising hundreds of millions of cells across organs, developmental stages, and conditions- as the Human Cell Atlas (Regev
et atal.
2017 are continuously updated as new research is published.
By separating the diffusion of different classes into distinct branches which are learned by a multi-task neural network, a branchod diffusion model casily accommodates the addition of new training data (e.g. from a recent experiment). Suppose a branched model has been trained on classes C, and now a never- rever-before-seee class e' has been introduced. Instead of retraining the model from scratch on C U (c') for the entire diffusion timeline, a branched model can be easily extended by introducing new branch while keeping the other branches the same (assuming d' is sufficiently similar to some existing class). For example, leveraging the intrinsic structure of cell types (Han et al. 2020), a branched diffiusion model can be fine-tuned on a new study-potentially including new cell types wwithout retraining the entire model. Formally, we extend an existing branched diffusion model by adding a new terminal branch (si,4,, C.) = (0, 4n (c()), where 4a is determined by the algorithm in Appendix [ The new neural network has parameters 0= (0,. fo.... O ) with shared parameters 0g and outputttask-speecffi parameters 0g (one for each branch). Let b be the branch index of the new terminal branch. Then we simply need to learn Oz. training only on c' for times & E [0, 4o]:
0% = argming, (Farclass())=dd (ecrs [C(o, 1, so(a, 1) 1))
(5)
To illustrate this extendability, we trained branched diffusion models on MNIST and on our large real-world RN dataset. For the MNIST experiment, we trained on three classes: Os, 4s, and 9s. We then introduced new class: 7s. To accommodate this new class, added single new branch to the diffusion model (Figure Db). We then fine-tuned only the newly added branch, freezing all shared parameters and parameters for other output tasks. That is, we only trained on 7s, and only on times E (8),4) for the newly added branch by- After fine-tuning, our branched model was capable of generating high-quality 7s without affecting the ability to generate other digits (Figure 2p).
In contrast, label-guided (linear) diffusion models cannot easily accommodate a new class. In our MNIST experiment, we trained linear model on Os, 4s, and 9s. After fine-tuning the linear model 78 the model
6

---

Published in Transactions on Machine Learning Research (08/2024)
a)
b)
Defore sdding
fne-umarg
O O
O
o
O 0
0 O
O 4
Branched
4
4
4 4
y
4 4 8
4 9
4
9
9
9
9
q
4
9 7
7
Betore adstng new
funing
ony
Suning on
Labal-gatded (inar))
0
0
0
0
O
7
7 0
7
A
0
O
O
4
4
4
4
4
7
7
11
4
{
4
4
9
9
7
7
7
7
7 7
9
9
7
7
7
7
7
7
7 7
7
7
C)
MNIST
Ealow After resiinree das
Single-coll RNA =seq
Lubal
150
20
100
Class
Class
Class
Class
Figure 2: Extending a branched model to new classes. a) Schematic of the addition of new digit class to an existing branched difffusion model MNIST. The introduction of the class is accomplished by adding singular new branch (purple dotted line). b) Examples of MNIST digits generated from a branched diffusion model (above) and label-guided (linear) diffusion model (below), before and after fine-tuning on the class. For the label-guided model, we also show examples of digits after fine-tuning on the whole dataset. c) On the MNIST dataset (left) and the single-cell RNA-seq dataset (right); we show the FID (i.e. generative performance) of each class, before and after fine-tuning on the new class. For the label-guided models, we also show the FIDs after fine-tuning on the whole dataset.
suffered from catastrophic forgetting: largely lost the ability generate the other digits (even though their labels were fed to the model generation time), and generated almost all 7s for any label (Figure 2 For the linear model to retain its ability to generate pro-cxisting digits, it must be retrained dataset, which is far more inefficient, especially when the number of classes is large. In our MNIST example, retraining the linear model on all data took approximately seven times longer than training the singular new branch on a branched model. Notably, even after retraining on all data, the lincar model still experienced inappropriate influence from the new task. We observed the same trend of extendability on the branched model (and inter-class interference on the linear model) in our experiments on the large real-world RNA RNA-seq dataset.
the entire
We then quantified this class-extension ability by computing the FID of branched and label-guided models before and after fine-tuning (Figure ) On both MNIST and the RNA-seq dataset, we found that the branched models achieved roughly the same FIDs on pre-existing classes after fine-tuning on the new data class. In contrast, fine-tuning the label-guided models on the new data class caused the FID of other classes significantly worsen. The label-guided models needed be trained on the entire dataset to recover the FIDs of pre-cxisting classes, although the FIDs were still generally worse than those of the branchod models.
Notably, in extension of branched models, were able to generate high-quality samples of the new class by only training a single branch. Although several upstream branches (i.e. at later diffusion times) also diffuse over the newly added class d' (along with pro-existing classes), we found that fine-tuning can be performed on only the new branch which diffuses solely only é, and the model still achieved high generative performance across all classes C' {c'). This a natural consequence of our explicit branch points, which are defined so that reverse diffusion at upstream branches is nearly identical between different classes. Note that even if fine-ture tuninu were done on all classes for all b, where d' E C,, the branched diffusion model is still more efficient to extend than linear model because most branches diffuse over only a small subset of classes in C.
7

---

Published in Transactions on Machine Learning Research (08/2024)
This highlights the advantage of branched diffusion models to accommodate new classes efficiently (i.e. with little fine-tuning) and cleanly (i.e. without affecting the generation of other classes) compared to the current. state-of-the-ar methods for conditional diffusson. Note that although we showed the addition of a brand new class c' & C, branched models can also easily accommodate new data of an existing class c E C' by fine-tuning only the appropriate branch(es).
6 Analogy-based conditional generation between classes
In diffusion model, we can traverse the diffusion process both forward and in reverse. In a branched diffusion model, this allows for unique ability to perform analogy-bosed conditional generation (or transmutation) between classes. That is, we start with an object of one class, and generate the analogous, corresponding object of a different class. Formally, consider the set of all branches {(s), t,, C.) Say we have an object xı of class C1. We wish to transmute this object into the analogous object of class c2. Let to be the first branch point (carliest in diffusion time) in which 01 and 02 are both in the same branch. Then, in order to perform transmutation, we first forward diffuse æ1 lo 2h ~ qt, (ar(²1). Then we draw an object from the conditional distribution Po(2|o2, where conditioning is both on class 0 and the noisy object ib (partially diffused from 2.). In summary:
4. == (siles E Ci} Tb 2t (x(m))
(6)
22 ~ 20)
In practice, sampling ve from ev (o(21(22, 26) is performed by reverse diffusing to generate an object of class C2 (Algorithm 2]. as if we started at time ts with object 2n ~ qt, (z|z1).
Conditional generation via transmutation is unique and novel way to harness branched diffusion models for more sophisticated generation tasks which go beyond what is possible with current diffusion models, which typically condition on a single class or property (Song et al. 2021 Ho et al. (021). In transmutation, we enable generation conditioned on both a class and a specific instance fwhich may be of another class). That is, conventional conditional generation samples from qu (o(=(o2). but transmutation samples from (0(e22, 21 E ct). This feature can support discovery in scientific settings. For example, given a model trained types C, with each cell type mcasured in certain conditions huy...,hms transmutation can answer the following question: "whal would be the expression of a specific cell zi of cell type and condition lR, if the cell type were сj instead?". A branched model is thus distinct from models which simply generate cells with cell type Cj and/or condition hg. For instance, to study how a novel cell type (such as a B-cell) reacts to a drug whose effects are known for another cell type (such as a T-cell), one can conditionaly generate a population of B-cells starting from a population of T-cells under the particular drug effect.
many cell
On our MNIST branched diffusion model, transmuted between 4s and 9s (Figure 3h). Intriguingly, the model learned to transmute based on the slantedness of digit. That is, slanted 4s tended to transmute to slanted 9s, and vice versa. To quantify analogous conditional generation between classes, we then transmuted between letters on our tabular branched diffusion model (Figure 3p). Transmuting between V and Y (and vice versa), we found that for every feature, there was a positive correlation of the feature values before versus after transmutation. That to say, letters with larger feature value tended to transmute to letters also with a larger feature value, even if the range of the feature is different between the two classes.
We then turned to our branched model trained on the large real-world RNA-seq dataset, and transmuted sample of CD16+ NK cells to classical monocytes, and vice versa. In both directions, transmutation successfully incrcasod critical marker genes of the target cell type, and zeroed the marker genes of the source cell type (e.g. when transmuting NK cells to monocyles, the expression of NK marker genes such as MS4A6A were zeroed, and the expression of monocyte marker genes such as SPON2 were elevated) (Figure 3d). Additionally, we found high correlation of expression many genes before and after transmutation, including CXCL10 (r = 0.20), HLA-DRA (r = 0.16), and HLA-DRB1 (r= 0.15). These genes are especially relevant, as they were explicitly featured in Lee et al. (2020) key inflammation genes that distinguish
8

---

Published in Transactions on Machine Learning Research (08/2024)
a)
d)
4
9 9-44
SPON2
MSHAGA
Anyelic - Cyclic
4 4 4
4
9
9
9
800
1250
600
1020
Inter.
Inieer
400
500
9
9
9
/
4
200
230
0.03
0.04
0.00
b)
v-y
Y-VV
PRF
VCAN
Feature
Teatiiee
Mhe rerers oittzon
Mite remere cilturen
U 1
400
200
0.30 0.0
Nonhhlogenateed Holagenated
0-0.0083
0.8354
HGHBP2
SERMINA1
1000
Bufore To ward diflas on
Before forward ciffusiin
1250
Correlations of features
Comelationn all testures
600
750
500
200
250
8.00
1.5
CD16+ HK (true)
Mono. CD16- HK (gen)
0.5
(gan.)
0.8
0.1
92
0.3
04 0.3
0.
01
0.2
63
Phaaroon conrelation
Figure 3: Transmutation between classes, a) From our branched diffusion model trained on MNIST, we show examples of 4s transmuted to 9s (left), and 9s transmuted to 4s (right). We also show the diffusion intermediate at the branch point. b) From our branched model trained on labular letters, we show the scatterplots of some feature values before and after transmutation from Vs to Ys (left), or Ys to Vs (right). For each of the 16 features, correlate the feature value before versus after transmutation and show histogram of the correlations over all 16 features in either transmutation direction. c) From our branched model trained on the single-cell RNA--sq dataset, transmuted between CD16+ NK cells and classical monocytes, and show the distribution of several marker genes before and after transmutation. The left column shows marker genes of classical monocytes, and the right column shows marker genes of CD16+ NK cells. d) From our branched model trained on ZINC250K, transmuted between acyclic and cyclic molecules, and between non-halogenated and halogenated molecules.
cells from healthy cells. Notably, we recovered these strong correlations even though the original expression of these genes showed little to no distinction between healthy and infected cells (Supplementary Figure S4.. This illustrates how our branched model successfully transmuted COVID-infecte cells of one type into COVID-infected cells of another type (and reflexively, healthy cells from one type into healthy cells of another type).
Finally, trained branched diffusion models on ZINC250K (Supplementary Methods), another large real- world dataset and an entirely different data modality: molecular graphs. We trained branched diffusion models to conditionally generate acyclic and eyclic molecules, or halogenated and non- hon-halogenated molecules. We then transmuted molecules from one property class to another, while largely retaining core functional groups (e.g. amines, esters, sulfonamides, etc.) (Figure 31). Quantitatively, transmutation from acyclic to cyelic molecules was 96.5% effective (i.e. from 0% of molecules having a cycle, we transmuted to 96.5% of molecules having a cycle). We then quantified the preservation of functional groups by computing the Jaccard index before and after transmutation. A t-test (compared to the whole set of cyclic molecules) returned p 7.57 x 10-10, On the transmulation of halogenated to non-halogenated molecules, our transmutation was 100% effective, and a t-test on the preservation of functional groups returned 4.55 x 10 10-11
Across our many datascts, these results together qualitatively and quantitatively show that transmutation in branched diffusion models is both: 1) effective defining features of the source class are removed and defining fcatures of the target class are generated; and 2) analogous- features which characterize the original object/instance (but do not directly define its class) are preserved.
7 Interpretability of branch-point intermediates
Interpretability is a particularly useful tool for understanding data, and is a cornerstone of AI for science. Unfortunately, there is limited work (if any) that attempts to improve or leverage diffusion-model inter-
9

---

Published in Transactions on Machine Learning Research (08/2024)
pretability. By explicitly encoding branch points, branched diffusion model offers unique insight into the relationship between classes and individual objects from generative perspective.
In the forward-diffusion process of a branched diffusion model, two branches meet at a branch point when the classes become sufficiently noisy such that they cannot be distinguished from each other. Symmetrically, in the reverse-diffusion process, branch points are where distinct classes split off and begin reverse diffusing along different trajectories. Thus, for two similar classes (or two sets of classes), the reverse-diffusion intermediate branch point naturally encodes features which are shared (or otherwise intermediate or interpolated) between the two classes (or sets of classes).
In particular, hybrid intermediates represent partially reverse-diffused objects right before a branch splits into two distinct classes. Formally, consider the set of all branches {(s;, li, C.) For two classes e and (2, let to be the first branch point (earliest in diffusion time) in which e1 and e2 are both in the same branch. We define a hybrid object 24 between classes 6₁ and 0g as an object sampled from the partially diffused distribution at from the conditional distribution of ei or e2:
4= = minfsrler,, E C.) oh ~ Pr. (x(c)) = Pt. (x)
In practice, sampling 2n from Pt, (2r, ei) Pt. (I, 02) is done by performing reverse diffusion following Algorithm 2 from time T until time to for either e1 or 02 (it does not matter which, since we stop at t6).
For example, on our MNIST branched diffusion model, hybrids tend to show shared characteristics that underpin both digit distributions (Figure b). On our branched model trained on tabular letters, we see that hybrids tend interpolate between distinct feature distributions underpinning the two classes, acting smooth transition state between the two endpoints (Figure
a)
4-9
1-7
b)
Hybrid averages
4
4
7
4
4-9
1-7
3-5
Hybrid
7
+
Hybrd
{ 7
9
9
7
9
9
7
7
1
7
c)
0-X
E-F
Festure
Feature
Festure 12
Feeture 16
0.8
0.8
0.8
0.8
Hybrid
Hybrid
0.6
0.6
0.4
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
Figure 4: Interpretable hybrids at branch points. a) From our branched model trained on MNIST, we show examples of hybrids between the digits classes and 9 (left), and between the digit classes 1 and 7 (right). Each hybrid in the middle row is the reverse-diflusion starting point for both images above and below it. We applied a small amount of Gaussian smoothing to the hybrids for of viewing. b) Averaging over many samples, the aggregate hybrids at branch points show the collective characteristics that are shared between MNIST classes. c) From our branched model trained tabular letters, show the distribution of some features between two pairs of classes O and X (left), and E and F (right)-and the distribution of that feature in the generated hybrids from the corresponding branch point.
Note that although a branched diffusion model can successfully generate distinct classes even with very conservative branch points (i.e. late in diffusion time), the interpretability of the hybrid intermediates is best when the branch points are minimal (carliest diffusion) times of indistinguishabliity. Taken to the extreme, branch points close to t == T encode no shared information between classes whatsoever, as the distribution of objects at time T is independent of class.
10

---

Published in Transactions on Machine Learning Research (08/2024)
8 Efficient multi-class sampling from branched diffusion models
In addition to advantages in continual learning, transmutation, and interpretabilityy branched models also offer a minor benefit in generative efficiency, because partially reverse-diffused intermediates at branch points can be cached and reused. Note that branched models and standard linear models take the same amount of computation to generate a single class, but branched models enjoy significant savings in computational efficiency when sampling multiple classes (Supplementary Table S9
9 Discussion
In this section, we explore some of the trade-offs and caveats of branched diffusion models.
Firstly, the multi-task neural network behind branched diffusion model is crucial to its efficient parameteri- zation. As previously discussed, the number of branches scales linearly with the number of classes, so the multi-task architecture is relatively elficient even for dalasets with large number of classes. Still, we recognize that an extremely large number of classes could become a bottleneck. In those cases, the branched diffussin architecture could benefit. from recent advancements in efficient multi-task parameterizations Vaadenhendd
et al. 2022
Additionally, the advantages and performance of branched difffusion models rely on appropriately defined branch points. We performed a robustness analysis and found that although the underlying branch points are important, the performance of branched diffusion models is robust to moderate variations in these branch points (Figure S5). Our branch-point discovery algorithm (Appendix A) is also agnostic to the diffusion process, and although it relies on Euclidean distance between noisy objects (which may be hard to compute for data types like graphs), the algorithm (and subsequent diffusion) can always be applied in latent space to guarantee well-defined Euclidean distances.
Finally, branched diffusion models may have difficulties learning on certain image datasets where the class- defining subject of the image can be in different parts of the image, particularly when data may be sparse. For datasets like MNIST, the digits are all roughly in the center of the image, thus obviating this problem. Of course, images and image-like data are the only modalities that suffer from this issuc. Additionally, this limitation on images may be avoided by diffusing in latent space.
10 Conclusion
In this work, we proposed a novel form of diffusion models which introduces branch points which explicitly encode the hicrarchical relationship between distinct data classos. Branched diffusion models are an alternative method of conditional generation for discrete classes. Compared to the current state-of-the-art conditional diffusion models, we showed numerous ad advantages of branched models in conditional generation. We showcased these advantages across many different datasets, including several standard benchmark datasets and two large real-world scientific datasets.
Firstly, we showed that branched models are easily extendable to new, never-before-seen classes through an efficient fine-tuning step which does not lead to catastrophic forgetting of pre-existing classes. This can enhance diffusion- modill training in online-learning settings and in scientific applications where new data is constantly being produced experimentally. Additionally, branched models are capable of more sophisticated forms of conditional generation, such as the transmutation of objects from one class into the analogous object of another class. Using Iransmutation, we demonstrated the ability of branched diffusson models discover relevant biology and chemistry. Furthermore, we showed that branched models can offer some insights into interpretability. Namely, reverso-diffusion intermediates at branch points are hybrids which encode shared or interpolated characteristics of multiple data classes.
Finally, because branched diffusion models operate on the same underlying diffusion process as a conventional linear model, they are flexibly applied to virtually any diffusion-mode paradigm (c.g. continuous or discrete time, SDE based or Markov-chain based, different noise schedules and definitions of the noising process, etc.). Branched models also easily combined with existing methods which aim to improve training/ saampinng
11

---

Published in Transactions on Machine Learning Research (08/2024)
efficiency generative performance (Kong & Ping 2021 Watson et al. 2021 Dockhorn et al. 2021 Song
et al.
2021 Xiao et al. 2022), other methods which condition based external properties (Song et al.
2021
Ho et al. )021).
Branched diffusion models have many direct applications, and we highlighted their usefulness in scientific settings. Further exploration in the structure of diffusion models (e.g. branched vs linear) may continue have resounding impacts in how these models are used across many areas.
12

---

Published in Transactions on Machine Learning Research (08/2024)
References
Nicole Almanzar, Jane Antony, Ankit S. Baghel, Isaac Bakerman, Ishita Bansal, Ben A. Barres, Philip A. Beachy, Danicla Berdnik, Biter Bilen, Douglas Brownfield, Corey Cain, Charles K.F. Chan, Michelle B. Chen, Michael F. Clarke, Slephanie D. Conley, Spyros Darmanis, Aaron Demers, Kubilay Demir, Antoine de Morree, Tessa Divita, Haley du Bois, Hamid Ebadi, F. Hernán Espinoza, Matt Fish, Qiang Gan, Benson M. George, Astrid Gillich, Rafael Gömez-Sjöberg, Foad Green, Geraldine Genetiano, Xueying Gu, Gunsagar S. Gulati, Oliver Hahn, Michael Seamus Haney, Yan Hang, Lincoln Harris, Mu He, Shayan Hosseinzadeh, Albin Huang, Kerwyn Casey Huang, Tal Tram, Taichi Isobe, Feather Ives, Robert C C. Jones, Kevin S. Kao, Jim Karkanias, Guruswamy Karnam, Andreas Keller, Aaron M. Kershner, Nathalie Khoury, Seung K. Kim, Bernhard M. Kiss, William Kong, Mark A. Krasnow, Maya E. Kumar, Christin S. Kuo, Jonathan Lam, Davis P. Lee, Song E. Lee, Benoit Lehallier, Olivia Leventhal, Guang Li, Qingyun Li, Ling Liu, Annic Lo, Wan Jin Lu, Maria F. Lugo-Fagundo, Anoop Manjunath, Andrew P. May, Ashloy Maynard, Aaron McGeever, Marina McKay, M. Windy McNerney, Bryan Merrill, Ross J. Metzger, Marco Mignardi, Dullei Min, Ahmad N. Nabhan, Norma F. Neff, Katharine M. Ng. Patricia K. Nguyen, Joseph Noh, Roel Nusse, Róbert Pálovies, Rasika Patkar, Weng Chuan Peng, Lolita Penland, Augela Oliveira Pisco, Katherine Pollard, Robert Puccinelli, Zhen Qi, Stephen R. Quake, Thomas A. Rando, Eric J. Rulifson, Nicholas Schaum, Joe M. Segal, Shaheen S. Sikandar, Rahul Sinha, Rene V. Sit, Justin Sonnenburg, Daniel Staehli, Krzysztof Szade, Michelle Tan, Weilun Tan, Cristina Tato, Krissie Tellez, Laughing Bear Torrez Dulgeroff, Kyle J. Travaglini, Carolina Tropini, Margaret Tsui, Lucas Waldburger, Bruce M. Wang, Linda J. van Weele, Kenneth Weinberg, Irving L. Weissman, Michael N. Wosczyna, Sean M. Wu, Tony Wyss-Coray Jinyi Xiang, Soso Xue, Kevin A. Yamauchi, Andrew C. Yang, Takshmi P. Yerra, Justin Youngyunpipatkul, Brian Yu, Fabio Zanini, Macy E. Zardeneta, Alexander Zee, Chunyu Zhao, Fan Zhang, Hui Zhang, Martin Jinye Zhang, Lu Zhou, and James Zou. A single-cell transcriptomic atlas characterizes ageing tissues in the mouse. Nature 2020 583:7817, 583:590- 595, 7 2020. ISSN 1476-4687. doi: 10.1038/ 11115560000-496-11 110-4-66800022229-1111
URL https: www. nature. www.a//ueew...uureee..mm com sew.cceeesseeeeeeeeee 020-0000-33000 2496-1
Scbastian Bischoff, Alana Darcher, Michael Deistler, Richard Gao, Franziska Gerken, Manuel Glocckler, Lisa Haxel, Jaivardhan Kapoor, Janne K Lappalainen, Jakob H Macke, Guy Moss, Matthijs Pals, Felix Pei, Rachel Rapp, A Erdem Sağtekin, Cornelius Schröder, Auguste Schulz, Zinovia Stefanidi, Shoji Toyota, Linda Ulmer, and Julius Vetter. A practical guide to statistical distances for evaluating generative models in science. 3 2024. URL https: //tr://rxxii org//bs////aeo.ooooooooee 2403. 12636v1
C. Domínguez Conde, C. Xu, L. B. Jarvis, D. B. Rainbow, S. B. Wells, T. Gomes, S. K. Howlett, O. Suchanek, K. Polanski, H. W. King, L. Mamanova, N. Huang, P. A. Szabo, L. Richardson, L. Bolt, E. S. Fasouli, K. T. Mahbubani, M. Prete, L. Tuck, N. Richoz, Z. K. Tuong, L. Campos, H. S. Mousa, E. J. Needham, S. Pritchard, T. Li, R. Elmentaite, J. Park, E. Rahmani, D. Chen, D. K. Menon, O. A. Bayraktar, L. K. James, K. B. Meyer, N. Yosef, M. R. Clatworthy, P. A. Sims, D. L. Farber, K. Sacb-Parsy, J. L. Jones, and S. A. Teichmann. Cross-tissue immune cell analysis reveals tissue-speeific features in humans. Science, 376, 2022. ISSN 10959203. doi: 10.1126/SCCI /S/CSSCINEEAII ABBLNNCCAABLL11/SSPPLL BtCH://UUPPLFLL//SSIIEEEEE FILE/SCIENCE, ABL5197 MDARR REPRODUCIBILITY SCHCCCCCCECCCHHCCCC PDF. URL httts:://ww.ccieet org/do1/10 htw.iiiencee..b11991
Prafulla Dhariwal and Alex Nichol. Diffussio models beat gans on image synthesis. Advances in Newral Information Processing Systems, 34:8780-877 8794, 12 2021.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with critically-dam langevin diffusion. 12 2021. doi: 10:66350/amxiiv2211.00006 URL httll/ttss//aaiiiorg//////////// 0706684
Joe Felsenstein. Phylip. URL
Peter W. Frey and David J. Slate. Letter recognition using holland-style adaptive classifiers. Machine doi:
Learnang, 6:161-182, 1991. ISSN 15730565.
URL https
//link.springg com www./rticle/////////oooo 1ttp://00..000///122200000000000
Adam Gayoso, Romain Lopez, Galen Xing, Pierre Boyeau, Valch Valiollah Pour Amiri, Justin Hong, Katherine Wa, Michael Jayasuriya, Edouard Mehlman, Maxime Langevin, Yining Liu, Jules Samaran, Gabriel Misrachi,
13

---

Published in Transactions on Machine Learning Research (08/2024)
Achille Nazaret, Oscar Clivio, Chenling Xu, Tal Ashnach, Mariano Gabitto, Mohammad Lotfollahi, Valentine Svensson, Eduardo da Veiga Beltrame, Vitalii Kleshchevnikov, Carlos Talavera-Lóóez, Lior Pachter, Fabian J. Theis, Aaron Streets, Michael I. Jordan, Jeffrey Regier, and Nir Yosef. A python library for probabilistic analysis of single-cell omics data. Nature Biotechnology 2022 40:2, 40:163-116, 2 2022. ISSN 1546-1696. URL
doi: 10. (00000//111870011--0100
Xiaoping Han, Ziming Zhou, Lijiang Fei, Huiyu Sun, Renying Wang, Yao Chen, Haide Chen, Jingjing Wang, Huanna Tang, Wenhao Gc, Yincong Zhou, Fang Ye, Mengmeng Jiang, Junqing Wu, Yanyu Xiao, Xiaoning Jia, Tingyue Zhang, Xiaojie Ma, Qi Zhang, Xueli Bai, Shujing Lai, Chengxuan Yu, Lijun Zhu, Rui Lin, Yuchi Gao, Min Wang, Yiqing Wu, Jianming Zhang, Renya Zhan, Saiyong Zhu, Hailan Hu, Changchun Wang, Ming Chen, He Huang, Tingbo Liang, Jianghua Chen, Weilin Wang, Dan Zhang, and Guoji Guo. Construction of human cell landscape at single-cell level. Noture 2020 581:7808, 581:303 309, 3 2020. ISSN 1476-4687. doi: 10. (00008//41500 111155880000222550
URL. httpp://www. http///ww...rrwo..m////ea nww.nattuueeoo.mm 000 2157-44
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabillisic models. Advances in Neural Infor- malion Processing Systems, 33:6840 6851, 2020. URL
Jonathan Ho, Google Research, and Tim Salimans. Classifier-free diffusion guidance, 11 2021.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J. Fleet. Video diffussin models. Advonces in Newral Information Processing Systems, 35:8633 8646, 12 2022. URL https://video- htt.s/fusinnnnngtu..oooooee io/.
John J. Irwin, Teague Sterling, Michael M. Mysinger, Erin S. Bolstad, and Ryan G. Coleman. Zine: A free tool to discover chemistry for biology. Journal of Chemical Information and Modeling, 52:1757- CI3E/C13001277 SI 001.PDF. URL
1768, 7 2012. ISSN 1549960X. doi: 10.1021 /T110//CT300227/SPPPP https://pubs. acs wtt.s.or///o//fl////ooo.o 1021/c13001277
Jaehyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based generative modeling of graphs via the system of stochastic differential equations. 2 2022. doi: 01.488550/axxiv22222..222 URL https: org/ / 2202.02514v2
Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabilistic models, 6 2021.
Akim Kotelnikov, Dmitry Baranchuk, Ivan Rubachev, and Artem Babenko. Tabddpm: Modelling tabular data with diffusion models. 9 2022. doi: 10.48550/araa 0155500arxvv.2099115553 2209.222221 URL https://arxiv wwg./rgg/////....ooo abs/2299
15421v1
Yann LeCun, Corinna Cortes, and Chris Burges. Mnist handwritten digit database. URL http://yann
lecun. www.auu.c.//eed////eeeee com/exxb/nnist//
Jeong Scok Lee, Seongwan Park, Hye Won Jeong, Jin Young Ahn, Scong Jin Choi, Hoyoung Lee, Backgyu Choi, Su Kyung Nam, Moa. Sa, Ji Soo Kwon, Su Jin Jeong, Heung Kyu Lee, Sung Ho Park, Su Hyung Park, Jun Yong Choi, Sung Han Kim, Inkyung Jung, and Eui Cheol Shin. Immunophenotyping of covid-19 and influenza highlights the role of type interferons in development of severe covid-19. Science Immanology, 5: 1554, 7 2020. ISSN 24709468. doi: 00011000SCIIMMCUNOLAAADSIISSS /ALLL FILE/TABLE_S9, S9.XLSX. URL https::///// http://wwwwwcwnccce/// science. www.iiinneerrg/do/ooeeee org//oi/////oooooooooo
abd1554
Mohammad Lotfollahi, Mohsen Naghipourfar, Malte D. Luecken, Matin Khajavi, Maren Büttner, Marco Wagenstetter, Ziga Avsec, Adam Gayoso, Nir Yosef, Marta Interlandi, Sergei Rybakov, Alexander V. Misharin, and Fabian Theis. Mapping single-cell data to reference atlases by transfer learning. Nature Biotechnology 2021 40:1, 40:121 8 2021. ISSN 1546-1696. doi: 10. 10.1008//1111880011000000-1111 URL
www.://www.w..cocccc cwm http:///ww....aaaaaooooooooaaoar
Zeyu Lu, Chengyue Wu, Xinyuan Chen, Yaohui Wang, Lei Bai, Yu Qiao, and Xihui Liu. Hicrarchical diffusion autoencoders and disentangled image manipulation. 4 2023. URL. https: //trx///ivvvv///////// org//bs////00 11829v2
14

---

Published in Transactions on Machine Learning Research (08/2024)
Bo Qiang, Yuxuan Song, Minkai Xu, Jingjing Gong, Bowen Gao, Hao Zhou, Weiying Ma, and Yanyan Lan. Coarse-to-fine: hierarchical diffusion model for molecule generation in 3d. 9 2023. URL https
//rxiv..rg//bs//2230 13266v2
Aaditya Ramdas, Sashank J. Reddi, Barnabas Poezos, Aarti Singh, and Larry Wasserman. On the high- dimensional power of linear-time kernel two-sample testing under mean-difference alternatives. 11 2014.
URL http://axxxv...g/bbbbbaa
Aviv Regev, Sarah A. Teichmann, Eric S. Lander, Ido Amit, Christophe Benoist, Ewan Birney, Bernd Bodemmiller, Peter Campbell, Piero Carninci, Menna Clatworthy, Hans Clevers, Bart. Deplancke, Tan Dunham, James Eberwine, Roland Eils, Wolfgang Enard, Andrew Farmer, Lars Fugger, Berthold Göttgens, Nir Hacohen, Muzlifah Haniffa, Martin Hemberg, Seung Kim, Paul Klenerman, Arnold Kriegstein, Ed Lein, Sten Linnarsson, Emma Lundberg, Joakim Lundeberg, Partha Majumder, John C. Marioni, Miriam Merad, Musa Mhlanga, Martijn Nawijn, Mihaí Netea, Garry Nolan, Dana Pe'er, Anthony Phillipakis, Chris P. Ponting, Stephen Quake, Wolf Reik, Orit Soxemblatt-Rosen Joshua Sanes, Rahul Satija, Ton N. Schumacher, Alex Shalek, Ehud Shapiro, Padmanee Sharma, Jay W. Shin, Oliver Stegle, Michael Stratton, Michael J.T. Stubbington, Fabian J. Theis, Matthias Uhlen, Alexander Van Oudenaarden, Allon Wagner, Fiona Watt, Jonathan Weissman, Barbara Wold, Ramnik Xavier, and Nir Yosef. The human cell atlas. eLife, 6, 12 2017. ISSN 2050084X. doi: N0.554//ELIFEE
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjorn Ommer. High-resolution image synthesis with latent diffusion models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10674-10685 6 2022. doi: 10.1109/CVP522688.22 20022.0042202.0114 URL httpps://ieeexploreee ieee. org/ /ttp//cmmmmme///8/4999///////
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. 32nd International Conference on Machine Learning, ICML 2015, 3:2246-2255 3 2015. doi: 10.48550/arxiv. 150.....55.0000 03555. URL httpss///axivi http://axxvvvooggooo abs 1503. 03585v8
Yang Song, Jascha Sohl-Dickstein, Google Brain, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. 2021.
Gábor J. Székely and Maria. L.. Rizzo. Energy statistics: A class of statistics based distances. Journal of Statistical Planning and Inference, 143:1249 1272, 8 2013. ISSN 0378-3758. doi: 10.1016/1..5P1120133.00000
Gido M. de Ven and Andreas S. Tolias. Three scenarios for continual learning. 4 2019. doi: 10.48550/ arxiv.190..07734. URL https://arxiv. org. wttp://wwwww///////000/////333
Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc Proesmans, Dengxin Dai, and Lue Van Gool. Multi-task learning for dense prediction tasks: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44:3614-366 2022. ISSN 19393539. doi: 10.1109/TPAMT.200 3054719.
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning to efficiently sample from diffusion probabilistic models. 6 2021. doi: 00.48550/arxiv.222 wtt.i...1060080022 URL https:///axxi: 03802v1
F. Alexander Wolf, Philipp Angerer, and Fabian J. Theis. Scanpy: Large-scale single-cell gene expression data analysis. Genome Biology, 19:1-5, 2018. ISSN 1474760X. doi: URL https://genomebiology oomddcenntrall
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans, 1 2022. URL
15

---

Published in Transactions on Machine Learning Research (08/2024)
A Branch-point discovery
As objects from two different classes are forward diffused, their distributions become more and more similar to each other. At the extreme, it is expected that all objects (regardless of class) reach the same prior distribution the time horizon T. Our goal is quantify the earliest time point when objects of two different classes are sufficiently similar in distribution, such that the reverse diffusion after that point can be predicted by the same model between the two classes (without specifying class identity).
More formally, we define branch point between two classes and e2 as the diffuusin time ls where the objects sampled from these two classes are "relatively indistinguishable" from each other. Because the data distribution is typically complex and high dimensional, we quantify indistinguishabblity as log-fold change of expected Euclidean distances, comparing data drawn belween the two different classes, and data within the same class. That is, the branch point to between classes e1 and Og is defined as the ménimum time such that:
1 mbll2 <e I|. xolle
]]
EugD, (EEE,,,,ED,,cee cED,, Baa-a(slle
log(
(8)
where Dj is the distribution of data with class label G;.
We discuss theoretical justifications and connections Appendix B
The outer expectations are taken over points in the dataset, and the inner expectations are over the forward noising process. When this log-fold change is small number e, then we consider the classes relatively indistigguisaable (i.c. the average distance between noisy objects of different classes is comparable to the average distance between noisy objects of the same class).
Of course, these expectations are extremely intractable to compute in closed-form because of the large dataset and dimensionality, so instead we approximate these quantities by using Monte Carlo sampling. That is, we take sample of objects from each class, apply forward diffusion at many times t spaced regularly between and T, and identify the smallest i such that log-fold change of the distance is sufficiently small.
This procedure gives branch point 4 for every pair of classes C;, Cj (i ())
In branched diffusion model, each branch 6, = (s;, t4, C₁) learns to reverse diffuse between times [st, tt) for classes in C. The branches form a tree structure (i.e. hierarchy) with the root at time and a branch for each individual class at time 0. In order to convert the branch points between all pairs of classes into a hierarchy, we simply perform greedy aggregation starting from individual classes and iteratively merge classes together by their branch points (from early late diffusion time) until all classes have been merged together.
To summarize, the full branch-point discovery algorithm is as follows:
1. Start with dataset of objects to generate, consisting of classes C.
2. For each class, sample objects randomly and without replacement.
3. Forward diffuse each object over 1000 time points in the forward-diffusson process (we used 1000 steps, as this matched the number of reverse-diffusion steps we used for sample gencration). The branched diffusion model which will be trained using these branch definitions employs an identical forward-diffusion process.
4. At each time point I, compute the average Euclidean distance of each pair of classes, resulting in |C| |C| matrix at each of the 1000 time points. For distinct classes Ca Cj (i j), the distance s(t, Gi Gj) is computed over the average of n pairs, where the pairs are randomly assigned between the two classes; for self-distance of class G, the distance s(t, c, ci) is computed over the average of n pairs within the class, randomly assigned such that the same object is not compared with itself.
For each pair of classes c, Cj (i may be equal to f), smooth the trajectory of s(t, G, C3)Cg) over time by applying a Gaussian smoothing kernel of standard deviation equal to and truncated to 4 standard deviations each side.
16

---

Published in Transactions on Machine Learning Research (08/2024)
6. For each pair of distinct classes C,, Gj (i * j), compute the earlist time in the forward-diffusion process such that the log-fold change of the average distance between and ej over the self-distance of G and Cj (averaged between the two) is at most a tolerance of e. That is, for each pair of distinct classes Ca Cj (i * j), compute the minimum such that log( ((o(h.na))Hllte..e)))) pair of distinct classes "minimal time of indistinguishabilitt"" Topey
s(be,,))
This gives each
7. Order the (2) minimal times of indistinguishabilityy т by ascending order, and greedily build hicrarchical tree by merging classes together if they have not already been merged. This can be implemented by a set of |C] disjoint sets, where each set contains one class; iterating through the times T in order, two branches merge into a new branch by merging together the sets containing the two classes, unless they are already the same set.
B Theoretical justification for branch-point-discover algorithm
The goal our branch-point discovery algorithm is to find the branch point between any pair of classes en and Cg. The branch point the earliest point in diffusion time where e1 and o2 are "indisinguushable" from each other. Formally, indistinguishabiiityy occurs when the conditional distributions of these two classes are close enough that they can be modeled by single reverso-diffusion process. That is, want the branch point ls to be such that gt, (x/(c)) x 4t, (z/c2))
Of course, this is difficult condition to formally define and satisfy, and there are several challenges:
We only have access to samples of objects from each class' conditional distribution.
2. The data can be high dimensional, and particularly because we only have limited samples, the curse of dimensionality emerges quickly, even with as few as 10 dimensions.
3. These noisy conditional distributions ge (z(cc) are not tractable to compute or even represent, particu- larly at carlier diffusion times (when the distribution (z(z) is so close to the unknown and complex data manifold (o(x)).
4. The noisy conditional distributions for two different classes will generally not be identical (until the time horizon T), and they slowly approach indistinguishabiity asymptotically.
In order to find the optimal to, we must compare q( c1) and (x(c2) at potentially all time points along the diffusion timeline, well as across all pairs of classes in the dataset, so this comparison needs to be extremely computationally efficient.
We designed our branch-point-discover algorithm based on these challenges and desiderata, leading to our definition a branch point in Equation g
Notably, Equation g is closely related eneryy distance (Székely & Rizzo 2013 Energy distance compares the expected distance between objects of two different distributions, to the average expected distance within the same distribution:
E(D,, D2) = DE..ED,
Di BED, Illa - -1/2] - EaE BED, |l|a - -b/2]] - EagD, Illa -b/2]]
We chose to use energy distance for several reasons:
Energy distance (by definition) compares the expected distance between two different distributions to the expected distance within the distributions. This allows us to altain a measure of "relative indistinguishabbiitty"" By picking a time point when the energy distance low, we ensure that the model upstrcam of the branch point has no more difficulty representing and learning both classes together, as would to learn only a single class (as in standard, linear diffusion model). This addresses Challenge 4.
17

---

Published in Transactions on Machine Learning Research (08/2024)
Energy distance is much more robust in high dimensions. In high dimensions, particularly with sparser samples, metrics such as MMD and significance testing such as t-tests can be unreliable (Randas et al. 2014 Bischoff et al. 2024 This addresses Challenges 1 and 2.
Energy distance does not require any distributional assumptions. Furthermore, we modify the energy distance formulation slightly, that instcad of taking the difference between expected inter- distribution and intra-distribution distances, we take the ralio, as this allows for the energy distance to be more comparable in magnitude across different classes and times, even as the distributions change drastically across diffusion time. It also allows us to select c in a more principled and comparative manner across different datasets. In contrast, most frameworks for significance lesting relies on strong distributional assumptions, particularly in high dimensions. This addresses Challenge 3.
Energy distance is also extremely efficient to compute. It involves computing expected distance between objects of the two distributions, as well as between objects of the same distribution. our algorithm, we perform Moute Carlo sampling over possible random couplings between objects lo compute these expected distances. This make energy distance far more suitable for our needs compared other metrics such as Wasserstein distance, which requires minimizing over all possible couplings. This addresscs Challenge 5.
Energy distance is calculated simply using Euclidean distance, and there intuitive reason for why Euclidean distance in diffusion space meaningful measurement distance ween distributions. The Euclidean distance can measure how much "work" diffusion model does push objects from one point to another. In particular,
view the reverse-diffusion process as taking sampled object from (Ix) and pushing the feature values to final generated sample. For single branch representing multiple classes over some diffusion time interval, distances between objects reflect the expected distance the branch will be pushing objects, and we would like this distance to be no more for objects of different classes versus objects of the same class in that branch.
18

---

Published in Transactions on Machine Learning Research (08/2024)
C
Supplementary Figures and Tables
Continuous-time diffusion model
Discrete-time diffusion model
O
/
1
1
E
2
2
R
a
2
3
2
2
2
2
2
2
2
2
A
2
2 3
3
2
5
3
5
3
3
3
3
9
3
3
3
3
3
3
3 4
3
4
4
4
4
4
4
4
4
4
4
4 4
4
4
4
|
4
4 5
6
5
5
3
5
D
5
5
5 6
5
5 5
5
5"
5
5
5
5
5 6 7 8 9 9
6
7
6
6
6
6
6
0
6
6
6
4 6
6
6
6
6
6
7
7
7
7
7
1
7
7
7
7
7
7
7
7
7
7
7 8
8
6
8
8
Q
8
8
f
F 8
8
8
8
8
8
7
7
9
9
9
q
9
9
3
7 0
9 7
Q
9
9
9
9
9
Figure S1: Examples of generated MNIST images. We show (uncurated) images of MNIST digits generated by branched diffusion models. Since branched diffusion models naturally output each class separately, generation individual classes does not require supplying labels pretrained classifiers. We show sample of digits generated from a contimuous-time (score-matching) diffusion model (Song et al.] 2021). and discrete-time diffusion model (denoising diffusion probabilistic model) (Ho et al. 2020). Branched diffusion models for multi-class generation fit neatly into practically any diffusion-model framework.
19

---

Published in Transactions on Machine Learning Research (08/2024)
a)
c)
Distribution of Wasserstein distance
True correlations
Generated correlations
between features
2.5
True generated features
True vs true features
1.0
2.0
Class H
0.9
1.5
0.8
1.0
0.7
0.5
0.0
0.0
0.5
1.0
1.5 2.0 2.5 3.0 3.5
Wasserstein distance
Class K
0.5
b)
P Q, Feature 12
True
Generated
1.4
True
0.3
1.2
Generated
1.0
0.8
Class X
0.1
0.6
0.4
0.2
0.0
-2
Figure S2: Examples of generated letters. We show some examples of distributions generated from a branched diffusion model trained on tabular data; English letters of various fonts, featurized by a hand-engineered set 16 features (Frey & Slate 1991). a) For each letter class and each of the 16 numerical features, computed the Wasserstein distance (i.e. earthmover's distance) between true data distribution and the generated data distribution. We compare this distribution of Wasserstein distances to the distribution of Wasserstein distances between different true average, branched diffusion model learned to generate features which are similar distribution to the true data. b) We show example of the true and generated feature distributions for particular feature, comparing two letter classes: and Q. Although the two classes show a very distinct distribution for this feature, the branched diffuusin model captured this distinction well and correctly generated the feature distribution for each class. c) Over all 16 numerical features, we computed the Pearson correlation between the features, and compared the correlation heatmaps between the true data and the generated examples. In each of these three classes, the branched diffusion model learned to capture not only the overarching correlational structure shared by all three classes, but also the subtle secondary correlations unique each class.
features as a baseline.
20

---

Published in Transactions on Machine Learning Research (08/2024)
a)
b)
FID between true and generated digits
FID between true and generated letters
Branched
20
100
Label-guided (linear)
Branched
80
15
Label-guided (linear)
FID
60
10
40
5
20
0
0
0
3
4
5
8
9
ABCDEFGH LIRQQPRSSS VWX
MNIST class
Letter class
C)
FID between true and generated cells
1500
Branched
Label-guided (linear)
1250
1000
FID
750
500
250
0
CD16 6+ NK
Cl. Mono.
Late Eryth.
Macroph.
Megakar.
Mem:
NK
Plasmabl
Tem/Eff. H. T
Cell class
Figure S3: Sample quality of branched diffusion vs label-guided (linear) diffusion. We compare the quality of generated data from branched diffusion models to label-guided (linear) diffusion models of similar capacity and architecture. For each class, we computed the Fréchet inception distance (FID) between the generated examples and a sample of the true data. A lower FID is better. We show the FID for generated a) MNIST digits; b) tabular letters; and c) single-cell RNA-seq. We find that our branched diffusion model achieved comparable sample quality compared to the current state-of-the-arl method of label-guided diffusion. In cases, the branched model consistently generated better examples.
21

---

Published in Transactions on Machine Learning Research (08/2024)
Table S1: MNIST branch definitions
Branch start (8))
Branch end (t,)
Branch classes (C,)
0.4855
1
0,1,2,3,445,66,7888,
0.4474
0,4855
1,2,3,4,5,6,778899
0.4334
0.4474
2,3,4,5,6,,,899
0.4164
0.4334
2,3,4,5,7,8,9
0.3744
0.4164
3,4,5,7,8,9
0.3684
0.3744
3,4,5,8,9
0.3524
0.3684
3,4,5,9
0.3483
0.3524
3,4,5
0.2713
0.3483
3,5
0.4855
0
0.4474
0.4334
0
0.4164
0
0.3744
0
0.3684
8
0
0.3524
9
0.3483
0.2713
0.2713
Branch definitions for model on all MNIST digits.
Table S2: MNIST (discrete) branch definitions
Branch start (s;)
Branch end (t,)
Branch classes (C.)
761
1000
0,1,2,33,4,566,7888,
760
761
0,2,3,4,5,6,7788,99
712
760
2,3,4,5,6,7,8,9
709
712
3,4,5,6,7,8,9
700
709
3,5,6,8
685
709
4,7,9
659
700
3,5,8
656
659
3,5
527
685
4,9
761
0
760
0
712
2
0
700
0
685
0
659
656
656
0
527
0
527
Branch definitions for discrete-Line model on MNIST digits.
22

---

Published in Transactions on Machine Learning Research (08/2024)
Table S3: MNIST branch definitions for 0, 4, 9
Branch start (8s)
Branch end (44)
Branch classes (C:)
0.5
0,4,9
0
0.5
0
0.35
0.5
4,9
0
0.35
0.35
Branch definitions for model MNIST digits 0, 4, and 9.
Table S4: MNIST branch definitions for 0, 4, 7, and 9
Branch start (st)
Branch end (ty)
Branch classes (C))
0.5
0,4,7,9
0
0.5
0.38
0.5
4,7,9
0
0.38
0.35
0.38
4,9
0
0.35
0
0.35
Branch definitions for model on MNIST digits 0, 4, 7, and 9.
23

---

Published in Transactions on Machine Learning Research (08/2024)
Table S5: Letters branch definitions
Branch start (s:)
Branch end (44)
Branch classes (C;)
0.5235
1
ABSDDERCCHLLLLLAANOQQQSSSUUUVUAA
0.5165
0.5235
ABCOBEFGHIAAKLANNOOQQSSSSUXXXX
0.5115
0.5165
SCIDEFGHHLJLAMNNPQRRRTUVVYYZ
0.4945
0.5115
SCDEEFGHIJJKMMOOPQRSSTTVVYYY
0.4795
0.4945
I,J
0.4725
0.4945
BCIDEFGHIKKMNORQQRSSTUVYYYYY
0.4565
0.4725
BCDEEGHIKKMNOOOQRSSSUUZX
0.4364
0.4725
F,P.T,V.Y
0.4174
0.4565
B.DDE.GHKKOOOKSS..XX
0.4134
0.4174
B.D.E.....KK..OOR.....ZZ
0.4094
0.4134
B.D.H..K..OQ....ZZZZZZ F,T,V,Y
0.4024
0.4364
0.3864
0.4094
B.D.C.I..K.OO...........
0.3814
0.3864
B.C.K.OO...S...
0.3734
0.3814
B.C.OOO........
0.3604
0.4024
F,T.Y
0.3564
0.4134
C,E
0,3534
0.3604
T,
0.3514
0.3734
B.R.S.X,
0.3413
0.3734
G.H.O.Q
0.3223
0.3514
B,S.X.Z
0,2763
0.3223
B,S,X
0.2643
0.3413
G,H.O
0,2573
0.2643
G,O
0.1562
0.2763
S,X
0.5235
W
0.5165
0.5115
L
0.4795
J
0.4795
I
0.4565
M
0.4364
0.4174
0.4094
0.4024
0.3864
D
0.3814
K
0.3604
F
0.3564
E
0.3564
0.3534
Y
0.3534
T
0.3514
R
0.3413
0.3223
BB
0.2763
HHH
0.2643
0.2573
0.2573
0.1562
S
0.1562
X
Branch definitions for model on all letters.
24

---

Published in Transactions on Machine Learning Research (08/2024)
Table S6: Single-cell RNA-seq branch definitions
Branch start (s:)
Branch end (4:)
Branch classes (C;)
0.6436
CD16+ NK, Cl. Mono., Late Eryth., Macroph., Megakar.,
Mem. B, NK, Plasmabl., Tem/ER EEf. H. T
0.5405
0.6436
CD16+ NK, Cl. Mono., Late Bryth., Macroph., Megakar.,
NK, Plasmabl., Tem/Eff. H. T
0.5085
0.5405 CL. Mono., Late Eryth., Macroph., Megakar,, NK, Tem/Eff.
H. T
0.4505
0.5405
CD16+ NK, Plasmabl.
0.3724
0.5085
CL. Mono., Late Eryth., Megakar,, NK, Tem/Eft H. T Megakar., NK, Tem/ /Ef/fff. H. T
0.3644
0.3724
0.2292
0.3724
Cl. Mono., Tate Eryth. Megakar., NK
0.1842
0.3644
0.6436
Mem. B
0.5085
Macroph.
0
0.4505
CD16+ NK
0
0.4505
Plasmabl.
0
0.3644
Tem/ EEf. H. T
0
0.2292
CL. Mono.
0
0.2292
Late Eryth.
0
0.1842
NK Megakar.
0
0.1842
Branch definitions for model on all single-cell RNA-seq cell types.
Table S7: Single-cell RNA-seq branch definitions for CD16+ NK and Classical Monocytes
Branch start (s;)
Branch end (t,)
Branch classes (Ct)
0.5796
1
CD16+ NK, Cl. Mono.
0
0.5796
CD16+ NK
0
0.5796
CL. Mono.
Branch definitions for model all single-cell ItNA-seq cell lypes CD16+ NK and Classical Monocyles
Table S8: Single-cell RNA-seq branch definitions for CD16+ NK, Classical Monocytes, and Memory Branch start (s:)
Branch end (4:)
Branch classes (C2)
0.6787
1
CD16+ NK, CL. Mono., Mem. B
0.5796
0.6787
CD16+ NK, Cl. Mono.
0.6787
Mem.
0.5796
CD16+ NK
0.5796
Cl. Mono.
Branch definitions for model on all single-cell RNA-seq cell types CD16+ NK, Classical Monocytes, and
Memory B
25

---

Published in Transactions on Machine Learning Research (08/2024)
CXCL10
HLA-DRA
HLA-DRB1
B
3
2
2
Normal
COVID-1 19 (Mild)
Flu
COVID-19 19 (Severe)
Asymptomatic
Normal
COVID-19 (Mild)
Flu
COVID-19 19 (Severe)
Asymptomatic
Normal
COVID-15 19 (Mild)
Flu
COVID 19 (Severe)
Asymptomatic
Figure S4: True expression of genes by cell population. We show the distribution true (normalized) expression of CXCLIO, HLA-DRA, and HLA-DRB1 across different cell populations (Le. /aatthhy/ooma,,,, COVID-19, etc.). The average expression in each population is denoted by a black line. Due to complex interactions between genes and the highly noisy measurements that characteristic of single-cell R.AA-eq, the distribution of the expression of these genes generally shows no trivial differences in cell populations infected with COVID-I 19 relative to healthy cells. However, important differences the distribution of these genes between COVID-19 and healthy cell populations can be recovered through more advanced computational methods, detailed in Lee et al. 2020
26

---

Published in Transactions on Machine Learning Research (08/2024)
Table S9: Efficiency of multi-class conditional generation
Datasel
Linear model
Branched model
MNIST
78.73 + 0.11
37.30 + 0.03
Letters
110.42 + 0.14
67.54 + 0.04
Single-cell RNA-seq
275.81 + 0.02
132.37 + 0.01
When generating data from multiple classes, intermediates at branch points can be cached in a branched diffusion model. For three datasets, we measure the time taken to generate one batch of each class from a
branched diffusion model, and from a label-guided (linear) model of ident ical capacity. Averages and standard errors are shown over 10 trials cach, All values are reported as seconds.
27

---

Published in Transactions on Machine Learning Research (08/2024)
a)
0.6
80
06
68
06
b)
Branched models
200
Label-guided (linear) model
20
Discovered trees Random trees
15
150
10
100
50
8.00 0.25 0.50
0.75 1.00
1.25
1.50
1.75
Branch
Class
Figure $5: Robustness branch points. a) We computed branch points and hierarchies for the MNIST dataset 10 times, each time resulting in slightly different branching structure. The variation results from randomness in sampling from the dataset, and stochasticity in the forward-diffusion process. The 10 branching structures vary not only in their branching times, but also in their topologies (above). To emphasize the variation in the hierarchies, we also overlay all 10 hierarchies on the same axes (below). b) Compared to randomly generated hierarchies, the branching structures generated by our algorithm (Appendix A have a much lower branch-score distance between themselves (p 10-25 25 by Wilcoxon lest). c) We trained branched diffusion model each of the hierarchies, and quantified generative performance using Fréchet inception distance (FID). Over all hierarchies, the FID from the branched models were relatively consistent with each other, and also generally better than the label-guided (linear) model.
28

---

Published in Transactions on Machine Learning Research (08/2024)
Table S10:
Relative size and capacity branched models vs relative training time
Branched model
Number of classes
Total branch length
Model capacity
Actual epochs
MNIST
10
11.49
8.48
Fashion MNIST
10
4.81
8.48
Letters
26
10.05
11.06
Single-cell RNA-seq
9
4.29
6.38
1.2
7INC250K
2
1.15
1.89
1
Because branched diffusion models separate the diffusion process into branches, they naturally often require more epochs to train than their linear counterparts. We found, however, that the time taken to train these models was still far less than what would be expected if the training time were an additive function of total branch length or model capacity. Here, we show (for each of our main branched models) the total branch
2/C1-1
length (i.c. the sum of diffusion times over all branches by: Σ (t)) the approximate model capacity,
and the actual training time in epochs. All values in the table above are relalive (i.e. divided by) the branched model's linear counterpart. Although training time for branched models was already far less than what is expected based on branch length or model capacity, we also suggest that training with larger batch sizes or accumulating gradient updates between batches (particularly for shorter branches) may also allow for these models to be trained even faster. We leave the exploration of this for future work.
a)
b)
120
Branched models
Label-guided (linear) model
100
80
20
0.0
0.2
0.4
0.6
0.8
1.0
0
2
4
5
7
9
Class
Figure S6: Robustness of e in branch-point discovery algorithm. The value is used in the branch-point discovery algorithm (Appendix A) determine when classes are sufficiently similar be combined branch. Although is relatively easy to select by simply choosing value where the branch points are not all too close to t = 0 or t T, here we explore the robustness of branched diffussin models different choices between 10-5 and 10-1 (uniformly sampled in logarithmic space), and computed branch points for each in our MNIST dataset. The two largest values of yielded hierarchies where the terminal branches were all length 0, so they removed from this analysis. We show overlay of the hierarchies. Note the similarity of the hierarchies here (which arise from different values of e) compared to the distribution of hierarchies Supplementary Figure S53 (which arise from random variation in sampling and forward diffusion from the same value of e). b) We trained branched diffusion model on each of the hierarchies, and quantified generative performance using Fréchet inception distance (FID). Over all hierarchies, the FID from the branched models were mostly cousistent with each other, and also mostly better than the label-guided (linear) model. The model which performed the worst arose from the largest value of e in the analysis, which resulted in the hierarchy with the shortest terminal branches (gray hierarchy in panel a).
a) sampled 10 values of
29

---

Published in Transactions on Machine Learning Research (08/2024)
e
only
only
2 only
O
o
O
0
O
0 4 9
O
O
O
O
O
0
0
Branched
4 4
4
4
4
4
4
4 9
4
4
4
4
4
4
4
4
4
4 9 7 1
4 9 7 1 2
9
9
9
9
9
9
9
9
9
9
9
9 7
9
9
9
9
7
7
7
7
7
7
7
7
7 1
7
7
1
1
1
3 2
2 2
Defore adding new wokks
After Fine-tuning 7 only
Afler fine-tuning on only
After fire-tuinng 2 only
After fine-tuiing on all data
0
0
O
0
0
7 Q 7 7 7
2
2 2 2 2 2 2 2 2 2
0
O
0
O
Label-guidee
4
4
4
4
7
7
7
4
7
2 V
L / 2 7
.
4 1
4
4
A
9
9
9
q
9
O
7
7
7
7 7
7 2 2 2
9
9
9
9 7
9
771
7
7
2
7 1
7 /
7
7
2
9
2
2
2
2
b)
Branched model
Label-guided (linear) model
300
Before fine-tuning
After fine-tuning 7
200
After fine-tuning 1
After fine-tuning 2
After fine-tuning on all classes
100
0
7
Digit class
Digit class
Figure S7: Multi-step continual learning. We continue the same analysis as shown in Figure but show multi-step continul--eearning scheme, in which 3 never-before-seet classes are iteratively introduced to a model. We started with branched diffusion model trained to generate Os, 4s, and 9s. We then iteratively introduced 7s, 1s, and 2s. With the introduction of each new digit class, we added single new terminal branch, using the branch point computed using the branch-point discovery algorithm (Appendix [A]. We then fine-tuned only that branch for only that class. For the label-guided model, we also began with a model trained to generate Os, 4s, and 9s. We then iteratively introduced 7s, 1s, and 2s. With the introduction of each new digit class, we fine-tuned the model on only that digit class. We also attempted fine-tume tomituning the label-guided model by training all digits (old and new). a) We show examples of digits generated after each fnn-ttuiig step. The branched model was able generate high-quality digits upon the addition of each new digit class, without affecting the ability to generate pre-existing classes. 'he label-guided model, however, suffered from catastrophic forgetting upon being fine-tuned each step, and largely forgot how to generate other digits. Upon fine-tuning all digit classes (which is highly inefficient procedure). the label-guided model was able to generate all classes once more, but still suffered from inappropriate crosstalk between the classes. b) We quantified the generative performance of each model using Fréchet inception distance (FID). The branched models achieved roughly the same FIDs upon fine-tuning on each new data class, whereas the FID the label-guided models suffered enormously result of catastrophic forgetting.
30

---

Published in Transactions on Machine Learning Research (08/2024)
9
Original digits
4
4
4
4
4
4
4
4
4
4
9
9
7
9
9
9
9
9
9
9
Branched model t 0.35
9
9
9
9
9
9
9
9
9
9
4
4
4
4
4
4
4
4
4
4
Línear model = 0.15
4
4
4
4
4
4
4
4
4
4
4
9
9
9
9
9
9
9
9
9
Linear model == 0.25
4
9
9
9
9
9
6
9
9
4
9
4
9
4
9
4
4
9
9
Linear model
9 9
G
9
9
9
2
9
9
9
4
4
4
4
4
4
4
4
4
4 4 4
Linear 0.55
9 9
9
9
9
9
4
9
9
9
4
4
4
4
+
4
4
+ 4
Linear model 0.75
9 9
9
8
9
9
9
9
9
9
U
4
4
4
Linear model t= 0.95
9
7
8
8
9
9
8
8 9
4
4
4
Figure S8: Transmmttation in (label-guided) linear model.. Transmutation is most naturally performed in branched diffuusion model, where branch points have been explicitly defined and trained with. However, possible to perform transmutation-like procedure in a label-guided (linear) model by first forward diffusing to some intermediate point (which we will call "turn-back point") and then reverse-diffusing from that point while supplying desired target label. a) We show a random sample of 4s and 9s from MNIST, and perform transmutation using branched diffusion model, with a well-defined branch point. The resulting digits show that transmutation the branched model was both efficacious (i.e. the digit successfully transformed from class the other) and analogous (i.e. non-class- features like slantedness were preserved). We then attempted trausmutation in a linear model, using various turu-back points. Notably, in linear model, the appropriate turn-back point is known ahead of time (unlike in branched model). As such, it is difficult to select the optimal turn-back point. Turn-back points which are too early cause transmutation to fail efficacy: the target class is not generated all. Turn-back points which are too late transmutation to fail at analogy: non-class-peecffic features like slantedness are no longer preserved. Additionally, some turn-bacd points (e.g. t= 0.25) cause some objects to be transmuted efficaciously, and others fail to generate the target class entirely. Furthermore, even by setting the turn-back point to be the branch point in the branched model (f = 0.35), transmuted results from the linear model 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.
a)
b)
Hybrid averages 6-9
A-W Feature 15
V-M
1-0
Feature 14
Feature 9
1.0
A
M
0.8
0
0.8
Hybrid
Hybrid
Hybrid
0.6
1-8
0.4
0.2
-2
D
-4 -2
2
-2
Figure S9: Interpreting hybrids of less-related classes. Branch points between less-related classes can still interpreted, but between very unrelated classes, the interpretations are naturally less meaningful, only very high-level (and relatively uninformative) features will shared between these classes. a) We show the distribution of feature values for highly distinct features between less similar letters. The hybrids show feature-value distribution which is intermediate and interpolated between the two classes, exhibiting properties of both. b) We show the hybrid resulting from interpreting the branch point between 6s and 9s between is and 8s (dissimilar digits which selected due to their late branch paints) from our branched model trained on MNIST. The resulting hybrids show the common features between each pair of digit classes, which are generally fairly high-level, including: the digits in the center of the image; and 2) areas which are generally empty, coinciding with the "holes" of how most people draw their 6s or 9s, or their 8s.
31

---

Published in Transactions on Machine Learning Research (08/2024)
a)
Fréchet inception distance between true and generated samples b)
Branched
300
Hybrids at branch points
Label-guided (linear)
Coat and Shirt
Sandal and Sneaker
150
100
2
Fashion MNIST class
C)
Shirt (original)
Coat (generated)
Figure S10: Branched models on Fashion MNIST. In addition to MNIST, we also show additional results on branched diffusion models trained more complex image dataset, but which does suffer from the problem of centering. a) We found significantly better performance in all classes (except for one) compared to the linear model, and this difference was oftentimes stark. b) We show hybrids between certain classes, which show shared shapes between certain types of clothing. c) Transmuting between shirts and coats, certain analogous features are preserved, such as overall color (light vs dark) and relative sleeve length (shorter vs longer).
Distance between classes at branch point
4 vS 4
Vs
4 VS 9
10
Frequency
8
6
4
0.25
0.30
0.35
0.40
0.45
Euclidean distance
Figure S11: Distribution of distance between objects branch point. We verify that distances between noisy objects of the same class are roughly the same compared to distances between noisy objects different classes at the branch points. Here, we show the distance between noisy 4s, between noisy 9s, and between noisy 4s and 9s the branch point. The distributions are almost identical, signifying that it takes no more "work" or "effort. for the diffusion model generate various objects of the same class (e.g. or 9s) compared objects different classes (e.g. The observation that these distributions match is a direct consequence of how branch points are defined (Equation S).
9s).
32

---

Published in Transactions on Machine Learning Research (08/2024)
D
Supplementary Methods
We trained all of our models and performed all analyses single Nvidia Quadro P6000.
D.1 Training data
We downloaded the MNIST dataset and used all digits from http://yana lecun. com/exdb/miist/ (LeCun et a..). We rescaled and recentered the values from [0, 256) to [-1, 1). This rescaling and symmetrization about 0 were to assist in the forward-disfusion process, which adds noise until the distribution approaches a standard (0-mean, identity-covariance) Gaussian.
We used the Fashion MNIST dataset as loaded from TorchVision.
We downloaded the tabular letter-recognition dataset from the UCI https://archive ics. www.uciic.. www...c.mmmm/a//ec... /ttta//////aaatteeeetttttttooooo tered and scaled each of the 16 tabular features to zero mean and unit variance (pooled over the entire dataset, not. for each individual letter class).
repository: We cen-
(Frey & Slate 1991.
We downloaded the single-cell RNA-seq dataset from GEO (GSE149689) Lee et al. 2020 We used Scanpy to pre-process the data, using standard workflow which consisted of filtering out low-cell-count genes and low-gene-counl cells, filtering out cells with too many mitochondrial genes, and retaining only the most variable genes and known marker genes (Wolf et al., 2018). We assigned cell-type labels using CeilTypist. [Conde et al. 2022). Of the annotated cell types, retained 9 non-redundant cell types for training: CD16+ NK cells, classical monocytes, late erythrocytes, macrophages, megakaryocytes/ platteless memory B cells, NK cells, plasmablasts, and TFM//ffector helper cells. After pre-processing, the dataset consisted of 37102 cells (i.e. data points) and 280 genes (i.e. features). To train our diffusson models, we projected the gene expressions down to a latent space of 200 dimensions, using the linearly decoded variational autoencoder in scVI (Gayoso et al., 2022). The autoencoder was trained for 500 epochs, with a learning rate of 0.005.
We downloaded the ZINC250K dataset and converted the SMILES strings into molecular graphs using RDKit. We kekulized the graphs and featurized according to (Jo el al. 2022). We explored two methods of labeling the molecules for branched diffusionn. First, we labeled molecules based on whether they were acyelic or had one cycle (molccules with multiple cycles were removed for simplicity). Secondly, we labeled molocules basod on whether or not they possessed a halogen element (i.e. F, CI, Br, I).
D.2 Diffusion processes
For all of our contimuous-time diffusion models, we employed the Evariance-preserviin stochastic differential equation" (VP-SDE) (Song et al., 2021). We used a variance schedule of ((t) 0.9t + 0.1. We set time horizon T (i.e. E [0, 1)). This amounts to adding Gaussian noise continuous time. Our ZINC250K models were an exception, and we used the same diffusion processes (different for the node features and adjacency matrix) that were used in (Jo et al.) 2022).
For our discrete-time diffusion model, we defined a discrete-time Gaussian noising process, following Ho et al. (2200) We defined B (1 10-4) + (1 x 10-5),. We set our time horizon 1000 (i.c. t [0, 1000)).
D.3 Defining branches
To discover branch points, we applied our branch-point discovery algorithm (Appendix Ab.
For continuous-time branched model on MNIST (and Fashion MNIST), used 0.005. For discrete-time branched model on MNIST, we used = 0.001. For our continuous-time branched model on tabular letters, we used = 0.01. For our single-cell RN A-seq dalaset, we used = 0.005. These values were selected such that the branch points were not all too close to or T.
The final branch definitions can be found in the Supplementary Figures and Tables. For our ZINC250K dataset, always used branch point of 0.15.
33

---

Published in Transactions on Machine Learning Research (08/2024)
D.4 Model architectures
Our model architoctures are designed after the architectures presented in Song ct al. (2021) and Kotelnikov et al.
)0022).
Our MNIST and Fashion MNIST models were trained on a UNet architecture consisting of 4 downsampling and upsampling layers. In our branched models, the upsampling layers were shared between output tasks. We performed group normalization after every layer. The time embedding was computed as [sin(2 / cos(2 ]] For each layer in the UNet, the time embedding was passed through a separate dense layer (unique for every UNet layer) and concatenated with the input to the UNet layer. For a label-guided model, we learned an embedding for each discrete label. As with the time embedding, the label embedding was passed through a separate dense layer (unique for every UNet layer) and concatenated to the input to each UNet. layer.
Our letter models were trained on a dense architecture consisting of 5 dense layers. In our branched models, the first two layers were shared between output tasks. The time embedding was computed (sim(2 (2 =)], where is set of Gaussian parameters that are not trainable. The time embeddings were passed through dense layer, and the output was added to the input after the first dense layer. For a label-guided model, we again learned an embedding for each discrele label. The label embedding was passed through a dense layer, and the result was concatenated to the summation of the time embedding and the input after the first. layer, before being passed to the remaining 4 layers.
Our single-cell RNA-sey models were trained on a dense residual architecture consisting of 5 dense layers. Each dense layer has 8192 hidden units. The input to each dense layer consisted of the sum of all previous layers' outputs. In our branched models, the first 3 layers were shared between output tasks. The time embedding was computed as [sin(2x y 2), cos (2 (2), where 2 is a set of Gaussian parameters that are not trainable. At each layer (other than the last), the time embedding was passed through a layer-specific dense mapping and added to the input to that layer. For a label-guided model, we again learned an embedding for each discrete label. The label embedding was passed through a dense layer and added to the very first layer's input.
Our ZINC250K branched models were trained on architecture almost identical to that presented in Jo et al. (022). In order to multi-task the model, we duplicate the last layers of the score networks for the node features or the adjacency matrix. For the node-feature score network, we multi-tasked only the final MLP layers. For the adjacency-matrix score network, we multi-tasked the final Attention ionLayer and the final MLP layers. To incorporate i, we computed a time embedding as [sin(2x cos 22),, where z is a set of Gaussian parameters that are not trainable. This embedding was passed through dense layer which projected this embedding down to a scalar, which was directly multiplied onto the final output of the score networks.
These neural-network architectures described above are the standard architectures used for linear (traditional) diffusion models for the associated data type (e.g. images, tabular data, or molecules). In order to turn these neural network architectures into multi-task architectures for branched diffusion models, we converted each model by multiplexing the last few layers (literally copying the last few layers 2|C|- - times for |C| classes), thereby turning the model into a multi-task model. The exact number of multiplexed layers multiplexed naturally depends on the data type and architecture, and is discussed below. This multiplexing was done such that the prediction path from the input to any single output head is architecturally identical to the standard single-task architecture.
For MNIST and Fashion MNIST, we multiplexed the last 4 layers of the U-Net (i.e. the upsampling layers). Note that because the U-Net concatenates downsampling (early) layers' outputs to the upsampler inputs, we effectively duplicate the downsamplers' outputs multiple times, as well, food into each output head. For letters, we multiplexed the last 2 layers. For RNA-seq, we multiplexed the last 3 layers. For molecules, we multiplexed the last, 4 layers for the adjacency-matrix score, and the last. 2 layers for the node-features score.
34

---

Published in Transactions on Machine Learning Research (08/2024)
D.5 Training schedules
For all of our models, we trained with a batch size of 128 examples, drawing uniformly from the entire dataset. This naturally ensures that branches which are longer (i.c. take up more diffusion time) or are responsible for more classes are upweighted appropriately.
For all of our models, we used a learning rate of 0.001, and trained our models until the loss had converged.
For our label-guided MNIST model, we trained for 30 epochs. For our label-guided letter model, we trained for 100 epochs. For our label-guided single-cell RNA-seq model, we trained for 100 epochs. In all cases, we noted that the loss had converged after training.
For our branched continuous-tim MNIST model, we trained for 90 epochs. For our branched discrete-time MNTST model, we trained for 200 epochs. For branched letter model (continuous-time), we trained for 100 epochs. For our branched single-cell RNA-seq model, trained for 120 epochs. For our branched ZINC250K model labeled by cyclicity, we trained for 200 epochs; for our branched ZINC250K model labeled by halogenation, we trained for 50 epochs. Again, we noted that the loss had converged after training, Our branched model on Fashion MNIST used the same training procodure as with MNIST.
For our analysis on extending branched models and label-guided (linear) models to new classes, we also trained MNIST models on a subset of the dataset (i.e. only 0/4/9 or only 0/4/7/9), and single-cell RNA--sq models on a subset of the dataset (i.e. only CD16+ /K/classical monocytes or only CD16+ NK/classical semnocytes/memon B cells). In these cases, we followed the same training parameters above, except we trained for fewer epochs. In the class-extension analysis on MNIST, we started with branched or label-guided models trained on Os, 4s, and 9s. These models we trained for 30 epochs each. In the class-extension analysis on single-cell RNA-seq, we started with a branched or label-guided model trained on CD16+ NK and classical monocytes. The branched model was trained for 120 epochs, and the tabel-guided model was trained for 100 epochs.
D.6 Sampling procedure
When generating samples from a continuous-ting time diffusion model, we used the predictor-correet algorithm defined in (ong et al. (2021), using 1000 time steps from to O. For our discrete-time diffusion model, we used the sampling algorithm defined in Нo al. (2020). Note that we employed Algorithm 2 for branched models.
D.7 Analyses
Sample quality
We compared the quality of samples generated from our branched diffuusin models to those generated by our label-guided (linear) diffusion models using Fréchet Inception Distance (FID). For each class, generated 1000 samples of each class from the branched model, 1000 samples of each class from the linear model, and randomly selected 1000 samples of each class from the true dataset. We computed FID over these samples, comparing each set of generated classes against the true samples. For the tabular letters dataset, there were not enough letters in the dataset to draw 1000 true samples of each letter, so we drew 700 of each letter from the true datasct. For the single-cell RNA-seq dataset, we generated 500 each cell type from the diffusion models, and we sampled many of each cell as possible from the true dataset (up to maximum of 500).
Class extension
For our MNIST dataset, we started with a branched diffusion model trained on Os, 4s, and 9s. To extend & new branch to reverse diffuse 7s, we simply created a new model with one extra output task and copied over the corresponding weights. For the new branch, we initialized the weights to the same as those on the corresponding branch (with s; = 0). We trained this new branch on only 7s, only for the time interval of that new branch. On the corresponding label-guided (linear) model, we fine-tuned on only or on Os, 4s, 7s, and 9s. In each experiment, we started with the linear model trained on 0s, 4s, and 9s.
35

---

Published in Transactions on Machine Learning Research (08/2024)
For our single-cell RNA-seq dataset, repeated the same procedure, but started with branched diffusion model trained on CD16+ NK cells and classical monocytes, and introduced memory B cells as a new task. We trained the new branch on only memory B cells, only for the time interval of the new branch. On the corresponding label-guided (lincar) model, we fine-tuned on only memory B cells or on all three cell types. To fine-tune on only memory B cells, we started with the linear model trained on CD16+ NK cells and classical monocytes only. To fine-tune on all three cell types, we again started with the linear model trained only on the original two cell types.
As above, always fine-tuned until the loss converged. We note that this took much longer for the linear models compared to the branched model.
Hybrid intermediates and transmutation
For certain pairs of MNIST digits or letters, we found the carliest branch point for which they belong the same branch, and generated hybrids by reverse diffusing to that branch point. To generate the average MNIST hybrids, we sampled 500 objects from the prior and reverse diffused to the branch point, and averaged the result.
Transmuted objects were computed by forward diffusing from one class to this branch point, and then reverse diffusing down the path to the other class from that intermediate.
To compute the preservation of functional groups in the transmutation of ZINC250K, we used the following list of functional groups:
Multi-class sampling efficiency
We computed the amount of time taken to generate 64 examples of each class from our branched diffusion models, with and without taking advantage of the branch points. We took the average time over 10 trials each.
When leveraging the branching structure to generate samples, we ordered the branches by start time s; in descending order. For each branch b; in that order, we reverse diffused down the branch, starting with cached intermediate at for the branch that ended at t,. For the very first branch (the root), we started reverse diffusion by sampling (Ix) This guarantees that we will have a cached batch of samples al every branch point before we encounter a branch that starts at that branch point. Eventually, this algorithm generates a batch of samples for each class. For cach branch, we performed reverse diffusion such that the total number of steps for any one class from ( = to = 0 was 1000.
To generate samples without leveraging the branching structure, we simply generated each class separately from the branched model, without caching any intermediates. Note that this takes the same amount of time purely linear model (of identical capacity and architecture) without any branching structure.
Robustness of branch points
To quantify the robustness of branched diffusion models to the underlying branch points, we computed the branch points for MNIST (continuous-time, all 10 digits) times, each time following the procedure Appendix A Variation in the branch points resulted from variation in the randomly sampled objects, and in the forward-diifusion process. For each set of branch definitions, we trained branched diffusson model using the procedure above. We then computed FID using the same procedure as above for other MNIST models, and compared the values to the FIDs of the corresponding label-guided MNIST model.
To quantify the similarity of the hierarchies, we computed the pairwise branch-score distance between all (20) pairs of hierarchies discovered from our algorithm. We then generated 10 random hierarchies in a greedy fashion: start with all classos, and uniformly pick a random partition; uniformly pick a branch point b, between 0 and 1; recursively generate the two hicrarchies below with a maximum time of b, until all class sets have been reduced to singletons. We used Wilcoxon test to compare the distribution of branch-score
36

---

Published in Transactions on Machine Learning Research (08/2024)
distances between the hierarchies discovered by proposed algorithm, and the distances between the random hierarchies. Branch-score distance was computed using PHYLIP (Felsenstein).
37


Reviews:
[{"review_id": "A6i2zqEhJh", "reviewer": "Reviewer_wkP3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The hierarchical branching of diffusion models is an interesting idea that the authors test in diverse scenarios: continual learning, interpretability, and interpolation.\n- Paper is clearly written and easy to read.\n- The appendix is full of additional experiments and details.\n- The supplementary material contains the code and several notebooks for reproducing the results.\n\nWeaknesses:\n\n- Evaluation is insufficient.\n    - The datasets used in this paper are fairly simple. Diffusion models are primarily used with images, and the image experiments here are run only on MNIST, which is a trivial dataset. It is not apparent if this approach scales up to Cifar or CelebA. The ZINC250K and RNA-seq datasets are nice additions, but I still think an extra image dataset would be useful.\n    - The baselines are not present or too weak. I don’t think we require the method to outperform all prior work but having some solid results for reference would be useful.\n        - In continual learning experiments, the architecture for branched diffusion consists of a shared encoder and a single head per branch. Only the heads are trained on new tasks, the encoder is frozen. On the other hand, it seems like label-guided diffusion trains the whole model without any freezing. In the label-guided setting, one still could have a bunch of shared parameters followed by a single head per label — I would expect this to work much better.\n        - In interpolation experiments, it would be nice to have some baselines. Interpolation with diffusion models is not the most trivial thing, but there have been several approaches for this purpose. [1, 2]\n        - In the interpretability experiments, there is no reference method. I don't have any great suggestions here, but one could, for example, use the label-conditioned diffusion with the label set at `[0.5, 0.5]` between relevant classes.\n- The technique for branching point identification is not under empirical scrutiny here. The authors describe the theoretical underlying of the approach in Appendix A, but there aren’t any experiments to check if the identified branching points $t_b$ work for all examples for a given task. I would imagine that for certain examples (noise samples), the class becomes visible much more quickly than for others, and even more so in complex images.\n    - For example, if we were to apply the branching to Cifar-10 and find the branching point between classes “bird” and “airplane” then examples with blue background (might be either a sky or a bird against a sky) would branch much later than examples with a green background (probably a bird in a forest).\n\n\n[1] Wang, Clinton, and Polina Golland. \"Interpolating between images with diffusion models.\" (2023). \\\n[2] Jain, Siddhant, et al. \"Video Interpolation with Diffusion Models.\" arXiv preprint arXiv:2404.01203 (2024).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces the idea of class-conditional hierarchical branching in diffusion models. Namely, there is a point in the diffusion generation process when the class of the example is decided, which we call the branching point. The authors try to find the average such point for each branch by considering when the objects of different classes are as different as objects of the same class. The authors then train such hierarchical branching diffusion models and show their application in different problems, in particular, continual learning, analogy-based conditional generation, and interpretability.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address the issues with the evaluation and identification of the branching point described in the weaknesses section. I would say that having a better baseline in the CL setting is of critical importance and the discussion of how well we can identify the branching point is very much needed. The other points are nice to have but not that critical, in my opinion.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "4ibwSN3uUq", "reviewer": "Reviewer_nEk7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n(1)\tThe model proposed in the paper is novel and is an elegant and intuitively-appealing strategy for performing class-conditional sampling using a diffusion model. \n\n(2)\tThe experiments use datasets from multiple modalities and illustrate the potential of the proposed approach. In addition to achieving good generative performance, the method has appealing capabilities to incorporate new classes or perform transmutation. \n\n(3)\tThe paper is well-written and describes the technique and experiments clearly (although there are some structural issues and some important information is excluded from the main body). \n\nWeaknesses\n\n(1)\tThe model relies on the selection of good branchpoints (although the paper does include some results exploring robustness). The paper presents the model as though the branchpoints are provided as input, but based on all of the presented experiments, it appears that the branchpoints are actually parameters of the model. All of the results rely on a procedure to choose these branchpoints, but that procedure is not discussed in the main text of the paper. \n\n(2)\tThe branch selection procedure incorporates design choices that are not particularly well-motivated (e.g., linear classifier, Euclidean distance). \n\n(3)\tSome of the structural choices of the paper could be improved with a view to facilitating comprehension for a reader.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a novel diffusional model that incorporates branching to facilitate class-conditional sampling. The paper provides brief algorithms for sampling and training (assuming that the branch structure is known). A heuristic for identifying the branch structure is provided in an appendix. The paper presents experimental results for multiple datasets of different modalities. First, there is a demonstration that the proposed branched model can achieve generative performance that is equivalent to (or perhaps even slightly better than) a traditional model. Second, the paper illustrates how the model can incorporate new classes without requiring complete re-training. Third, the paper provides examples of conditional generation via transmutation. Finally, the paper discusses interpretability of the branchpoints and the potential for efficient multi-class sampling.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "(1)\t[Critical] The main body of the paper should include some discussion of the choice of branch points in the methodology section. It is fine to leave the details to an appendix, but there should be acknowledgement in the methodology that this is a critical component of using such a model (and underpins all of the presented results. \n\n(2)\t[Critical] The decision to base the branching point on indistinguishability according to a linear classifier appears odd. This would seem to be a low bar in the sense that the distributions could still be structurally very different (e.g, circle enclosed by a ring) and yet the linear classifier would struggle. Please add some text that provides a better motivation for this design decision. In particular, why is it not possible to use more standard approaches for testing whether two distributions are equivalent? (e.g., something like an MMD test, such as the one in Schrab et al., “MMD Aggregated Two-Sample Test”, JMLR 2023). Is there a reason that the diffusion process means that the classifier (Euclidean distance) is a sensible strategy? \n\n(3)\t[Recommended, not mandatory] The paper includes experimental results in Section 3.1, a subsection of Section 3, which describes the model and algorithms. I think the paper would be easier to digest if the reader can separate the methodology from the experimental work, and I would recommend moving Section 3.1 to Section 4. Sections 3.1, 4 and 5 could be subsections of an “Applications and Experiments” section. Section 7 is very brief and contains very little information. If it is not going to present any results, or even a summary of results, then it should be included as a paragraph in the Discussion section.", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "SdheQmx0vw", "reviewer": "Reviewer_GyoJ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\\+ conceptually nice motivation\n\n\\+ diverse-ish experiments\n\n\\+ ablations over branch points\n\n\\++ class extensions and analogies are new capabilities\n\n0 limited discussion of computational intensity\n\n0 foundation of branch point a bit ad-hoc\n\n\\- model details of implementation of heads not clear from text\n\n\\- evaluation on low dimensional data mainly, covid data only \"real\" task and limited performance there\n\n\\-- possible inconsistency: main body talks about euclidean distance(in latent space), appendix talks about latent space  of a classifier *AND* the margin of prediction.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper outlines an approach to conditional modeling of a dataset with $\\vert C \\vert$ labels in which the similarity between instances is used to inform a specialization into  $2\\vert C \\vert -1$ heads,\neach head being responsible for a branch in a tree, with he root being the single shared \"trunk\" near the noise limit and lower levels of the tree being incrementally more specialized diffusion processes.\nThis tree is constructed via similarities of the $\\ell_2$ distance of an auxiliary classifier by defining the branch point as the point where inter-class and intra-class similarity are matched in expectation and using a monte-carlo estimate to compute the branch point.\n\nThis architecture is claimed to enable class extension (add a new class into the diffusion process without requiring retraining on *all* data) while adding a new head (assuming the new class is sufficiently similar to existing classes), analogy generation (diffusion until the branch point, then reverse diffusion into a new class) and interpretable hybrids at the branch points, yielding insights into class feature structure.\n\nExperiments are performed on MNIST, a tabular letter dataset, a single-cell RNA-seq gene expression dataset and  ZINC250K, with ablations being performed to the specific branch point timing", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "RFC\n\n\nfor the appendix\n\n1. evaluate your branch point estimation by analysing whether a https://en.wikipedia.org/wiki/Welch%27s_t-test or another suitable statistical test would yield the same branch point method (and its error rate=\n2. clarify the architecture used and the parameter count of vanilla vs your method. Also, report compute and timings of all methods in appendix\n\nfor the main body and/or appendix\n\n1. I am concerned about the lack of improvement on the COVID task, can you run this on e.g. fashion MNIT, cifar10 and ciaf100 and resized imagenet-1k to ablate over increasing complexity and number of classes, and also investigate scalability? If not possible to perform, at least analyze the complexity of this method\n2. critical: clarify whether the branch point computation is done with margin of probability, l2 or latent space l2\n3. an ablation over class extension with a hold-out class of varying similarity ot existing classes would be an interesting nice to have (i.e., add a very dissimilar class vs. a very similar class to compare corners of the process).", "broader_impact_concerns": "I think adding BICs is not required", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Two reviewers recommended acceptance, one reviewer weakly opposed acceptance. The submission satisfies the TMLR criteria for acceptance.
Decision: Accept (poster)