PaperlayerView exact revision on arXiv

Source-linked AI summary

Complexity-Balanced Diffusion Splitting

Noam Issachar, Dani Lischinski, Raanan Fattal

arXiv:2606.06477v1cs.CV

TL;DR

Continuous-time generative models deploy monolithic networks across heterogeneous denoising regimes, motivating a more efficient temporal allocation of capacity. CBS uses approximation theory and de Boor’s equidistribution principle, with tractable spatial and geometric complexity monitors, to partition the timeline. Across architectures and datasets, it improves synthesis quality without increasing per-step inference cost and avoids costly temporal-boundary searches.

  • Problem

    Monolithic diffusion networks are deployed uniformly across heterogeneous denoising regimes, while existing temporal partitions rely on heuristics or expensive searches.

  • Method

    CBS partitions the diffusion timeline into equal-approximation-burden segments using Dirichlet-energy and trajectory-acceleration monitor functions.

  • Results

    CBS consistently improves synthesis quality across SiT, JiT, and UNet without increasing per-step inference cost, including ~35% FID improvement with CFG on SiT-XL versus naive splitting.

  • Takeaways & Limitations

    Complexity-based boundaries can be estimated before training with negligible overhead, eliminating costly empirical searches for temporal partitions.

  • Takeaways & Limitations

    CBS focuses on temporal splitting; extending equidistribution to spatial splitting is left for future work, and deriving suitable monitor functions is expected to be challenging.

Abstract

from arXiv · show

Standard continuous-time generative models rely on monolithic architectures that must navigate vastly different signal regimes, from isotropic noise to intricate data distributions. While scaling model capacity improves performance, deploying a massive network uniformly across the entire generative timeline is inherently inefficient. In this work, we propose Complexity-Balanced Splitting (CBS), a principled framework for temporal capacity allocation that distributes the generative workload across multiple specialized sub-networks. Grounded in function approximation theory and de Boor's equidistribution principle, CBS partitions the diffusion timeline into segments of equal approximation burden, allocating more representational capacity to regions where the generative dynamics are more difficult to model. To estimate this local complexity, we introduce two complementary and tractable monitor functions: a spatial measure based on the flow's Dirichlet energy, and a geometric measure based on the acceleration of the sampling trajectories. Using a lightweight auxiliary model to estimate these complexity profiles, our approach eliminates the need for heuristic temporal splits or computationally expensive search procedures. Extensive evaluation across multiple architectures (SiT, JiT, and UNet) and datasets demonstrates that CBS consistently improves synthesis quality without increasing per-step inference cost. In particular, CBS improves FID by ~35% on SiT-XL with CFG relative to naive temporal partitioning. Project page is available at https://noamissachar.github.io/CBS/.

1 Introduction

CBS addresses the inefficiency of monolithic diffusion models by partitioning the generative timeline according to approximation burden. It estimates local complexity with spatial and geometric monitors and improves synthesis quality without increasing per-step inference cost.

  • Monolithic diffusion networks handle heterogeneous denoising regimes uniformly, although no individual regime warrants the full scaled-up capacity.
  • Existing temporal-specialization methods use heuristic splits or expensive searches over candidate boundaries, lacking a principled allocation criterion.
  • CBS applies de Boor’s equidistribution principle to divide the diffusion timeline into segments with equal approximation burden.
  • CBS estimates local approximation burden using flow Dirichlet energy and sampling-trajectory acceleration.
  • ~35% FID improvement with CFG on SiT-XL relative to naive splitting accompanies consistent synthesis gains without higher per-step inference cost.The reported improvement is ~15% without CFG and ~35% with CFG.

2 Preliminaries

The paper formulates continuous-time generative modeling as velocity-field approximation and uses domain decomposition to distribute approximation burden across temporal intervals. De Boor’s principle places finer partitions where complexity is higher, while minimizing maximal local error targets sampling quality.

  • Continuous-time interpolants bridge a Gaussian noise prior and target data through a time-augmented state trajectory.
  • Flow matching and v-prediction diffusion train a neural network to approximate the trajectory’s instantaneous velocity field.
  • The target velocity field has complexity that varies over time, motivating approximation-theoretic allocation of error across temporal subdomains.
  • Domain decomposition assigns separate localized models to disjoint intervals and chooses splitting points so intervals have equal modeling complexity.
  • De Boor’s principle clusters nodes densely where the target function is complex and sparsely where it is smooth, using a monitor function to quantify local burden.
  • Equal monitor-function integrals across sub-intervals equidistribute the approximation bound, giving each local model a comparable representational burden.
  • For path-integrated flow fields, a large local error can permanently divert the sampling trajectory, so minimizing maximal error is directed toward better sampling quality.
  • Standard denoising objectives average error across the denoising interval rather than minimizing the maximum instantaneous error.

