A Deep Generative Approach to Stratified Learning

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

• The paper introduces a novel stratified learning framework using a mixture-of-VAEs and score-based diffusion models.
• It derives non-asymptotic convergence rates and error bounds in Wasserstein and Hellinger metrics under varying noise levels.
• The study provides statistically consistent methods for intrinsic dimension and strata recovery, validated on synthetic and molecular datasets.

A Deep Generative Approach to Stratified Learning

Introduction and Motivation

The paper "A Deep Generative Approach to Stratified Learning" (2604.10650) fundamentally advances the theoretical and algorithmic treatment of data distributively organized on stratified spaces, i.e., unions of manifolds with possibly distinct dimensions and intersection singularities. The manifold hypothesis, prevalent in classical machine learning, postulates that high-dimensional data concentrates near a single, smooth, low-dimensional manifold. However, recent empirical and theoretical work—including on LLM embeddings and scientific data—demonstrates that stratified spaces (unions of manifolds with different dimensions) are the correct geometric framework for modeling many real-world domains. The study of stratified learning thus raises intricate statistical and algorithmic challenges: singularities at intersections, spatially varying intrinsic dimension, and identification of underlying generative distributions. Existing geometric and statistical approaches (local Poisson mixtures, persistent homology-based clustering, discriminative deep networks) lack generative sampling capabilities and often fail to robustly recover the intrinsic geometry in the presence of measurement noise.

Generative Modeling on Stratified Spaces: Sieve MLE and Diffusion Approaches

The authors introduce two complementary deep generative frameworks for estimating distributions on stratified spaces:

1. Sieve Maximum Likelihood Estimator (MLE) via Dimension-Aware Mixture-of-VAEs: This approach approximates the data-generating process by a mixture-of-experts generator, with each expert (VAE) modeling samples from a distinct stratum. The ambient distribution is expressed as a convolution of the mixture with isotropic Gaussian noise, capturing realistic sampling/nuisance effects. The generator function $f$ is expressed hierarchically (stratum selection, chart selection, local coordinates) and is theoretically shown to be well-approximated in $L^2$ by sparse deep ReLU networks (details in Lemma~1).
2. Score-Based Diffusion Model Framework: The second framework treats the learning problem via forward-backward time-dependent SDEs (Ornstein-Uhlenbeck process) on the ambient data, where the forward process iteratively injects Gaussian noise, automatically circumventing the need for careful manual regularization. The reverse SDE generates samples by following the vector field given by a deep network-trainable score, estimated by denoising score matching.

The paper systematically derives non-asymptotic convergence rates for ambient (noisy) and intrinsic (deconvolved) distribution estimators in Wasserstein and Hellinger metrics, relating sample complexity to the minimal smoothness and largest intrinsic dimension across all strata. The analysis leverages deep function approximation theory, covering/entropy calculations, and Laplace asymptotics around singularities.

Statistical Guarantees and Regime-Dependent Behavior

Likelihood-Based Mixture-of-Experts Sieve

The sieve MLE estimator is shown to achieve rate-optimal convergence to the true distribution provided the noise level is not vanishing too quickly relative to sample size. Notably, the convergence rate for recovering the intrinsic distribution $Q_\ast$ (in Wasserstein-1) is determined by the worst-case ($d_j / (\alpha_j + 1)$) among all strata:

• For $n$ samples, the optimal rate (neglecting log-factors) is $$n^{-[(\alpha_{j_\ast}+1)-d_{j_\ast}\gamma]/[2(\alpha_{j_\ast}+1) + d_{j_\ast}]}$$, where $d_{j_\ast}$ and $\alpha_{j_\ast}$ are the dimension and density regularity of the bottleneck stratum, and $\gamma$ is the exponent for noise decay. This regime collapses if noise level decays too rapidly, requiring artificial injection of perturbation (data regularization).
• The theoretical construction explicitly parametrizes a generator via locally supported neural network charts, with partition-of-unity gating—a realization not previously implemented in generative stratified learning. Approximation and bracketing entropy bounds are worked out for this neural mixture class.

Score-Based Diffusion Model

