PaperlayerView exact revision on arXiv

Source-linked AI summary

ComboStoc: Combinatorial Stochasticity for Diffusion Generative Models

Rui Xu, Jiepeng Wang, Hao Pan, Yang Liu, Xin Tong, Shiqing Xin, Changhe Tu, Taku Komura, Wenping Wang

arXiv:2405.13729v3cs.LGcs.AIcs.CVcs.GR

TL;DR

Diffusion models can undersample off-diagonal combinations of dimensions and attributes, limiting coverage of structured data during evaluation. ComboStoc samples this combinatorial space uniformly, improving training across images and 3D shapes while enabling asynchronous, controllable generation.

  • Problem

    Existing diffusion training schemes undersample regions of the corrupted-sample path space, leaving some dimension–attribute combinations insufficiently covered.

  • Method

    ComboStoc uniformly samples the combinatorial space by using vectorized diffusion timesteps across dimensions and attributes.

  • Results

    ComboStoc improves training convergence across image and structured 3D-shape generation, with lower ImageNet FID scores and substantially better 3D generation quality than baselines.

  • Takeaways & Limitations

    The approach supports asynchronous timestep inference, enabling graded control over attributes in applications such as image inpainting and 3D part-level assembly.

  • Takeaways & Limitations

    ComboStoc has limited expected impact when data dimensions are nearly independent and lack strong structural or combinatorial information.

Abstract

from arXiv · show

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, additional attributes are combined to associate with data samples. We show that the space spanned by the combination of dimensions and attributes can be insufficiently covered by existing training schemes of diffusion generative models, potentially limiting 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. Our code is available at: https://github.com/Xrvitd/ComboStoc

1 Introduction

ComboStoc addresses biased coverage in diffusion training by sampling the combinatorial space of high-dimensional data with asynchronous time steps. This improves image and 3D structured-shape generation while enabling graded, asynchronous test-time control.

  • Motivation: Existing diffusion training often follows a single transport path, which can bias sampling density across corrupted-sample space.This creates a mismatch between the training scheme and the space evaluated at test time.
  • Method: ComboStoc samples combinatorial complexity by assigning asynchronous diffusion time steps to patches, parts, attributes, and feature dimensions.The modification fully samples subspaces spanning combinations of source and target data points.
  • Results: ComboStoc improves generative models across images and 3D structured shapes by incorporating combinatorial sampling into training.The paper reports systematic FID-50k improvements over baseline SiT and DiT on ImageNet, while the scheme is indispensable for a working 3D structured-shape model.
  • Test-time capabilities: Asynchronous generation enables flexible conditioning on partial observations and graded control across image patches and channels.The framework supports control beyond binary masks by varying schedules for different patches, parts, or attributes.
  • Contributions: ComboStoc vectorizes diffusion time steps during training to better capture and reason about structured and combinatorial data.The stated contributions include consistent performance across image and 3D structured-shape domains and accelerated ImageNet training.

2 Related Works

Prior work improves diffusion training efficiency and inference control, while comparatively little attention has been given to combinatorial complexity. Related research also develops structured 3D generation and orthogonal methods for accelerating or stabilizing diffusion and flow-based training.

  • Image diffusion: Image-generation research refines diffusion loss weighting and time-step schedules, uses distillation for acceleration, and enforces sampling-path consistency.These advances leave the role of combinatorial complexity in diffusion training relatively underexplored.
  • Inference control: Spatially or temporally varying noise schedules support finer-grained diffusion inference control, including stroke-based image editing and adaptive noise conditioning.SDEdit applies a global noise level to coherently blend user strokes into generated images.
  • Structured 3D generation: Structured 3D generation research includes hierarchical shape representations, VAE-based structured variation, rewriting-based cross-category generation, and controllable part-level decomposition.BANG produces temporally coherent exploded states, while X-Part uses bounding-box prompts and pointwise semantic features for editable 3D assets.
  • Training acceleration: REPA accelerates diffusion training by distilling self-supervised visual representations from clean images into intermediate latent representations of noisy inputs.Other recent studies investigate diffusion- and flow-based training acceleration or stabilization from orthogonal perspectives.

3 Background on Diffusion Models

