Complexity-Balanced Diffusion Splitting

2026-06-04

AI Summary

This work addresses the inefficiency of conventional continuous-time generative models, which uniformly allocate computational resources across time despite the time-varying complexity of signals. To overcome this limitation, the authors propose the Complexity-Balanced Segmentation (CBS) framework, which adaptively partitions the diffusion process into intervals of balanced modeling difficulty by leveraging function approximation theory and de Boor’s equidistribution principle. CBS dynamically assigns subnetworks with enhanced representational capacity to high-complexity segments. A key innovation is the introduction of a computable monitoring function based on Dirichlet energy and trajectory acceleration, enabling temporal capacity allocation without heuristic rules or costly search procedures. Coupled with a lightweight auxiliary model to estimate the complexity profile, CBS seamlessly integrates into architectures such as SiT, JiT, and UNet, significantly improving generation quality—e.g., yielding a ~35% FID improvement over naive partitioning in SiT-XL+CFG—without increasing per-step inference cost.

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

Problem

• generative models
• diffusion models
• model efficiency
• temporal partitioning
• computational complexity

Innovation

• Complexity-Balanced Splitting
• temporal capacity allocation
• diffusion generative models
• Dirichlet energy
• trajectory acceleration