The diffusion-based estimator is shown to be minimax rate-optimal for singular measures and enjoys parametric convergence rates $O(n^{-1/2})$ for the ambient distribution in the presence of constant-variance Gaussian noise. In the small-noise or noiseless (singular support) regime, convergence to the intrinsic distribution holds without artificial regularization or noise injection, as denoising is integrated via the forward process.

The analysis details how the structure of the score field adapts to mixture components, and demonstrates that, unlike the likelihood-based estimator, the score-based model is robust to singularity. Crucially, the optimal choice of annealing/early-stopping time for the reverse SDE is deduced analytically from dimensions and smoothness parameters.

Geometric Structure Learning: Intrinsic Dimension and Strata Recovery

A hallmark theoretical contribution is the development of statistically consistent, diffusion-based estimators for both the number of strata ($L^2$0) and the dimension of each stratum, using only sampled data:

• The local score field around regular (non-singular) points aligns with the normal space of the nearby manifold as $L^2$1 (Theorem~\ref{th:normalalign}), and its second-moment exhibits a spectral gap at the codimension (Lemma~\ref{lm:spectral_gap}).
• Applying singular value analysis to collected score vectors after local diffusion steps allows recovery of local dimension, and clustering these dimension estimates infers the number of strata.
• The method’s statistical consistency is proven for regime-appropriate schedules of noise and sample complexity, extending prior work (e.g., [stanczuk2024diffusion]) to the multiple-strata, singular-intersection setting.

Figure 1: Stratified dimension estimation yields correct clustering by local dimension via singular value analysis of the score field.

Experimental Validation and Regime Analysis

Extensive experiments validate the theory on synthetic unions of low-dimensional manifolds with controlled intersections, and on molecular dynamics trajectory spaces. Baseline manifold learning and classical local dimension estimators (Levina-Bickel MLE, Local PCA) quickly degrade in low or moderate ambient noise, whereas the diffusion-based estimator remains accurate even with challenging intersection singularities and provides stratification by intrinsic dimension.

Figure 2: Ground truth, algorithmically classified points, and reconstructions for the mixture of a circle and a 2D plane. The method robustly separates and reconstructs each stratum.

Practical error in Wasserstein-1 between the estimated and ground truth distribution is examined for both mixture-VAEs and diffusion models over varying sample sizes and noise levels. As noise increases, mixture-of-expert VAEs attain lower error up to moderate levels, after which their performance degrades relative to diffusion (which nevertheless models the convolved, not intrinsic, distribution).

Figure 3: Wasserstein-1 distance between learned and true intrinsic distributions versus noise, contrasting diffusion and mixture-of-VAEs. Crossover occurs as noise increases.

Theoretical and Practical Implications

These results provide a rigorous foundation for generative modeling when data geometry is fundamentally stratified, rather than (almost) manifold-constrained. The sieve-MLE (likelihood) approach is appropriate in moderate-noise, deconvolution regimes, while diffusion models are strongly preferable in low-noise or purely singular settings, as they sidestep instability due to sharp concentration near support intersections. The ability to perform automatic, consistent estimation of both the number and local dimensions of strata situates score-based diffusion not only as a generative modeling paradigm, but also as an unsupervised manifold topology learner.

This framework paves the way for practical generative models in applied science (e.g., chemical configuration spaces, complex biological measurements, LLM embedding analyses) where data is stratified and intersection singularities are intrinsic. On the theoretical level, the asymptotic analysis provides minimax-optimal rates under geometric regularity, and clarifies subtle interactions between intrinsic geometry, sample size, and ambient noise.

Conclusion

The paper delivers a comprehensive theoretical and algorithmic foundation for unsupervised, generative stratified learning, unifying advanced sieve likelihood and diffusion-based generative frameworks. The interplay of geometry and noise is precisely elucidated, and statistically consistent methodologies for geometric structure learning are established. Future directions include generalizing geometric and noise assumptions, optimizing practical network architectures for large-component mixtures, and characterizing minimax detection/estimation rates under weaker identifiability constraints. The insights and tools developed here substantially extend generative modeling beyond the manifold setting, offering practical and theoretically sound methodologies for stratified, singular-support data.