Diffusion and flow-matching models transform simple source distributions into target data distributions along stochastic interpolation paths. This section identifies non-uniform path-space coverage as a source of outliers and slow convergence, motivating ComboStoc’s desynchronized sampling strategy.

  • Diffusion-model background: Diffusion, score-based, and flow-matching models transform a simple source distribution into a target distribution through dynamics specified by stochastic differential equations.The paper uses a linear one-sided interpolant for conceptual simplicity and strong practical performance.
  • Sampling bias: These deficiencies are especially relevant for structured data with insufficient samples and unstructured data with large-scale datasets, including 3D shapes and images.The passage gives 18K samples for PartNet 3D shapes and 1.3M samples for ImageNet images.
  • Sampling bias: Standard interpolation produces non-uniform path-space sampling density, with regions farther from target data points receiving less training coverage.The density grows toward target points, creating shrinking coverage along the generation path.
  • Sampling bias: Non-uniform coverage can produce non-converging outliers because sparsely trained regions are encountered during numerical integration.Toy particle simulations illustrate this effect for flow matching.
  • ComboStoc motivation: Desynchronizing interpolation schedules and sampling all combinations of attributes and feature dimensions makes sampling density uniform within subregions, enabling robust low-data generation and faster rich-data convergence.This strategy is introduced as the response to biased sampling.

4 Combinatorial Stochastic Process

ComboStoc addresses the limited coverage of standard synchronous diffusion training by assigning asynchronous schedules across dimensions and attributes, producing uniform sampling over combinatorial subregions. This broadens network coverage, encourages cross-attribute correlations, supports graded control at generation time, and defines a proper generative flow model.

  • Motivation: Standard diffusion models treat all dimensions and attributes homogeneously and synchronously, leaving low-density regions insufficiently covered.This limitation arises for high-dimensional data such as image latents and structured 3D shapes with varying parts, positions, and bounding boxes.
  • Method: ComboStoc assigns independently sampled asynchronous diffusion times to dimensions and attributes, yielding uniform sampling within subregions spanned by source–target pairs.The schedule tensor has the same shape as the data sample, and its entries are sampled uniformly within [0, 1].
  • Benefits: The resulting augmented samples broaden network coverage, improve test-time robustness and quality, and train the network to learn correlations among dimensions and attributes.The network is trained to synchronize dimensions and attributes toward final data points, rather than only following synchronized paths.
  • Benefits: Asynchronous schedules enable different dimensions and attributes to receive varying degrees of finalization in the synthesized result.This provides a mechanism for flexible control over the generation process at test time.
  • Toy analysis: In a toy simulation, ComboStoc produces broader spatial coverage, fewer outliers, and more concentrated particle trajectories during convergence.The velocity field is illustrated through particle-based simulations because the summed field over source-derived spans lacks a closed-form expression in that setting.
  • Theoretical validity: The scheme defines a conditional vector field on the rectangular span between source and target samples and satisfies the continuity equation, establishing a proper generative flow model.Unlike FlowMatching’s diagonal support, ComboStoc uses the full rectangular subspace spanned by x0 and x1 and conditions the field on both endpoints.

5 Results and Discussion

Section 5 shows that ComboStoc improves diffusion-model training convergence for images and structured 3D shapes, while its asynchronous time steps enable novel applications. It distinguishes these training and inference contributions, noting that standard generation-quality improvements use synchronized inference timesteps.

  • Training convergence: ComboStoc improves training convergence for diffusion generative models on images and structured 3D shapes.These results are presented in Sec. 5.2.
  • Asynchronous inference: Asynchronous time steps enable novel applications for diffusion-model generation.These applications are demonstrated in Sec. 5.3.
  • Training and inference contributions: Standard generation-quality improvements are obtained with synchronized timesteps during inference, separating ComboStoc’s training and inference contributions.The paper reports these improvements in Figs. 6–7 and Tab. 4.

5.1 Implementation Details