3 Method

CBS derives temporal splits from local approximation burden rather than heuristic boundaries, using complementary spatial and trajectory-based monitor functions. A lightweight auxiliary flow model estimates these monitors, and de Boor-style equidistribution determines boundaries used to train and switch specialized networks.

  • Monitor functions: CBS uses two complementary monitors: Dirichlet energy measures spatial flow complexity, while trajectory acceleration measures geometric path complexity.The spatial monitor is motivated by approximation bounds and estimated through spatial variation; the trajectory monitor uses the second-order time derivative.
  • Spatial flow complexity: The Dirichlet-energy monitor provides a computable substitute for spectral complexity, which cannot be assessed directly for the high-dimensional denoising flow.Parseval-based reasoning links rapid spatial variation to spectral energy, and a bandwidth bound connects this quantity to approximation error.
  • Trajectory complexity: The trajectory monitor uses k = 2 acceleration because it filters constant-velocity displacement and isolates path curviness.The resulting second-order derivative approximates curvature when trajectory velocity magnitudes are relatively constant.
  • Monitor estimation: A smaller auxiliary network trained on 10% of the dataset estimates monitor values from K sampling trajectories over a 100-point temporal grid.The auxiliary model only estimates temporal boundaries, so an approximate architecture and reduced training budget are sufficient.
  • Time splitting and deployment: CBS selects split points by equally partitioning the cumulative monitor value, then uses those boundaries for specialized-network training and inference.Each network is trained only on its designated interval, and inference switches networks across the resulting temporal segments.

4 Experiments

CBS is evaluated across architectures, datasets, monitor functions, and network counts, with complexity-derived temporal partitions improving synthesis quality under matched inference costs. The experiments also test boundary optimality and the practicality of estimating complexity with lightweight auxiliary models.

  • Experimental settings: CBS is evaluated on latent ImageNet-256 with SiT, pixel-space ImageNet-64 with JiT, and unconditional CIFAR-10 with UNet.The study includes transformer and convolutional architectures across latent, pixel-space, conditional, and unconditional settings.
  • Generative performance: CBS improves SiT synthesis quality across model capacities while maintaining the same active parameter count and per-step inference cost as the monolithic baseline.Table 1 reports improvements in FID, IS, Recall, and Precision with and without CFG.
  • Generative performance: CBS improves pixel-space JiT synthesis under an identical budget of 131M activated parameters and 25 per-step GFLOPs.This tests whether complexity-based partitioning handles raw high-frequency spatial gradients beyond latent-space generation.
  • Network scaling: Increasing the number of specialized networks from N = 1 to N = 4 consistently improves both FID and Inception Score, with N = 3 used by default.The default balances synthesis quality with total training overhead, while N = 4 demonstrates continued scaling.
  • Boundary optimality: Temporal splits with equal cumulative complexity consistently achieve the lowest FID among perturbed boundaries, supporting the monitor as a proxy for empirical learning burden.The optimality study perturbs each derived boundary separately before comparing the resulting models.
  • Monitor comparison: The path-acceleration monitor achieves better FID than the Dirichlet-energy monitor in the reported comparison and is therefore used by default.Both monitors are near-optimal on SiT-B/2, while path acceleration is attributed a closer connection to final sampling accuracy.
  • Estimation efficiency: Lightweight auxiliary configurations produce nearly identical complexity curves and exactly the same temporal boundaries.This addresses the practical overhead of relying on a pre-trained auxiliary model for complexity estimation.

5 Related Work

