Complexity-Balanced Diffusion Splitting

Abstract: 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/.

• The paper introduces Complexity-Balanced Splitting (CBS), a method that uses equidistribution principles to allocate generative workload evenly across temporal sub-networks.
• The approach partitions the diffusion process based on complexity monitor functions such as Dirichlet spectral energy and path acceleration to optimize sample quality.
• Empirical tests on datasets like ImageNet and CIFAR-10 demonstrate that CBS significantly outperforms uniform splits and monolithic designs without increasing inference cost.

Complexity-Balanced Diffusion Splitting: Principled Temporal Capacity Allocation in Continuous-Time Generative Models

Overview

"Complexity-Balanced Diffusion Splitting" (2606.06477) introduces a mathematically principled methodology for temporal capacity allocation in continuous-time generative models, specifically diffusion models. The core contribution is Complexity-Balanced Splitting (CBS), a framework that distributes generative workload across multiple specialized sub-networks, each assigned to a distinct segment of the diffusion timeline. Unlike heuristic or computationally expensive search-based partitioning schemes, CBS leverages function approximation theory—especially de Boor’s equidistribution principle—to partition the generative process so that each sub-network handles an equally challenging portion of the trajectory, as rigorously quantified via tractable complexity monitor functions.

CBS demonstrates strong improvements in sample quality compared to uniformly split architectures, achieving up to a 35% reduction in FID on large models with classifier-free guidance, while keeping per-step inference cost unchanged.

Motivation and Theoretical Foundation

Large diffusion models conventionally use a single, monolithic neural network to denoise across the entire signal trajectory, even though the statistical and geometric complexity of these phases varies dramatically. This global approach is suboptimal: scaling up improves the aggregate modeling capacity but fails to efficiently allocate resources to the most challenging regions of the generative process. Previous efforts to temporally specialize sub-networks have been either heuristic (naive, uniform splits) or reliant on expensive empirical searches over the timeline.

CBS formalizes this temporal specialization by treating it as a domain decomposition (partitioning) problem in approximation theory. Invoking de Boor’s equidistribution principle, the generative timeline $[0,1]$ is split so that each sub-network operates on a temporal segment with equal integrated approximation burden. The monitor function $m(t)$ quantifies, at each point in time, the local complexity of modeling the generative flow field. The partition points $\{t_i\}$ are then derived so that

$$\int_{t_{i-1}}^{t_i} m(t)\, dt = \frac{1}{N} \int_0^1 m(t)\, dt,$$

ensuring balanced representational requirements among all sub-networks.

Complexity Monitor Functions

CBS introduces two primary complexity measures:

• Dirichlet Spectral Energy: This uses the Dirichlet energy of the velocity field, providing a computable upper bound (via Barron's theorem and frequency analysis) on how difficult it is for a neural network to fit the flow at each time point. The monitor function is approximated efficiently with randomized trace estimators to sidestep the computational infeasibility of full-Jacobian evaluation in high-dimensional spaces.
• Path Acceleration: This geometric measure captures the temporal nonlinearity of the sampling trajectory by considering the magnitude of the trajectory's second derivative (acceleration) with respect to time. This reflects the "curviness" of the generative path—the more complex the acceleration profile, the more difficult the sub-task.

Both measures require an auxiliary, lightweight network trained briefly over the full timeline to provide rough trajectory and flow approximations adequate for boundary computation. The acceleration-based monitor is found empirically to be more directly linked to sample quality.

Figure 1: SiT (ImageNet-256).

CBS Training and Inference Pipeline

The full CBS pipeline consists of the following stages:

1. Auxiliary Model Training: A small, inexpensive model is trained on a small subset of data to approximate the flow.
2. Complexity Profiling: This auxiliary model is used to generate trajectories and evaluate the monitor function $m(t)$ on a temporal grid.
3. Timeline Partitioning: Boundaries $\{t_i\}$ are computed by discretely enforcing equidistribution of integrated monitor values.
4. Sub-Network Specialization: For each temporal segment $[t_{i-1}, t_i]$, a dedicated neural network is trained and deployed exclusively within its assigned interval.
5. Inference Switching: During sampling, the process dynamically switches between the specialized sub-networks at the computed boundaries.

This modular splitting maximizes sample fidelity without increasing per-step computational requirements, as at each timestep only one sub-network is active.

Empirical Results

CBS is validated across multiple generative settings: latent diffusion (SiT on ImageNet-256), pixel-space diffusion (JiT on ImageNet-64), and unconditional diffusion (UNet on CIFAR-10). Quantitative results establish that complexity-based splitting substantially outperforms both monolithic and uniformly split ensembles in FID, IS, precision, and recall metrics—especially at large model scales and with classifier-free guidance. Improvements are not a result of increased per-step compute: parameter count and FLOPs per step are held constant across all baselines.

Notable findings include:

• FID improvement of up to 35% on large-scale conditional synthesis tasks (SiT-XL, CFG).
• Performance gains hold consistently as the number of sub-networks ($N$) is increased; CBS maintains optimal boundaries even when $N$ scales, unlike heuristic schemes.
• Empirical ablation confirms that CBS-derived boundaries are near-optimal; minor perturbations result in degraded synthesis quality.
• Both spatial (Dirichlet) and geometric (acceleration) monitors achieve near-optimal partitions, with the acceleration metric being marginally superior in practice.

Efficiency of Complexity Estimation

A critical practical aspect is the negligible computational cost of boundary estimation: the complexity profile and temporal splits are robust to the scale and accuracy of the auxiliary model, dataset subset size, and even training duration. The same splits are reproduced whether using a fully trained large network or a lightly trained, reduced model, due to the intrinsic geometric nature of the underlying generative flows.

Implications and Future Directions

CBS provides a rigorous approach to temporal specialization in diffusion models, offering an effective mechanism to decouple global parameter scaling from per-step inference cost. This enables the design of larger, more sample-efficient models without increasing runtime or hardware requirements.

On the theoretical front, CBS is a practical realization of classical approximation bounds for high-dimensional, dynamical neural architectures, bridging theory and implementation. The method could be extended to other axes of specialization: spatial decomposition (e.g., per-token routing), adaptive architecture allocation, and hybrid goal-based partitions. Identifying suitable, tractable monitor functions for such generalized settings is an open challenge likely to drive future research.

Conclusion

"Complexity-Balanced Diffusion Splitting" establishes a search-free, theoretically grounded solution to temporal workload balancing in continuous-time generative modeling. By leveraging function approximation theory and introducing tractable complexity monitors, CBS unlocks significant improvements in synthesis quality while preserving inference efficiency. The framework's flexibility and principled structure suggest broad applicability to future advances in model specialization, neural scaling, and adaptive generative architectures.