ComboStoc: Combinatorial Stochasticity for Diffusion Generative Models

## Abstract

Abstract In this paper, we study an under-explored but important factor of diffusion generative models, i.e. , the combinatorial complexity.
Data samples are generally high-dimensional, and for various structured generation tasks, there are additional attributes which are combined to associate with data samples.
We show that the space spanned by the combination of dimensions and attributes is insufficiently sampled by existing training scheme of diffusion generative models, causing degraded test time performance.
We present a simple fix to this problem by constructing stochastic processes that fully exploit the combinatorial structures, hence the name ComboStoc .
Using this simple strategy, we show that network training is significantly accelerated across diverse data modalities, including images and 3D structured shapes.
Moreover, ComboStoc enables a new way of test time generation which uses asynchronous time steps for different dimensions and attributes, thus allowing for varying degrees of control over them.

## 1 Introduction

Diffusion generative models rely heavily on modeling the desired behavior over the whole space of possibilities, so that the generative models cover all data distributions systematically.
However, the current training schemes generally focus on a single transport path from the source pure noise distribution to the target data distribution [, , , ].
The training therefore gives insufficient sampling of large regions of the entire space of possibilities, which nevertheless can be stamped upon because of stochastic sampling during evaluation and produce inaccurate behavior, leading to poor generation results.

To solve this mismatch between training scheme and test time evaluation, we propose to fully sample the space of combinatorial complexity.
To see why the space of possibilities has a combinatorial structure, we note that the data samples are most likely residing on high dimensional spaces with clear combinatorial structures.
For example, the most powerful generative models so far use transformers as the network model [, ], and treat an image sample as a collection of patch tokens to be generated in parallel; moreover, each of the patch token is encoded as a vector of high dimensions.
The combination of patches and their feature vectors present highly complex spaces, over which the diffusion generative models must learn to evolve toward data samples where patches and feature vectors are correlated nontrivially.
In addition, for generative tasks in more structured domains, for example, 3D shapes with semantic parts, the combinatorial complexity is even more pronounced: each part has different attributes encoding properties like its existence, bounding box and part shape, in addition to the part/patch decomposition and multiple feature channels analogous to images.

**Table 1: Improvements over SiT across iterations.**
| Model | Params(M) | Training Steps | FID |
| --- | --- | --- | --- |
| DiT-XL | 675 | 400K | 19.5 |
| SiT-XL | 675 | 400K | 17.2 |
| ComboStoc | 673 | 400K | 15.69 |
| DiT-XL | 675 | 800K | 14.3 |
| SiT-XL | 675 | 800K | 12.6 |
| ComboStoc | 673 | 800K | 11.41 |
| DiT-XL (cfg=1.5) | 675 | 7M | 2.27 |
| SiT-XL (cfg=1.5) | 675 | 7M | 2.06 |
| ComboStoc (cfg=1.5) | 673 | 800K | 2.85 |

We sample the spaces of such combinatorial complexity by a simple modification of typical transport plans.
In particular, instead of using a synchronized time schedule for each data sample, we apply asynchronous time steps for each of the patches/parts, attributes and feature vector dimensions, which allows for full sampling of a subspace spanning the various combinations of each pair of source and target data points.

We show that by simply enhancing the training scheme to incorporate the combinatorial sampling, the generative models for images and 3D structured shapes can be significantly improved (Fig. 1).
In particular, for images from ImageNet, we obtain systematic FID-50k improvements along different training iterations than the baseline SiT model (Tab. 1).
For 3D structured shapes which have even stronger combinatorial complexity, we show that our training scheme is indispensable for obtaining a working generative model.

In addition to the improved performances, the training scheme exploiting combinatorial stochasticity enables new modes of using the trained generative models.
Specifically, we can now generate different patches/parts/attributes in asynchronous time schedules.
This means that for example we can condition the final sample on flexible partial observations of a reference sample beyond fixed modes like binary masks.
Instead, for images we can apply graded conditioning across patches and channels.
For structured shapes we can also specify the shapes of some parts only, and let the model generate the rest parts and attributes.
These new modes of generation have the potential to unify specialized image and shape editing solutions.

## 2 Background on diffusion generative models

The problem of generative modeling aims at capturing the complete distribution of a set of data samples.
Its state-of-the-art solutions include denoising diffusion probabilistic models [], score-based models [] and flow matching [, ], all of which transform a simple source distribution (e.g. the unit normal distribution) into the target distribution following the dynamics specified by variations of stochastic differential equations.
Remarkably, the different formulations can be unified through the framework of stochastic interpolants [, ].
In particular, the stochastic interpolants framework defines the processes of turning data samples into source distributions and vice versa as a simple interpolation between the two distributions, augmented with random perturbations during the processes.
Without loss of generality, we reproduce the formulation of a simple linear one-sided interpolant process below:

$$ $\displaystyle\mathbf{x}_{t}=(1-t)\mathbf{z}+t\mathbf{x}_{1},\ t\in[0,1]$ (1) $$

where $\mathbf{z}\sim N(0,\mathbf{1})$ samples the source distribution, $\mathbf{x}_{1}\sim D$ samples the target data distribution, $t\in[0,1]$ is the interpolation schedule.
The network model $f_{\theta}(\mathbf{x}_{t},t)$ can be trained to recover the interpolation velocity $\frac{\partial\mathbf{x}_{t}}{\partial t}=\mathbf{x}_{1}-\mathbf{z}$, the target data sample $\mathbf{x}_{1}$, or the noise $\mathbf{z}$ [].
To generate data samples, one then starts from a random sample $\mathbf{z}$, follows the velocity field and integrates them numerically into the final samples.

Remarkably, on modeling large scale image datasets like ImageNet, a scalable transformer architecture implementing the above process [] shows state-of-the-art performance and outperforms alternative formulations, including DDPM [] implemented via the same network [].

## 3 Combinatorial stochastic process