The paper uses SiT-based image and 3D shape models with modality-specific architectures, training settings, and sampling procedures. Inference is synchronized for standard generation and asynchronous for graded control and inpainting, with mask values determining per-dimension preservation and evolution rates.

  • Image generation: The image model modifies SiT-XL/2 with 28 layers, 1152 hidden dimension, 2 × 2 patches, and 16 attention heads, trained for 800K iterations.Training uses AdamW, fixed learning rate 10−4, batch size 256, and 4 Nvidia H100 GPUs; evaluation uses 250-step SDE integration.
  • Structured 3D shape generation: The 3D shape model uses SiT-small with 12 layers, 384 hidden dimension, 256 part tokens, and 6 attention heads, trained for 1.5K epochs.Training uses AdamW, fixed learning rate 10−4, batch size 16, and 4 Nvidia A100 GPUs; evaluation uses 500 iterative-sampling iterations with a 0.5 existence threshold.
  • Synchronized inference: Standard image and 3D generation uses a synchronized schedule that interpolates one scalar timestep into 250 values and broadcasts it across all dimensions.This schedule follows the synchronized procedure in Algorithm 1(A).
  • Asynchronous inference: Hierarchical control and inpainting use an asynchronous schedule initialized as x(0) = (1 −m) ⊙z + m ⊙x1, where m specifies per-dimension preservation.Each mask entry determines both the initial state and its effective timestep, with larger m_i values starting closer to x1 and using smaller step sizes.
  • Asynchronous inference: Different dimensions follow different effective timestep schedules, yet all complete evolution within the same 250 integration steps, enabling graded control and spatially varying inpainting.Larger m_i values evolve more slowly, while smaller values evolve more freely.

5.2 Improved Training of Diffusion Models

ComboStoc improves diffusion-model training by exploiting combinatorial structure through asynchronous timestep schedules, with gains across images and structured 3D shapes. The approach is especially important for highly combinatorial 3D generation, while remaining efficient in actual runtime.

  • Overall findings: Across image and structured-shape experiments, configurations exploiting combinatorial complexity improve over baselines without asynchronous time schedules.The paper reports that stronger combinatorial complexity makes the scheme increasingly important for training a working model.
  • Images: The unsync_none image setting performs slightly worse than baseline SiT because its timestep embedding module is smaller.Unsync_none and baseline SiT use the same architecture elsewhere; more capable embedding layers might close this gap, while unsync_all already outperforms the baseline.
  • Images: For ImageNet generation, unsync_all consistently improves baseline SiT and significantly improves over reference DiT, including under a fair wall-clock budget.ComboStoc is slightly slower per training step than SiT because of tensorized timestep embedding, but the improvement persists after converting training steps to wall-clock time.
  • Efficiency: Despite a moderate increase in GFlops, ComboStoc is as efficient as baseline SiT and DiT models in actual runtime.The computational comparison is provided alongside the image-generation results.
  • Structured 3D shapes: For structured 3D shapes, exploiting more combinatorial complexity improves performance, while unsync_none almost entirely fails to produce meaningful shapes.Spatial-part unsynchronization is particularly helpful because the task composes flexible numbers of parts; the unsync_all setting achieves the best overall result.
  • Structured 3D shapes: Part-level combinatorial stochasticity significantly enhances structured-shape generative performance, and unsync_all shows the best overall result on the reported chair-category metrics.The evaluation uses FPD, COV, and MMD; the model is also compared with StructRe and StructureNet in these metrics.

5.3 Applications

ComboStoc’s asynchronous timesteps enable test-time control over how strongly different image dimensions, feature channels, and structured-shape attributes are preserved. The applications demonstrate graded image preservation and coherent 3D shape completion and assembly.

  • Applications: Asynchronous timesteps let users assign different preservation strengths to dimensions and attributes during test-time generation.The process starts from a partially preserved sample and increases individual time steps asynchronously.
  • Images: Increasing image preservation weights strengthens reference-content preservation while allowing the remaining regions to generate from scratch.The red-panda example illustrates how larger t0 values increase constraint capability.
  • Images: Spatially varying preservation weights preserve image quadrants to different degrees while generating coherent content across regions.This provides controlled generation beyond binary inpainting.
  • Images: Channel-varying preservation reveals that earlier VAE latent channels emphasize image structures, whereas later channels emphasize color distributions.The result combines spatially varying preservation across quadrants with channel-specific weights.
  • Structured 3D shapes: In structured 3D shapes, setting chair bases to t0 = 0.9 enables diverse class-consistent completion, while preserving part shapes and scales enables position-only assembly.The assembly setting fixes part rotations and leaves rotating shape parts for future work.