Prior work addresses temporal specialization, approximation complexity, and model scaling, while CBS uses temporal capacity allocation to preserve active inference cost as total capacity grows.

  • Temporal Specialization in Diffusion Models: Temporal specialization assigns different denoising phases to specialized networks, but existing methods require exhaustive searches to identify transition boundaries.Cascaded models partition by spatial resolution rather than time.
  • Approximation Theory in Neural Networks: Approximation theory links neural-network error and required capacity to the target function’s spectral or high-order spatial complexity.Barron’s and Yarotsky’s results provide theoretical foundations for complexity-aware allocation.
  • Scaling Up Generative Models: Scaling laws motivate larger generative architectures, but CBS distributes expanded capacity across time while keeping active parameters constant at each timestep.This decouples total parameter scaling from per-step inference cost.

6 Conclusion

CBS frames temporal partitioning as approximation-theoretic domain decomposition, using complexity monitors to allocate capacity across diffusion time. Evaluations report improved quality without higher per-step inference cost, while spatial splitting remains future work.

  • Core contribution: CBS divides the diffusion timeline into segments of equal representational burden using de Boor’s equidistribution principle.The burden is quantified with spatial Dirichlet energy or temporal path acceleration.
  • Empirical outcome: Empirical evaluations across SiT, JiT, and UNet report improved synthesis quality without increasing per-step inference costs.Complexity metrics also allow pre-training boundary estimation without costly empirical searches.
  • Scope and future work: The current work focuses on temporal splitting and reaches close to optimal solutions only within that axis.Spatial token routing and its monitor functions are left for future work.
  • Practical implication: CBS decouples total parameter capacity from inference costs by concentrating computation where generative dynamics demand it most.This supports scaling model capacity while maintaining constant active computation per timestep.

A Bounding Spectral Complexity via Dirichlet Energy

The appendix bounds spectral complexity using Dirichlet energy through Cauchy–Schwarz and Parseval’s identity, under an effectively band-limited flow-field assumption.

  • Derivation: The derivation establishes a relationship between spectral complexity C_vt and the Dirichlet energy E_D(v_t) using the Cauchy–Schwarz inequality.The spectral-complexity definition is the starting point for the bound.
  • Assumption: Bounding the relevant frequency integral requires a decaying weight, bounded support, or an effectively band-limited flow field.The practical assumption restricts integration to a frequency ball of radius Ω_max.
  • Cauchy–Schwarz bound: Applying Cauchy–Schwarz over B(0, Ω_max) separates the integral into the ball’s volume and a gradient-related spectral term.The latter can be extended to R^d when energy outside the ball is negligible.
  • Parseval connection: The volume term is denoted V(Ω_max), while Parseval’s identity connects the remaining gradient term to Dirichlet energy.Substitution yields the final inequality relating spectral complexity to the energy measure.

B Cumulative Path Acceleration Analysis

Cumulative path acceleration is nonlinear across the evaluated architectures and datasets, so CBS places narrower intervals where trajectory acceleration is steepest.

  • Compared environments: Figure 2 visualizes cumulative path acceleration for SiT on latent ImageNet-256, JiT on pixel-space ImageNet-64, and UNet on CIFAR-10.The figure compares temporal dynamics across three architectures and data modalities.
  • Observed dynamics: The accumulated trajectory acceleration is distinctly nonlinear across all evaluated architectures and data modalities.Equal-acceleration boundaries therefore adapt to dataset-specific dynamics.
  • Capacity allocation: CBS assigns narrower time intervals and higher localized parameter capacity to the steepest phases of the cumulative acceleration curve.This provides a visual explanation for its complexity-based temporal splits.

C Implementation Details

The experiments use NVIDIA H200 GPUs with PyTorch DistributedDataParallel across SiT, JiT, and UNet architectures. Training and evaluation configurations are standardized within each generative domain, while Figure 2 visualizes cumulative path acceleration across datasets and models.

  • Hardware and Compute Resources: Experiments used NVIDIA H200 GPUs and PyTorch DistributedDataParallel to scale training across SiT, JiT, and UNet.The software environment used standard PyTorch and CUDA releases.
  • Hyperparameters and Configuration: Models within each generative domain used identical base training, optimization, and sampling configurations unless explicitly stated otherwise.The complete hyperparameter set is provided in Table 7.
  • Hyperparameters and Configuration: Table 7 lists training and sampling hyperparameters for the baseline architectures evaluated in the experiments.
  • Diagnostic Visualizations: Figure 2 visualizes cumulative path acceleration over t ∈[0, 1] for SiT, JiT, and UNet across datasets and models.
Loading 2606.06477v1…