Figure: Figure 2: ComboStoc enables better coverage of the whole path space. Assuming two-dimensional data samples. (a) the standard linear one-sided interpolant model reduces its density as it approaches individual data samples; the low density regions are not well trained and once sampled would produce low-quality predictions. (b) using ComboStoc, for each pair of source and target sample points, a whole linear subspace spanned with their connection as the diagonal will be sufficiently sampled, so that there are fewer low-density regions not well trained. (c) when the network is trained to predict velocity $\mathbf{x}_{1}-\mathbf{z}$, on an off-diagonal sample point $\mathbf{x}_{\mathbf{t}}$, a compensation drift ($\mathbf{v}_{cmpn}$, green vector) is needed to pull the trajectory back to diagonal.

Most interesting data samples are high dimensional. For example, state-of-the-art generative models encode images as latent patches with both spatial and feature dimensions [, ].
3D shapes structured as part ensembles include even more attributes in addition to spatial and feature dimensions, for example, the varying numbers of parts and their bounding boxes and positions []; to generate such data requires the handling of more flexible dimensions.

No matter how many dimensions and attributes a data sample has, standard diffusion generative models treat them homogeneously and in synchronization.
For example, for the stochastic interpolants, the generative model is trained on samples distributed on densities with shrinking coverage along the transport paths connecting the source distribution and each target data sample, as illustrated in Fig. 2(a).
The problem with this homogeneous and synchronized sampling scheme, however, is that there are low-density regions closer to the data samples that are insufficiently trained, and once they are sampled in test stage via solving stochastic differential equations, the network tends to produce poor results.

To address the above problem, we emphasize the combinatorial complexity of individual dimensions and attributes of data samples.
In particular, we purposefully sample points with asynchronous diffusion schedules of dimensions and attributes, as shown in Fig. 2(b).
To implement the asynchronous schedules is rather simple.
We turn the interpolation schedule $t$ of Eq. (1) into a tensor of the same shape as $\mathbf{x}$, and use different values within $[0,1]$ for different dimensions or attributes:

$$ $\mathbf{x}_{\mathbf{t}}=(1-\mathbf{t})\odot\mathbf{z}+\mathbf{t}\odot\mathbf{x }_{1}$ (2) $$

Correspondingly, the network $f_{\theta}(\mathbf{x}_{\mathbf{t}},\mathbf{t})$ is conditioned on the tensorized schedule $\mathbf{t}$ and trained to predict velocity, or the target data sample, etc.

The benefits of using these augmented samples from combinatorial stochasticity lie in three folds:

- •
We make sure the network coverage is broader than the synchronized schedule, so that during test stage the network performs more robustly and with higher quality.
- •
We encourage the network to learn the correlations of different dimensions and attributes, as the network is trained to synchronize them to reach the final data points.
- •
The trained network enables more flexible conditioning, where different dimensions and attributes can be given varying degrees of finalization as specified by their interpolation schedules.

Next, we discuss the detailed adaptations and achieved effects through generative tasks from two different domains, i.e. images and structured 3D shapes.

#### Images.

For image generation, we take on the baseline of SiT [] which applies highly scalable transformer networks and achieves state-of-the-art performance on ImageNet scale generation.
In particular, a given image is encoded via the VAE encoder from [] as a latent image $\mathbf{x}_{1}$ of shape $C{\times}H{\times}W$, and the network is trained to predict velocity given the diffused latent image $\mathbf{x}_{t}$ (Eq. (1)) and conditioned on the interpolation schedule $t$ and image class $c$, i.e. $f_{\theta}(\mathbf{x}_{t},t,c)=\mathbf{x}_{1}-\mathbf{z}$.

Correspondingly, we make several simple adaptations to implement the ComboStoc scheme.
First, we construct $\mathbf{t}$ with the same shape of $C{\times}H{\times}W$, and update the time-step embedding module of SiT to accommodate this change (see Appendix A).
Second, importantly, we note that for velocity prediction, the samples with asynchronous $\mathbf{t}$ should not only predict the original velocity $\mathbf{x}_{1}-\mathbf{z}$, but also incorporate a compensating drift that pulls the trajectory back to the diagonal path connecting the source and target points with synchronized schedules, as illustrated in Fig. 2(c).
Without this off-diagonal compensation, an integration of the velocity field would lead to divergence from the target sample once any off-diagonal point is encountered.
The compensation drift $\mathbf{v}_{cmpn}$ is simply obtained as the projection of the target drift $\mathbf{x}_{1}-\mathbf{x}_{\mathbf{t}}$ orthogonal to diagonal, i.e.

$$ $\mathbf{v}_{cmpn}(\mathbf{t})=\mathbf{x}_{1}-\mathbf{x}_{\mathbf{t}}-\frac{( \mathbf{x}_{1}-\mathbf{x}_{\mathbf{t}})\cdot(\mathbf{x}_{1}-\mathbf{z})}{|| \mathbf{x}_{1}-\mathbf{z}||^{2}}(\mathbf{x}_{1}-\mathbf{z})$ (3) $$

Thus the total velocity becomes $\mathbf{v}(\mathbf{t})=\mathbf{x}_{1}-\mathbf{z}+\mathbf{v}_{cmpn}(\mathbf{t})$. Note that for on-diagonal points, the compensation drift vanishes as expected.

#### Structured 3D shapes.

We use the generative modeling of structured 3D shapes as a new task to further demonstrate the importance of exploiting combinatorial complexity.
Indeed, structure 3D shapes have even stronger combinatorial complexity than images, as shown in its varying numbers of parts, their positions and bounding boxes, as well as the detailed shape variations for each part.
Precisely, we denote a structured 3D object as a collection of object parts, i.e., $\mathbf{x}=\{\mathbf{p}_{i}\},i\in[L]$, where we set $L=256$ to cover the maximum number of parts in a dataset.
An object part is further encoded as $\mathbf{p}=(s,\mathbf{b},\mathbf{e})$, where $s\in[0,1]$ indicates the existence of this part, $\mathbf{b}=(x,y,z,l,w,h)$ denotes the bounding box center $(x,y,z)$ and length $l$, width $w$ and height $h$, and $\mathbf{e}\in\mathbb{R}^{512}$ is a latent shape code encoding the part shape in normalized coordinates.
Note that under this representation, a permutation of the part indices does not change the 3D shape, which is quite different from images represented as a feature grid of fixed order and size.