5.4 More Discussions

This section analyzes off-diagonal drift during asynchronous integration and compares two velocity adaptations, selecting drift minimization for its better convergence and structural preservation. It also reports that ComboStoc converges faster and remains more robust than SiT in a low-data regime.

  • Off-diagonal drift: Asynchronous tensorized timesteps can make the original velocity x1 − z miss the target data point when integration encounters an off-diagonal sample.Integration proceeds until the slowest timestep finishes, with t0 = min(t0).
  • Off-diagonal drift: Off-diagonal drift minimization compensates the velocity by following the negative gradient of a drift potential, promoting convergence to target data points.The drift potential is Φ(δ(x_t)) = 1/2 ∥δ_t∥^2.
  • Velocity adaptations: The cone-shaped velocity field also converges to target points, but its normalization by the slowest dimension introduces scaling and may degrade velocity regression.For synchronized schedules, it reduces to the original velocity x1 − z.
  • Ablation study: In ablations, off-diagonal drift minimization converges better than no compensation, produces more seamless asynchronous image transitions, and preserves structural similarity better than the alternatives.The cone-shaped field converges more slowly and yields poorer results, while uncompensated generation shows a midline discontinuity.
  • Insufficient data: ComboStoc-S/2 consistently achieves lower training loss than SiT-S/2 on 1,000 ImageNet images despite using slightly fewer parameters.ComboStoc-S/2 has 32.36M parameters, while SiT-S/2 has 32.58M.
  • Insufficient data: ComboStoc produces higher-quality, more stable images earlier and demonstrates faster convergence and stronger robustness than SiT in low-data training.SiT often remains blurry after 200K steps or produces chaotic early outputs, whereas ComboStoc yields clear, semantically coherent samples midway through training.

5.5 Computational Complexity Analysis

The analysis compares ComboStoc with SiT and DiT under matched XL/2 image-generation settings across model size, speed, memory, and computation. ComboStoc uses fewer parameters and slightly less training memory, while its tensor-shaped conditioning differs from the vector conditioning used by DiT and SiT.

  • Computational Complexity Analysis: The comparison evaluates parameter count, training speed, memory usage, inference speed, and GFlops for XL/2 models on one Nvidia A100-80G GPU.Tests use 256×256 inputs and batch size 256; GFlops are calculated with DeepSpeed.
  • Computational Complexity Analysis: ComboStoc has fewer parameters than DiT and SiT because it uses a smaller timestep embedding module.The paper attributes this architectural difference to the timestep embedding design shown in Fig. 5.
  • Computational Complexity Analysis: ComboStoc’s training GPU memory cost is slightly lower than SiT’s, despite using tensor-shaped conditioning rather than a channel-sized conditioning vector.DiT and SiT condition with vectors of the image tensor’s channel size, whereas ComboStoc uses a tensor matching the image tensor’s shape.

6 Conclusion

The paper addresses combinatorial complexity in diffusion generative models, where one-sided stochastic interpolants can undersample off-diagonal or asynchronous regions of path space. ComboStoc targets this issue by sampling the whole space, but is most effective when data contains combinatorial and structural information.

  • Problem: Diffusion models can undersample off-diagonal or asynchronous regions of path space for high-dimensional, multi-attribute data.The paper frames this as a combinatorial-complexity problem involving dimensions and attributes in samples such as images and 3D shapes.
  • Method: ComboStoc fixes the undersampling issue by sampling the whole space.The proposed scheme applies to one-sided stochastic interpolants that model many diffusion- and flow-based models.
  • Limitations and future work: ComboStoc has significant effects only when data contains combinatorial and structural information, such as image patches or 3D-shape parts.This is identified as a limitation and condition for exploiting correlations under the combinatorial schedule.
  • Limitations and future work: When vector-space dimensions are nearly independent, exploiting dimension correlations and training a model for different-dimensional combinatorial schedules is difficult.In that setting, individual dimensions may be ideal candidates for separate treatment.
Loading 2405.13729v3…