# A Deep Generative Approach to Stratified Learning

arXiv: https://arxiv.org/abs/2604.10650 · HTML: https://mlsanity.com/paper/2604.10650/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 56 references that are in the map.
- First posted on arXiv 12 Apr 2026.

## Builds on (most cited first, 17 of 17)

- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Understanding disentangling in β-VAE](https://mlsanity.com/paper/1804.03599/) (2018)
- [The Intrinsic Dimension of Images and Its Impact on Learning](https://mlsanity.com/paper/2104.08894/) (2021)
- [Convergence of denoising diffusion models under the manifold hypothesis](https://mlsanity.com/paper/2208.05314/) (2022)
- [Deep ReLU network approximation of functions on a manifold](https://mlsanity.com/paper/1908.00695/) (2019)
- [Convergence of Diffusion Models Under the Manifold Hypothesis in High-Dimensions](https://mlsanity.com/paper/2409.18804/) (2024)
- [Verifying the Union of Manifolds Hypothesis for Image Data](https://mlsanity.com/paper/2207.02862/) (2022)
- [ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing](https://mlsanity.com/paper/2412.14711/) (2024)
- [Linear Convergence of Diffusion Models Under the Manifold Hypothesis](https://mlsanity.com/paper/2410.09046/) (2024)
- [On gauge freedom, conservativity and intrinsic dimensionality estimation in diffusion models](https://mlsanity.com/paper/2402.03845/) (2024)
- [Improving the Euclidean Diffusion Generation of Manifold Data by Mitigating Score Function Singularity](https://mlsanity.com/paper/2505.09922/) (2025)
- [Token embeddings violate the manifold hypothesis](https://mlsanity.com/paper/2504.01002/) (2025)
- [Unraveling the Localized Latents: Learning Stratified Manifold Structures in LLM Embedding Space with Sparse Mixture-of-Experts](https://mlsanity.com/paper/2502.13577/) (2025)
- [On Convolutions, Intrinsic Dimension, and Diffusion Models](https://mlsanity.com/paper/2506.20705/) (2025)
- [Estimating Distributions with Low-dimensional Structures Using Mixtures of Generative Models](https://mlsanity.com/paper/2301.00890/) (2023)
- [Learning collective variables that preserve transition rates](https://mlsanity.com/paper/2506.01222/) (2025)

## Built on by (most cited first, 2 of 2)

- [Minimax Lower Bound for Estimating Diffusion-based Local Intrinsic Dimension](https://mlsanity.com/paper/2609.04822/) (2026)
- [Nonparametric In-Context Learning under Growing Geometric Complexity: Minimax Optimality and Local Geometry-Adaptivity of Transformers](https://mlsanity.com/paper/2609.31458/) (2026)

## Abstract (arXiv)

While the manifold hypothesis is widely adopted in modern machine learning, complex data is often better modeled as stratified spaces -- unions of manifolds (strata) of varying dimensions. Stratified learning is challenging due to varying dimensionality, intersection singularities, and lack of efficient models in learning the underlying distributions. We provide a deep generative approach to stratified learning by developing two generative frameworks for learning distributions on stratified spaces. The first is a sieve maximum likelihood approach realized via a dimension-aware mixture of variational autoencoders. The second is a diffusion-based framework that explores the score field structure of a mixture. We establish the convergence rates for learning both the ambient and intrinsic distributions, which are shown to be dependent on the intrinsic dimensions and smoothness of the underlying strata. Utilizing the geometry of the score field, we also establish consistency for estimating the intrinsic dimension of each stratum and propose an algorithm that consistently estimates both the number of strata and their dimensions. Theoretical results for both frameworks provide fundamental insights into the interplay of the underlying geometry, the ambient noise level, and deep generative models. Extensive simulations and real dataset applications, such as molecular dynamics, demonstrate the effectiveness of our methods.