To generate structured 3D shapes with semantic parts, we train a stochastic interpolant model.
In particular, given a structured 3D shape $\mathbf{x}_{1}=\{\mathbf{p}_{i}\}$ and its diffused sample $\mathbf{x}_{\mathbf{t}}$ (Eq. (2)), we make the network predict the target data sample directly for simplicity, i.e., $f_{\theta}(\mathbf{x}_{\mathbf{t}},\mathbf{t},c)=\mathbf{x}_{1}$, where $c$ is the class condition of the 3D shape.
Note that here $\mathbf{t}$ assigns different time steps for all the different attributes and dimensions of each object part.
We validate the generative model for structured 3D shapes by training on the PartNet dataset [], as discussed in Sec. 4.

## 4 Results and discussion

In this part, we show that ComboStoc improves the training convergence of diffusion generative models for both images and structured 3D shapes.
We also demonstrate the novel applications enabled by the asynchronous time steps of ComboStoc.

**(a) Image domain configurations**
| Feature<br>Spatial | w/o patch | w/ patch |
| --- | --- | --- |
| single code | insync_none | insync_patch |
| feature vector | insync_vec | insync_all |

### 4.1 Improved training of diffusion generative models

We explore the combinatorial complexities for both images and structured 3D shapes, and build corresponding configurations which exploit these complexities to compare with baseline configurations that do not apply asynchronous time schedules.
We show that the different configurations improve over baseline configurations universally.
In addition, we show that the higher degree of combinatorial complexity, the more important our scheme is for training a working model.

Figure: Figure 3: Comparison on image generation with respect to training steps. (a) plots the baseline SiT and our model, as well as DiT for reference; all models are of the scale XL/2 []. (b) plots the different settings using varying degrees of combinatorial stochasticity.

Figure: Figure 4: Results of image generation at different training steps. Settings with stronger combinatorial sampling produce well-structured images earlier; e.g. see the koala bear faces and cat eyes.

#### Images

Following SiT [], we train on ImageNet [] for class-conditioned image generation.
To fully explore the effects of combinatorial stochasticity, we enumerate four settings with different levels of combinatorial flexibility in diffusing the data samples (see also Tab. 2(a)).
In particular, we use insync_none, insync_patch, insync_vec, and insync_all to denote no splitting of time steps, using different time steps for latent image pixels, for latent image channels, and for both image pixels and channels.
We run the different settings on top of the SiT-XL/2 baseline model.
Considering the difficulty posed by ImageNet data size, in each batch we apply the split time steps only to half of the samples and leave the other half unchanged with synchronized time steps, which balances between samples along and off diagonal paths (Fig. 2)(^1^11This batch mixing scheme may not be optimal. In preliminary tests (A) we found that blending the split time steps with synchronized ones gives even better results. The search of optimal sampling is left for future work.).
Plots of FID-50K [] with respect to training steps are shown in Fig. 3, where classifier-free guidance is not used.

First, as shown in Fig. 3(a), our scheme (using insync_all) shows consistent improvement of the baseline SiT model, and significant improvement over the reference DiT model.
Second, as shown in Fig. 3(b), the different settings of time step insynchronization behave differently.
Overall, the finest split by insync_all obtains the best performances consistently, followed by insync_vec and insync_patch which split along feature and spatial dimensions and have almost indistinguishable performances.
The worst performance is obtained by insync_none, i.e. the setting using no combinatorial stochasticity.
Fig. 4 visualizes the results of different settings along training steps, where we see better structured images emerge earlier for settings using stronger combinatorial complexity.
The comparison among these four settings shows that fully utilizing the combinatorial complexity indeed helps network training.

Note that due to the differences introduced in the time-step embedding module, insync_none has slightly worse performance than the baseline SiT, probably because the time-step encoding vector has a smaller size (see Appendix A).
While it may be possible to align insync_none with baseline SiT by introducing more capable embedding layers, insync_all already outperforms the baseline with significant margins (Fig. 3(a)).
More results by insync_all can be found in the Appendix A.

Figure: Figure 5: Results of structured shape generation by different settings. Semantic parts are colored randomly. Settings exploiting stronger combinatorial stochasticity show better results. In comparison, insync_none that does not apply ComboStoc nearly fails to generate meaningful shapes.

#### Structured 3D shapes

We show that for the task of structured 3D shape generation, which has even stronger combinatorial complexity due to the flexible parts and their multiple attributes, our scheme becomes more important to the extent of being indispensable.

For structured 3D shape generation, we identify combinatorial complexity in the following axes: attributes/feature vectors, and spatial parts.
Therefore, we obtain $3{\times}2=6$ settings, i.e., insync_none and insync_part which apply the same or different time schedules to parts respectively, insync_att and insync_att_part which use attribute level schedules, and insync_vec and insync_all which use the most finely divided feature vector level schedules.
See Tab. 2(b) for a summary of the 6 configurations.
Because of the relatively small size of the PartNet dataset (18K shapes in total, mostly in chair and table classes), we deem it easier to learn and simply apply the corresponding asynchronous time steps to all samples in each batch, in contrast to the mixing scheme of ImageNet training.
We report results at 1.5K epochs, since earlier results cannot be decoded into valid manifold shapes for evaluation in settings like insync_none.

As shown in Fig. 5, the more combinatorial complexity we exploit, the better performance of the trained network.
In comparison, the baseline setting without combinatorial stochasticity, insync_none, almost entirely fails to produce meaningful shapes.
Moreover, since this task models the highly flexible composition of various numbers of parts, applying the spatial part insynchronization (Tab. 2(b)) helps obviously, as shown through the three pairs of columns in Fig. 5 (e.g., part vs none, att_part vs att, and all vs vec).

