Mixture-of-Depth Rectified Flow for Heterogeneous Multimodality

Hierarchical Rectified Flow models data generation through a hierarchy of coupled ordinary differential equations (ODEs), where each depth solves a flow-matching problem over progressively transformed target directions. We show that the ambiguity of the conditional target-direction distribution varies across both time and hierarchy depth, and that increasing depth progressively concentrates this distribution, making the corresponding regression problem simpler. This heterogeneous multimodality implies that a fixed hierarchy depth can allocate computation inefficiently: some space-time regions require deeper hierarchical modeling, while others can be sufficiently modeled at shallow depth. To address this, we propose _Mixture-of-Depth Rectified Flow_ (MoD-RF), which adaptively selects the hierarchy depth at each outer integration step through a sequential stopping mechanism trained from depth-dependent velocity reconstruction quality, without ground-truth depth labels. Under a fixed number of function evaluations (NFEs), MoD-RF allocates deeper computation only where needed while preserving cheaper shallow updates elsewhere. Experiments on synthetic 1D and 2D distributions, MNIST, CIFAR-10, and ImageNet-32 demonstrate improved quality-compute trade-offs over standard rectified flow and fixed-depth hierarchical inference, with particularly strong gains on CIFAR-10 and ImageNet-32.