We report quantitative results in Tab. 3 using the chair category. Following [] we use three metrics, including Frechet Point Distance (FPD) that measures the FID on sampled point clouds, coverage (COV) that measures how well each GT sample is covered by the closest generated sample, and minimum matching distance (MMD) that measures how well each generated sample resembles the closest GT sample.
The numerical results again show that the part level combinatorial stochasticity enhances generative performance significantly, and insync_all shows the best overall results.

In Fig. 6 and Appendix  A we show more random samples generated by the insync_all setting.
In Appendix A, we also compare with other works that generate structured shapes by taking a hierarchical refinement process, and find that our results are within their performance ranges.

**Table 3: Quantitative evaluation of structured shape generation by different settings. Chair category is used. Best scores are marked in bold and underlined; second best scores in bold.**
|  | insync_none | insync_part | insync_att | insync_att_part | insync_vec | insync_all |
| --- | --- | --- | --- | --- | --- | --- |
| FPD$\downarrow$ | 7.99 | 4.71 | 7.47 | 3.51 | 4.62 | 4.04 |
| COV$\downarrow$ | 1.32 | 1.03 | 1.83 | 0.85 | 0.97 | 0.86 |
| MMD$\downarrow$ | 1.23 | 1.95 | 1.38 | 1.04 | 0.63 | 0.68 |

Figure: Figure 6: Class-conditioned generation of structured 3D shapes. From top to bottom the classes are: chair, laptop, table and display.

### 4.2 Application enabled by combinatorial stochastic process

The asynchronous time steps for different dimensions and attributes of ComboStoc enables a novel test time application, namely the chance to specify different degrees of preservation of a data sample to its dimensions and attributes.
Specifically, given $\mathbf{t}_{0}$ specifying the weights in $[0,1]$ to preserve the data of $\mathbf{x}$, we sample the generative process starting from

$$ $\mathbf{x}_{0}=(1-\mathbf{t}_{0})\odot\mathbf{z}+\mathbf{t}_{0}\odot\mathbf{x},$ (4) $$

and increase the time steps for individual dimensions and attributes via $\frac{\mathbf{1}-\mathbf{t}_{0}}{N}$ for $N$ steps.
Examples of such asynchronous generative processes are shown in Figs. 7, 8 for images, and Figs. 9, 10 for structured 3D shapes.

#### Images

In Fig. 7 we show that giving different $\mathbf{t}_{0}$ to a half of a reference image while leaving the other half to generate from scratch, we can achieve different degrees of preservation of the reference images.
In particular, as the preservation weight increases from 0 0 to $1$, the preservation of reference content is strengthened.
We note that at 0.5 the weight is good enough to preserve most of the reference content.

In Fig. 8, we use different preservation weights encoded by $\mathbf{t}_{0}$ for the four quadrants of each image, and show intermediate results along the iterative SDE integration process.
From the three examples we can see that stronger weights cause better preservation of reference regions, and the different regions are filled with coherent content despite the spatially varying time schedules.
This mode of controlled generation is novel, compared with the binary inpainting mode proposed for standard diffusion models [], where regions of an image are divided into two discrete types, i.e. those to preserve and those to generate from scratch.

Additionally in Appendix A, we show that channel-varying $\mathbf{t}_{0}$ reveals interesting observations about the different contents of the latent image encoding [].

Figure: Figure 7: Image generation using different weights of preservation. Each reference image (right) is split into two vertical halves (left), and the left half is given the preservation weights while the right region starts from scratch.

Figure: Figure 8: Image generation with spatially different preservation weights. As shown in the left column, the four quadrants use $t_{0}=0,0.25,0.5,0.75$, respectively. The sampling iterations converge to results that preserve the corresponding quadrants from the reference images (right) differently.

#### Structured 3D shapes

By controlling different parts and attributes of structured 3D shapes, we can achieve diverse effects, including shape completion and part assembly.
In Fig. 9, we fix the bases of chairs by giving them $t_{0}=0.9$, and complete them with meaningful but diverse structures that satisfy the class condition.
The given bases have the chance of being slightly updated to adapt to the completed shapes.
In Fig. 10, we randomly position a set of parts, and let the network arrange them into proper shapes, by giving the part shape codes $\mathbf{e}$ and bounding box sizes large preservation weights ($t_{0}=0.9$) and making the rest attributes free to be generated.
Here we have considered a simplified setting where the part rotations are given, and would like to explore more challenging applications where shape parts can be rotated in the future.

Figure: Figure 9: Structured shape completion. Given base parts (left), the network can complete the missing parts conditioned on a shape category name (chair in this example). While the completed parts show great diversity, the given parts are preserved faithfully.

Figure: Figure 10: Assembly of semantic parts. Given parts in random positions (left and right), the network assembles them into complete shapes (middle). We solve this part-assembly problem via preserving the attributes of part shapes and scales and only generating the attribute of part positions.

## 5 Related works

For image generation, while there are numerous optimizations of training schemes, including training loss weights and time schedules [], speedup by distillation [], and sampling path consistency [], few have noticed the factor of combinatorial complexity.
An exception is [], which attributes the slow training of DDPM based DiT model [] to the pixel-wise regression loss, which does not emphasize the structural correlation of different patches sufficiently.
To address the problem, [] design a mask-and-diffusion scheme that masks out portions of the input diffused images during training to encourage learning the patch correlation, implemented by a complex encoder-decoder network with additional side-interpolation modules.
In comparison, our scheme is simple and requires minimum changes of baseline networks, but can already improve training of SiT models significantly.
Notably, SiT models already surpass DiT models in performance [].

While 3D diffusion generative models are increasing [, ], few works have been done for structured shape generation.
[] studies the representation learning of hierarchically structured shapes and proposes to generate variations using a VAE model.
Compared to [], [] proposes a rewriting model to enable generalizable cross-category generation.
In comparison, we focus on generating flatly structured 3D shapes with leaf level semantic parts.
Moreover, by specifying parts and attributes independently, our model enables diverse tasks like shape completion and assembly.
Previously, such diverse applications have been studied individually via specialized solutions [, ], but here we show they can be potentially unified by a single model generating highly structured data.

## 6 Conclusion

We have proposed to focus on the problem of the combinatorial complexity of high-dimensional and multi-attribute data samples for diffusion-like generative models.
In particular, we note that for one-sided stochastic interpolants that model many variants of diffusion and flow based models, there exists the problem of under-sampling regions of the path space where the dimensions/attributes are off-diagonal or asynchronous.
We propose to fix this issue by properly sampling the whole space spanned by combinatorial complexity.
Experiments across two data modalities show that indeed by fully utilizing the combinatorial complexity model, performances can be enhanced, and new generation paradigms can be enabled where different attributes of a data sample are generated in asynchronous schedules to achieve varying degrees of control simultaneously.
We hope that our work can inspire future works that look through the combinatorial perspective of generative models.

#### Limitations and future work

We did not quantify the severity of the undersampling problem for standard diffusion generative models that do not exploit the combinatorial complexity of different attributes and dimensions.
Nevertheless, we conjecture that when the data samples are very densely distributed, the problem of under-sampling combinatorial complexity can be reduced.
But even in that case, ComboStoc training still has the benefit of enabling varying degrees of conditioning in test stage applications, and encourages the learning of correlated structures among different dimensions/attributes.
Practical issues like better batch time step scheduling for image generation training will also be investigated.

## Appendix A Appendix

### A.1 SiT adaptation for image generation

Figure: Figure 11: Adaptation of time step embedding for ComboStoc. Left: the original time step embedding module of SiT, where $N$ is batch size, $C_{F}$ is the sine/cosine frequency embedding length, and $C_{H}$ is the hidden dimension of SiT transformer. Right: the adapted time step embedding module for ComboStoc. $\mathbf{t}$ is now of shape $(N,C,H,W)$. The first two layers remain the same as the original module, applying to each entry of tensor $\mathbf{t}$ and producing a compressed time step encoding of dim $C_{C}$. Given the result tensor of shape $(N,C,H,W,C_{C})$, we further transpose it to combine the channel dimensions and use the same patchwise embedding layer as SiT (and ViT) to embed the local patches into vectors of dim $C_{H}$. Suppose the patch size is $L\times L$, then $T=H\times W/L^{2}$.
Refer to caption: x1.png

As shown in Fig. 11, given a tensorized time step $\mathbf{t}$ of shape $(N,C,H,W)$ that is the same as the latent encoding of input images, we not only encode each of the time steps for different dimensions as done before, but also embed the result feature map of time steps in the same way as image embedding, i.e. the patch-wise embedding originally from ViT [].
This design ensures that the different dimensions are conditioned on their corresponding time steps.
Note that to avoid introducing large embedding layers, we have used $C_{C}=4$ to encode a time-step scalar, which is significantly smaller than the $C_{H}$ of SiT.
This can be the reason why insync_none performs slightly worse than the baseline SiT (Sec. 4), both of which have exactly the same network else where.

### A.2 Encoding for structured 3D shape generation

We have adopted the pretrained part shape encoding network from [].
In particular, Wang et al. [] design a point cloud VAE to encode 3D shapes into a sparse set of latent codes, and on top of the latent set, they train another transformer VAE to compress them into a single latent code.
Therefore, each part shape from the PartNet dataset [] is normalized into unit size and encoded into a single code, which allows us to represent structured 3D shapes as a collection of parts as detailed in Sec. 3.

The embedding modules for part existence and bounding box follow the same design as time step embedding.
That is, we first turn each of the scalar dimensions into frequency codes using the sine/cosine embedding, and then embed them into vectors of dim 4 (cf. Fig. 11), before finally embedding each of the collective attributes as a whole into vectors of hidden dim 384, through respective FC layers.

### A.3 Details of tensorized time steps

In Tab. 4 we give the details of time step specifications for all configurations, across images and structured 3D shapes.
We rely on the broadcast semantics of numpy and pytorch to assign synchronized time steps to multiple dimensions.

**(a) Image generation**
| Setting | $\mathbf{t}$ |
| --- | --- |
| insync_none | $(N,1,1,1)$ |
| insync_patch | $(N,1,H,W)$ |
| insync_vec | $(N,C,1,1)$ |
| insync_all | $(N,C,H,W)$ |

### A.4 Implementation details

The image generation model is modified from SiT-XL/2, i.e. the large model with 28 layers, 1152 hidden dimension, $2\times 2$ patch size, and 16 attention heads.
We trained the model using the default settings of SiT, with AdamW solver and fixed learning rate $10^{-4}$, and batch size 256, on 4 Nvidia H100 gpus.
The training takes 7.5 days for 800K iterations.
Evaluating the models uses the SDE integrator with 250 steps.
The use of classifier-free guidance (CFG) or not is specified at corresponding results.
For comparison with baselines in terms of FID-50K, CFG is not used unless otherwise specified.
In the result gallery figures, CFG is used with guidance strength 4.0.

The structured 3D shape generation model uses a network of SiT small model, i.e. the model has 12 layers, 384 hidden dimension, $256$ tokens for parts and 6 attention heads.
We trained the model using the AdamW solver with a fixed learning rate of $10^{-4}$ and batch size 16.
We trained the model on 4 Nvidia A100 gpus, which takes 3 days for 1.5K epochs.
Evaluating the models uses iterative sampling with 500 iterations; in each iteration, the predicted part existence is binarized via threshold 0.5 before being diffused back for next iteration.
Class conditional sampling without CFG is always applied.

### A.5 More results

Fig. 12 shows another example of image generation where we use varying degrees of data preservation across both spatial dimensions (the four quadrants) and feature channel dimensions.
In particular, we assign spatial preservation weights according to the left column in the figure, and additionally assign 0.5 to the specified channel index $C$ and 0 to other channels, as shown in the middle four columns.
Interestingly, we see that the different channels of the stable-diffusion VAE latent space [] have very different content.
For $C=0$ the first channel, the generated results mostly preserve the spatial structures of the reference images, and the color cues are largely lost.
From $C=1$ to $C=3$, the generated results increasingly preserve the color cues of the reference images but lose more of the structures.
The findings suggest that earlier channels of the VAE latent space emphasize on structures and later ones on image-level color distributions.

Figure: Figure 12: Spatially and channel varying $\mathbf{t}_{0}$. For the spatial dimensions $\mathbf{t}_{0}[:,:,i,j]$, the assignment is specified in the left column. For the feature channel dimensions $\mathbf{t}_{0}[:,C,:,:]$, the $C$-dim is given 0.5 and the rest given 0. Therefore, we obtain results that correspond to the reference images ($\mathbf{x}_{1}$) in complex ways. Notably, earlier channels correspond more to image structures and later channels to image colors.

Fig. 13 shows an example of using a manually specified mask with graded preservation weights to achieve controlled generation of images.
We have assigned a stronger weight ($t_{0}=0.75$) to the face of the red panda, and a lighter weight ($t_{0}=0.5$) to the region around, and let the model generate variations freely at the other regions.
As expected, the generated images show different bodies for the same red panda face, with smooth transitions around the face regions.

Figure: Figure 13: Customized graded $\mathbf{t}_{0}$. Left shows the $\mathbf{t}_{0}$ map in pixels, where each pixel corresponds to a $8\times 8$ patch of the original image. The darker region uses $t_{0}=0.75$ and the lighter region uses $t_{0}=0.5$. Right overlays the map over the reference image. Middle shows generated images.

Tab. 5 gives the comparison between our structured 3D shape generation model and two baselines, i.e., StructRe [] and StructureNet [], in terms of FPD, COV and MMD.
Shapes in PartNet are labeled into semantic parts that are organized into trees, i.e., coarse parts can be decomposed into fine parts by following the tree.
Exploiting this hierarchical data, the two baselines expand coarse parts into fine parts progressively, which helps constrain the generated shapes toward better regularity.
In comparison, our network does not use this hierarchical information and directly generates the leaf level parts.
Nevertheless, the results by insync_all show performances within the baseline results.
Visually, we find our results generally show stronger diversity than the shapes by [, ].
Finally, it is an interesting topic to study how to combine the approaches of hierarchical generation and diffusion generative models, which have differing advantages in aspects of structure regularity and diversity.

**Table 5: Comparison on structured 3D shape generation. Our results are comparable to the baselines that additionally use the hierarchies of shape parts to constrain generations. Best scores are marked in bold and underlined; second best scores in bold.**
| Category | Method | FPD$\downarrow$ | COV$\downarrow$ | MMD$\downarrow$ |
| --- | --- | --- | --- | --- |
| chair | StructureNet | 4.67 | 0.89 | 0.58 |
| StructRe | 2.63 | 0.70 | 0.65 |  |
| Ours | 4.04 | 0.86 | 0.68 |  |
| table | StructureNet | 6.07 | 1.43 | 0.55 |
| StructRe | 1.98 | 0.66 | 0.53 |  |
| Ours | 3.43 | 1.20 | 0.72 |  |

Fig. 14 shows results from preliminary tests on a different scheme of batch mixing for image generative model training (Sec. 4).
For training this model named ComboStoc-TB, we simply blend the asynchronous time steps and the synchronized ones for a whole batch.
In addition, we try to align the time step embedding module with the baseline SiT by setting $C_{C}=C_{H}$ (Fig. 11).
The two modifications combined lead to even larger improvements over baselines.
Indeed, ComboStoc-TB speeds up the convergence of SiT by $\approx 1.75\times$.
We plan to investigate these modifications thoroughly in the future.

Figure: Figure 14: A different scheme of batch mixing for training image generative model. Plots include our baseline model (ComboStoc) and the new batch mixing model (ComboStoc-TB), as well as DiT and SiT for reference; all models are of the scale XL/2 [].

Figs. 15 to 30 show more results generated by ComboStoc models for both structured 3D shapes and images.

### A.6 Broader impact

In this paper we have presented ComboStoc which improves and extends baseline diffusion generative models, across tasks of image generation and structured 3D shape generation.
Image generation can be misused potentially, although our model as well as the baseline model uses very coarse level class name conditioning that prevents highly targeted applications.
Structured 3D shapes are mostly furniture like daily objects, so their generation is unlikely to be misused.
In terms of methodology, we have advocated the importance of sampling the combinatorial flexibility for both model performance and new applications.
The combinatorial complexity of high-dimensional and multiple-attribute data samples can be further explored theoretically based on our work, for example from the perspective of ergodicity [] as emphasized by our dense sampling of all possible interpolation points.

Figure: Figure 15: Class-conditioned structured 3D shapes generated by ComboStoc. Class label is chair.

Figure: Figure 16: Class-conditioned structured 3D shapes generated by ComboStoc. Class label is bed.

Figure: Figure 17: Class-conditioned structured 3D shapes generated by ComboStoc. Class label is table.

Figure: Figure 18: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “promontory, headland, head, foreland” (976).

Figure: Figure 19: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens” (387).

Figure: Figure 20: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “geyser” (974).

Figure: Figure 21: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca” (388).

Figure: Figure 22: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “alp” (970).

Figure: Figure 23: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “Arctic fox, white fox, Alopex lagopus” (279).

Figure: Figure 24: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “volcano” (980).

Figure: Figure 25: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “monarch, monarch butterfly, milkweed butterfly, Danaus plexippus” (323).

Figure: Figure 26: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “valley, vale” (979).

Figure: Figure 27: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “macaw” (88).

Figure: Figure 28: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “snail” (113).

Figure: Figure 29: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “Eskimo dog, husky” (248).

Figure: Figure 30: $256\times 256$ samples of ComboStoc-XL/2 800K. Classifier-free guidance scale = 4.0. Class label = “Siamese cat, Siamese” (284).

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: See weaknesses above.
soundness: 2
strengths: The viewpoint of exploiting the combinatorial complexity of data for improving generative models is novel.
summary: The paper proposes improving the diversity of the data distribution used to train stochastic interpolant models by leveraging the combinatorial nature of our data. The authors achieve this through sampling non-synchronised noised latents by taking the time component to be a vector. The authors then validate this approach experimentally on image and 3D shape generation tasks reporting improved performance.
weaknesses: The paper is not well structured and the arguments given by the paper are also unclear. 

1. The paper mainly considers the conditional interpolation paths for the motivation of their approach, however, the vector field learned by flow matching will in general not be linear as this is defined in terms on an expectation over all conditional vector fields. For example, in appendix a.2, the reasoning is in terms of integrating a linear vector field which doesn't happen in practice.
2. The terminology of the paper is unclear. The authors refer to the class of model they are working with as diffusion and stochastic interpolants in different places. They should be clear that they are using stochastic interpolants with the choice of the linear path and a Gaussian coupling. 
3. The paper does not summarise their new training and sampling procedure, instead leaving the reader to infer this. For example, what exactly is the training objective used in the paper? For images, the authors training a neural network to predict the difference between the data and noise, which for 3D shapes, this predicts the data instead. Further, how is "off-diagonal drift" incorporated during training and sampling and what is the reasoning behind this approach?  
3. The improvement reported by the authors could instead be due to a form of data augmentation enriching their data.
4. How is the form of \rho(x) derived in appendix A. This presentation seems to be non-standard.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - The abstract mentions that the space is spanned by a combination of "dimensions" and "attributes," yet these terms lack clarity at first glance. It would improve readability if the authors could define these terms early in the paper, perhaps in the introduction, and include specific examples to help readers grasp the intended meaning of these concepts.
- Figure 7 indicates a lack of global consistency between the left and right sections of the images. For example, in the third row, when $\lambda > 0.5$, the colors differ on either side of the vertical center line and look unnatural. If ComboStoc effectively learns to denoise across different dimensions with asynchronous timesteps, shouldn’t the transition in the images be seamless?
- In L319, it states that half of the training batches in ImageNet are constructed synchronously, while this is not the case for 3D shape data. Could the authors elaborate more on the relationship between data size and the mixing scheme? 
- In L181, the explanation about samples with asynchronous $\mathbf{t}$ should not predict the original velocity is unclear. Also, how do the two methods presented in Appendix A.2 differ? Are both methods regressing on $x_1 - x_{t_0}$​​ as the target velocity? Does this optimization ensure that the learned model captures the marginal velocity field in a manner similar to existing stochastic interpolation?
soundness: 3
strengths: - The paper is well-written and easy to follow.
- The proposed modification is simple yet effective in exploring the entire path space, which has particular advantages for learning data with combinatorial structures.
- The clear improvements over strong baseline (SiT) models showcase the effectiveness and applicability across various generation tasks.
summary: This paper proposes ComboStoc, a simple but novel fix of diffusion generative models, which replaces the scalar timestep $t$ to tensor. This adjustment aims to effectively cover the entire path space, improving the model’s capacity to handle complex combinatorial structures. Experimental results on class-conditional image generation in ImageNet demonstrate clear improvement over baseline SiT models. The experiment on 3D part generation highlights that ComboStoc is well-suited for data with combinatorial structures. Moreover, ComboStoc also enables asynchronous denoising across different dimensions, allowing various generation scenarios such as image inpainting or 3d part assembling.
weaknesses: - The intuition behind the method is not clear. It is difficult to fully understand why insufficient coverage of low-density regions is critical in synchronous generation in previous approaches. When we follow a synchronous time schedule as in Figure 2(a), it will cover all the paths from $z \rightarrow x_1$. In this context, could the authors elaborate further on why these off-diagonal paths are problematic in synchronous generation?
- Although the proposed method demonstrates empirical effectiveness, it is still unclear why asynchronous timestep schedules should perform better than synchronous counterparts. L155-156 states that ”it encourages the network to learn the correlations of different dimensions and attributes…” but there is no further justification or empirical evidence on this claim. After reading this claim, it seems like such a scheduling strategy might have connections to masked modeling in generative models [1]. By varying noise levels across pixels (or sub-parts) in asynchronous timestep scheduling, the network may more effectively leverage correlations between sub-parts of the data to denoise the entire image. Moreover, this paper seems to provide a similar observation (Figure 3) to [1], where masked modeling accelerates learning structures and makes convergence much faster compared to standard diffusion training schemes. A clearer explanation of the claim in L155-156 would provide valuable insight and better understanding of the proposed method.
- While the proposed method is empirically effective, its technical novelty appears limited, as it essentially replaces a scalar timestep with a tensor representation. 

[1] Gao et al., Masked Diffusion Transformer is a Strong Image Synthesizer, in ICCV 23.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: What does the training algorithm look like in practice, and what are the effects on training time? 

Is tensorized time sampled at train time and/or at test time? 

How relevant is asynchronous time sampling during test time?

By "during testing synchronized timesteps will be used if no graded control is applied", do you mean, that at test time, you generate new images completely from scratch and that all parts of the data share the same time step? 

Is there a potential for spillover or contamination effects in these evaluations?

Is there an alternative approach to account for combinatorial complexity without relying on an asynchronous schedule? How does the asynchronous schedule specifically relate to managing combinatorial complexity?

Have you witnessed any sort of mode collapse and how diverse are the samples with ComboStoc? 

Regarding the number of time steps in section 4.2, do all parts of the data receive the same number of steps, or is there variation? In specific, does the phrase “increase the time steps for individual dimensions and attributes via 1−t0 N for N steps” imply a rescaling of time steps between t0 and 1 so that these samples have the same number of denoising steps?
soundness: 3
strengths: The idea of asynchronous time steps to enhance diffusion models is innovative and introduces a novel approach to accelerate the training of diffusion models. 
The paper provides interesting experiments, especially in section 4.2 with the newly enabled applications, which demonstrate the effectiveness of their method for both image and 3D data.
The way they integrate the tensorized time step embedding with patch embedding is well-explained and integrated into the time embedding module of the SiT model. 
The experiments are supported by clear visualizations, enhancing the understanding of the model's performance. 
This approach appears to yield better qualitative and quantitative results compared to baseline models, without significantly increasing computational costs.
summary: The authors introduce ComboStoc, a new approach that allows diffusion models to better account for the combinatorial complexities in high-dimensional data. Traditional models often underperform because they do not explore all possible data combinations during training. In contrast, ComboStoc implements asynchronous time steps across different parts of the data during training to take advantage of nontrivial correlations. This method enables the model to fully explore the combinatorial space of features and attributes, thus, accelerating training and improving the quality of generated data in tasks involving images and 3D shapes compared to baseline models. Furthermore, the model gains better control over specific features during test time. This enables new generation methods, such as selectively modifying or controlling certain parts of an image or shape.
weaknesses: Some parts of the paper lack clarity and can be confusing to read. 

For instance, it is somewhat unclear what the training algorithm looks like. How do you sample the asynchronous time steps at train time? How does it differ at test time? Do you have to apply the compensation drift at test time as well? Providing a clearer distinction between the training and test time setups would enhance understanding of how ComboStoc operates in typical use cases.
The paper could offer more intuition as to why asynchronous time steps are beneficial and why they encourage the network to learn the correlations of different dimensions and attributes. 

There are some smaller grammatical errors (e.g., "each of the patch token is…"). 
Some abbreviations are introduced without explanation, which could confuse readers unfamiliar with them.
The sentence "To see why the space of possibilities has a combinatorial structure, we note that the data samples are most likely residing on high-dimensional spaces with clear combinatorial structures" is somewhat tautological.
The phrase "Correspondingly, the network fθ(xt) is trained to predict velocity, or the target data sample, etc." is vague, as it’s unclear when to predict velocity versus the target sample, and what "etc." refers to. However, it is good that you later specify that you predict velocity for image data and the target sample for 3D data.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - I am curious whether the effectiveness of the proposed method essentially depends on data scarcity, i.e., whether the gain is more significant in limited data regime. The under-sampling issue the paper described might become less important with enough data.
- Table 5: How could the proposed method maintain the same training/inference speed, despite of its increased GFLOPS, e.g., 237.34 → 352.46?
soundness: 2
strengths: - The paper is easy to follow.
- The proposed method is notably simple.
- The proposed method shows more effectiveness on structured 3D generation tasks (compared to well-studied ImageNet), an area where more study has been demanded.
- The experimental results are consistent in showing the effectiveness of the method.
- The application suggested from the method, i.e., asynchronous reconstruction, looks novel and has potential to be applied to other tasks.
summary: The paper shows that introducing multi-dimensional independent timesteps, each of which is responsible for different parts of input, e.g., per patch/feature basis, etc., can improve the overall training of diffusion models. Intuitively, this approach can help to extend the sampling coverage of input space during diffusion training, considering that the standard training only samples input along its corresponding diffusion trajectory. Experimental results on ImageNet (for image generation experiments) and PartNet (for 3D generation), mainly based on SiT, primarily report consistent improvements of FID. The paper also suggests a new application of the proposed method as a byproduct, viz., asynchronous reconstruction of input, which can be useful for image editing, fine-grained part control in 3D, etc.
weaknesses: - The overall experiments are limited to a single architecture, namely SiT. This raises a question whether the effectiveness of the method generalizes to other diffusion architectures.
- In case of ImageNet, the effectiveness of the method is solely supported by FID.
- Table 1: I think it is more fair to add SiT-XL (cfg=1.5) @ 800k result when evaluating ComboStoc (cfg=1.5) @ 800k.
- Discussions about computational efficiency are missing in the main text, while it seems important to have one as the method introduce extra complexity in diffusion model architectures (and may introduce overhead). For example, Figure 3 may also have a wall-clock comparison.
- Figure 8: There is no “baseline” results, thereby it is hard to judge whether the method indeed has effectiveness for the proposed application (or whether the standard diffusion model could just do the task as well). I think it would be also good to add some quantitative results in Section 4.2.
- It seems to me that at least some part of Appendix A.2 should be in the main text rather than Appendix. To my understanding, the “off-diagonal drift minimization” is a crucial component to enable the method to work, while its whole explanation is deferred to Appendix.

---

confidence: 5
contribution: 4
initial_rating: 8
presentation: 4
questions: The issues raised in this paper are very interesting, and I'm curious if the authors could do some experimentation or discussion of the more commonly used Noise Schedules (e.g., Coseine Schedules).
soundness: 4
strengths: + The combinatorial complexity in training of the diffusion model is not explored. This paper find this proboem and propose a sample way to simple fix to this problem. I think this find is really  interesting.
+ The author validat the method on both advanced DiT and SiT architectures, which fully illustrates the effectiveness of the methodology. Since SiT and DiT have different Noise Schedules, which also illustrates the robustness of the method.
+ The paper is well written and motivated with a clear description. And the code is given which makes the work very solid.
summary: This paper finds the combinatorial complexity problem of diffusion model training and proposes the ComboStoc method to train diffusion models, the core motivation of the author is that the existing interpolation methods lead to insufficient training of the model in the different timestep (i.e., $x_1|t=0.1. x_2|t=0.1$, in each training iter timestep is **same on all pixel**), so the authors design a new interpolation scheme in different pixel to improve it (i.e., $x_1|t=0.1. x_2|t=0.8$). Extensive experiments illustrate the effectiveness of the propsed method.
weaknesses: + Some related works[1,2] have found contradictory gradients for different timesteps, while others[3, 4] find non-uniform variations of different timesteps during diffusion, which is of some relevance to this paper (improve the training of the diffusion models), and I suggest that the authors add this part of the discussion to the related work. 

  [1] Addressing Negative Transfer in Diffusion Models NeurIPS-2023

  [2] Efficient Diffusion Training via Min-SNR Weighting Strategy ICCV-2023

  [3] A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training Arxiv-2024

  [4] Beta-Tuned Timestep Diffusion Model ECCV-2024

+ Figure 11 exemplifies the non-uniformity of the existing method, and the authors would be well advised to provide an optimized comparison plot, thus allowing the validity of the method to be better verified.