Diffusion Models for High-Dimensional Clustered Data: Intrinsic-Dimension Adaptivity via Bayesian Classification

Research Preprint

August 20, 2026

DIFFUSION MODELS FOR HIGH-DIMENSIONAL CLUSTERED DATA: INTRINSIC-DIMENSION ADAPTIVITY VIA BAYESIAN CLASSIFICATION By Yuga Iguchi and Paul Fearnhead

arXiv:2608.19067v1 [stat.ML] 19 Aug 2026

School of Mathematical Sciences, Lancaster University, Lancaster, UK

Abstract The empirical success of diffusion models in generative modelling has motivated theoretical work, including quantitative error bounds and qualitative analyses that characterise the different phases of denoising. We bring these two areas together by studying the adaptivity of diffusion models to the structured geometry of multimodal high-dimensional data that consists of multiple clusters in R𝐷 , each with its own low-dimensional structure, and inter-cluster separation depending on 𝐷. We employ 𝐾-mixture Gaussian distributions as a canonical framework to capture this geometry and establish two theoretical results. First, we interpret denoising as a dynamical Bayesian classifier: the mixture score is a posterior-weighted average of cluster-wise scores, and we show that, with high probability, the posterior class probabilities concentrate on a single cluster once the signalto-noise ratio reaches the scale Θ(log(𝐾 𝐷)/𝐷). Second, by separately analysing the denoising process in its mixing and cluster-commitment phases, we prove that the KL error bound depends linearly on the maximum intrinsic dimension of a cluster, up to a logarithmic factor, even when 𝐾 grows polynomially with 𝐷. This improves on ambient-dimensional bounds and extends existing low-dimensional adaptivity analyses to multimodal distributions with heterogeneous, approximately low-rank covariances.

Key words. Bayesian classification; Diffusion models; Dimensional scaling; Intrinsic dimension; Phase transitions.

1. Introduction Denoising Diffusion Probabilistic Models (DDPMs) [Song and Ermon, 2019, Ho et al., 2020, Song et al., 2020] are general methods for generative modelling that have produced state-of-the-art results across many applications [Dhariwal and Nichol, 2021, Popov et al., 2021, Lugmayr et al., 2022, Trippe et al., 2023]. They define a process that takes data and progressively adds noise to it to produce a realisation from (a close approximation) to some simple distribution. They then learn the reverse, denoising, process, so that we can generate new data examples by denoising realisations from the simple distribution. Relatedly, there has been much interest in understanding why DDPMs are so effective. Much of this work has been theoretical, producing bounds on the accuracy of DDPM. Initial work gave bounds in terms of the Kullback-Leibler (KL) divergence between the true data distribution and the distribution sampled by a DDPM [Benton et al., 2024, Conforti et al., 2025], though these have been extended to other metrics such as total variation distance [Liang et al., 2025, Brešar and Mijatović, 2025], 1-Wasserstein distance [De Bortoli, 2022] or 2-Wasserstein distance [Gao et al., 2025, Silveri and Ocello, 2025, Arsenyan et al., 2026, Koike, 2026]. Within this body of work there has been particular interest in understanding why DDPMs scale well to high-dimensional distributions. Under KL divergence, standard bounds on the accuracy of DDPMs [Benton et al., 2024, Conforti et al., 2025] have terms that, up to a logarithmic factor, scale linearly with the ambient dimension, 𝐷, and suggest that the computational cost, such as the number of time-steps used 1

2

Y. Iguchi and P. Fearnhead

when denoising, needs to scale with 𝐷. While earlier works have established such a bound under minimal conditions on the target, e.g., bounded second moment for data distribution, more recent work has made assumptions that the data lie on a low-dimensional manifold with a bounded support and then shows that in such cases the error can scale with the dimension of the manifold rather than the ambient dimension. For example, Potaptchik et al. [2025] assume that data lies on a bounded manifold of intrinsic dimension 𝑘 and show that the KL-bound is then linear in 𝑘. Li and Yan [2024] define the dimension of the manifold, 𝑘, in terms of metric entropy (which implicitly also implies the data distribution has bounded support) and show that to bound the total variation error the number of time-steps when denoising needs to scale quadratically in 𝑘. This was subsequently improved to linearly in 𝑘 by Li and Yan [2025] and Huang et al. [2026]. Separately, work has focussed on whether DDPMs generalise or just memorise the data [Li et al., 2023, Dupuis et al., 2025, Bonnaire et al., 2025, Maillard and Goldt, 2026]. That is, does simulating from a DDPM produce new examples from the underlying data distribution, or just reproduce members of the data set it was trained on. Many applications of DDPMs involve data distributions with separate clusters, for example image data with images of different types. Biroli et al. [2024] studied the noising/denoising process for such data and was able to give rigorous qualitative insight into its behaviour albeit assuming exact simulation of the denoising process using the true empirical score function. They show that the denoising process can be split into three phases, one where there is general exploration, a second where the process fixes the cluster it will sample from, and then a final stage where the process will be attracted to a specific data point from the training sample. The last stage corresponds to memorisation, and can be avoided by early stopping of the denoiser. Achilli et al. [2026] further extended this study of DDPMs for mixture distributions. 1.1. Our Contributions Our focus is on analysing the geometric adaptivity of DDPMs for target data from a 𝐾-mixture of distributions on R𝐷 , exhibiting both low dimensionality and distinct clustering. This is motivated by the notion of the union of manifolds hypothesis [Brown et al., 2022]: that high-dimensional data such as images lie on a disjoint union of manifolds of varying intrinsic dimensions. Our contributions are summarised as follows: • We rigorously extend the general framework of Biroli et al. [2024] and Achilli et al. [2026] to study a situation where each cluster distribution samples from, or is close to, a low-dimensional manifold. Under the structured geometry specified later, we show that the denoising process fixes on a single cluster to sample from with high probability after a suitably defined signal-to-noise ratio (SNR) reaches a threshold Θ(log(𝐾 𝐷)/𝐷). • Exploiting the above characterisation of cluster commitment in terms of SNR, we give rigorous bounds on the error of DDPM accounting for the various approximations such as time-discretisation of the denoising process, which, under suitable regularity conditions, is shown to scale linearly with the maximum intrinsic dimension of any of the clusters, rather than the ambient dimension 𝐷. As such our work brings together the mixture framework of [Biroli et al., 2024] and Achilli et al. [2026] with the literature on KL-bounds for DDPMs, and the ability of DDPMs to adapt to low-dimensional structure within each distinct cluster. To demonstrate the intrinsic adaptivity in the KL error despite the distinct multimodality, we first characterise the phase transition where the DDPM switches from exploring between clusters to fixing on a single cluster. By leveraging this, and separately analysing the behaviour of the DDPM before and after the phase transition, we are able to show that the KL error bound depends linearly on the intrinsic dimensions up to a logarithmic factor of 𝐷, even when the number of clusters grows polynomially in the ambient dimension 𝐷. A detailed description of the link between the phase transition and the intrinsic dimension adaptivity in KL error is provided in Section 3.1.

Diffusion Models for Clustered Data

3

In order to get our results, we make a number of assumptions on the data distribution. The strongest of which is that each component of our mixture is Gaussian. i.e., Gaussian mixture models (GMMs). This is in line with recent papers that study simplified models in order to gain insight into DDPMs [Li et al., 2026, Oymak and Gulcu, 2021, Shah et al., 2023, Wu et al., 2024, Wang et al., 2025]. In particular, Li et al. [2026] establish dimension-free convergence of DDPM in total variation under the common isotropic covariance structure. In contrast, we consider GMMs with heterogeneous and anisotropic covariances and assume that each cluster is well-separated, in that the squared distances between cluster means scale linearly with 𝐷, while each cluster has its own low-dimensional structure. This inter-cluster separation induces the phase transition described above, leading to our intrinsic-dimension adaptivity result under the heterogeneous and anisotropic covariance geometry. In our work, the intrinsic dimensionality is imposed by an assumption that the spectrum of the cluster covariance matrix decays quickly: for cluster 𝑖, there is some 𝑀𝑖 such that the sum of all but the 𝑀𝑖 -largest eigenvalues is bounded. This is a weaker condition than assuming the data in cluster 𝑖 lies on an 𝑀𝑖 -dimensional manifold. Importantly, we do not require that the support of the data distribution is bounded. We stress that our main result focuses on the dimensional scaling of the noise or time-discretisation schedule for DDPMs or associated reverse SDEs, given a sufficiently accurate estimate of the score, which analytically explains why a reasonable number of steps is enough in practice to sample from high-dimensional clustered data. In contrast, Wang et al. [2025] focuses on score estimation, and clarifies the relation between the training samples and the intrinsic dimensionality to minimise a training loss when the target is a mixture of exactly low-rank Gaussians. This paper is organised as follows. Section 2 summarises the basic set-up of diffusion models and introduces the structured geometry under GMMs. Section 3 presents our two main analytical results on the phase transition and intrinsic-dimension adaptivity in KL error. We provide empirical experiments using real data in Section 4 to demonstrate that our analytical findings on the phase transition can be observed beyond GMMs. We prove the second main result in Section 5 and then conclude in Section 6. √︃Í √ 𝐷 2 Notation. For a vector 𝑥 ∈ R𝐷 and matrix 𝐴 ∈ R𝐷×𝐷 , we write ∥𝑥∥ 𝐴 = 𝑥 ⊤ 𝐴𝑥 and ∥𝑥∥ = 𝑖=1 𝑥 𝑖 . √︁ ⊤ Denote by √︁ 𝜆𝑖 ( 𝐴) the 𝑖-th largest eigenvalue of the matrix A and write ∥ 𝐴∥ 2 = 𝜆1 ( 𝐴 𝐴) and 𝐷 ⊤ ∥ 𝐴∥ 𝐹 = tr( 𝐴 𝐴). For two vectors 𝑥, 𝑦 ∈ R , ⟨𝑥, 𝑦⟩ denotes the inner product. For two matrices 𝐴, 𝐵 ∈ R𝐷×𝐷 , 𝐴 ⪯ 𝐵 mean that 𝐵 − 𝐴 is positive semidefinite. For positive integer 𝑀, we write [𝑀] ≡ {1, . . . , 𝑀 }. For 𝑥 > 0, we write 𝑓 (𝑥) = Θ(𝑥) if there exist constants 0 < 𝑐 1 ≤ 𝑐 2 < ∞ such that 𝑐 1 𝑥 ≤ 𝑓 (𝑥) ≤ 𝑐 2 𝑥. For two sequences {𝑎 𝑛 } 𝑛 and {𝑏 𝑛 } 𝑛 , we write 𝑎 𝑛 = O (𝑏 𝑛 ) if there exists a constant 𝐶 > 0 such that 𝑎 𝑛 ≤ 𝐶𝑏 𝑛 for all sufficiently large 𝑛. We also write 𝑎 𝑛 = 𝑜(𝑏 𝑛 ) if lim𝑛→∞ 𝑎 𝑛 /𝑏 𝑛 = 0. For any two probability measures P and Q with P being absolutely continuous with respect to Q, we write
∫ 𝑑P KL P || Q = log 𝑑Q (𝜔)𝑑P(𝜔) for the KL divergence between them. We write 𝛾 𝐷 : B (R𝐷 ) → [0, 1] ∫ 2 as the standard Gaussian measure defined as 𝛾 𝐷 ( 𝐴) = (2𝜋) −𝐷/2 𝐴 exp(− ∥ 𝑧2∥ )𝑑𝑧, 𝐴 ∈ B (R𝐷 ). Also, for 𝑚 ∈ R𝐷 and 𝑉 ∈ R𝐷×𝐷 , N (𝑚, 𝑉) denotes the normal distribution with mean 𝑚 and covariance 𝑉. For a random variable 𝑋, Law(𝑋) denotes the probability law of 𝑋.

2. Setup 2.1. Score-Based Generative Models We first review score-based generative models [Song et al., 2020], which provide a continuous-time framework for sampling from a target distribution. Let (Ω, F , P) be the probability space where the 𝐷-dimensional standard Brownian motion (𝐵𝑡 )𝑡 ≥0 is equipped, and 𝑃data denote the target distribution on R𝐷 . Denoising diffusion models are designed to generate samples from 𝑃data by reversing a forward noising process, where the latter process can be formulated as the following Ornstein-Uhlenbeck

4

Y. Iguchi and P. Fearnhead

(OU)-type SDE: √ → − → − 𝑑 𝑋 𝑡 = − 𝑋 𝑡 𝑑𝑡 + 2𝑑𝐵𝑡 ,

→ − 𝑋 0 ∼ 𝑃data ,

𝑡 ∈ [0, 𝑇],

−𝑝 : B (R𝐷 ) → [0, 1] the marginal distribution of where 𝑇 > 0 is a fixed time horizon. We denote by → 𝑡 → − 𝑋 𝑡 , given as ∫ → −𝑝 ( 𝐴) = → −𝑝 ( 𝐴|𝑥)𝑃 (𝑑𝑥), 𝐴 ∈ B (R𝐷 ), 𝑡 data 𝑡 |0
− → − → − → − −𝑝 ( 𝐴|𝑥) ≡ P → 𝑋 𝑡 ∈ 𝐴 | 𝑋 0 = 𝑥 . The forward noise 𝑋 𝑡 conditional on the initial data point 𝑋 0 where → 𝑡 |0 → − → − follows a Gaussian distribution, and specifically, the solution is expressed as 𝑋 𝑡 = 𝑚 𝑡 𝑋 0 + 𝜎𝑡 𝑍, 𝑍 ∼ N (0, 𝐼 𝐷 ), where 𝑚 𝑡 and 𝜎𝑡 represent the strength of signal and noise, defined as √︃ (1) 𝑚 𝑡 = exp(−𝑡), 𝜎𝑡 = 1 − 𝑚 2𝑡 . −𝑝 as the associated Lebesgue density. Similarly, we write → −𝑝 With a slight abuse of notation, we write → 𝑡 𝑡 |0 → − for the Lebesgue density of the marginal law of 𝑋 𝑡 if it exists. Then, the denoising SDE is defined as: √ ← − ← − ← −
−𝑝 𝑑 𝑋 𝑡 = 𝑋 𝑡 + 2∇ log→ 𝑇 −𝑡 ( 𝑋 𝑡 ) 𝑑𝑡 + 2𝑑𝐵𝑡 ,

← − → − 𝑋 0 ∼ Law( 𝑋 𝑇 ).

(2)

← − The classical results of [Anderson, 1982, Haussmann and Pardoux, 1986] show that ( 𝑋 𝑡 )𝑡 ∈ [0,𝑇 ] admits a ← − → − ← − → − weak solution so that Law( 𝑋 𝑡 ) = Law( 𝑋 𝑇 −𝑡 ) for any 𝑡 ∈ [0, 𝑇]. In particular Law( 𝑋 𝑇 ) = Law( 𝑋 0 ), so exact simulation of the reverse SDE (2) allows for sampling from the target distribution 𝑃data . DDPMs [Ho et al., 2020] can be viewed as a tractable alternative to (2), which involves three approximations: (i) The use of an estimated score – the intractable score function is replaced with a function 𝑠ˆ : [0, 𝑇] × R𝐷 → R𝐷 parametrised with a neural network, obtained by optimising an appropriate loss function [Yang et al., 2023]; (ii) The use of the Gaussian prior N (0, 𝐼 𝐷 ) for initialisation → − rather than Law( 𝑋 𝑇 ); and (iii) Time-discretisation of the continuous-time process (2). Even with these approximations, it can be difficult to approximately simulate from (2) for 𝑡 ≈ 𝑇 as the ← − −𝑝 drift, specifically the term 2∇ log→ 𝑇 −𝑡 ( 𝑋 𝑡 ), blows-up as 𝑡 → 𝑇. To overcome this, simulation is often performed up to some time 𝑇 − 𝜀, for a suitable choice of early-stopping parameter 𝜀 > 0. It is simple to ← − ← − bound the error between Law( 𝑋 𝑇 − 𝜀 ) and Law( 𝑋 𝑇 ) in terms of 𝜖 [Benton et al., 2024].

2.2. Bayesian Classification Perspective and Our Analytic Framework Given the above setup, our primary goal is to rigorously investigate the generative mechanism and geometric adaptivity of the reverse process, characterised by (2), when 𝑃data is a union of 𝐾 clusters, written in the form 𝑃data ( 𝐴) =

𝐾 ∑︁

[𝑘] 𝜋 𝑘 · 𝑃data ( 𝐴),

𝐾 ≥ 2, 𝐴 ∈ B (R𝐷 ),

𝑘=1

Í [𝑘] where {𝜋 𝑘 } satisfying 𝐾 𝑘=1 𝜋 𝑘 = 1 are the prior weights and 𝑃data denotes the distribution associated with the 𝑘th class. In the mixture setting, the key observation is that the denoising procedure can be viewed as a dynamical Bayesian classification, as advocated also in recent work [Achilli et al., 2026]. −𝑝 More specifically, the score function ∇ log→ 𝑇 −𝑡 (·) contained in the drift of (2) is expressed as a weighted

Diffusion Models for Clustered Data

5

average of the cluster-wise scores, that is, −𝑝 ∇ log→ 𝑇 −𝑡 (𝑥) =

𝐾 ∑︁

−𝑝 [ 𝑘 ] (𝑥), 𝑊𝐷[ 𝑘 ] (𝑡, 𝑥)∇ log→ 𝑇 −𝑡

𝑥 ∈ R𝐷 ,

𝑘=1

− → − −𝑝 [ 𝑘 ] (𝑥) is the Lebesgue density of → where → 𝑋 𝑇 −𝑡 given initial data point 𝑋 0 belonging to the 𝑘-th cluster, 𝑇 −𝑡 i.e., ∫ [𝑘] → −𝑝 [ 𝑘 ] (𝑥) = → −𝑝 𝑇 −𝑡 |0 (𝑥|𝑧) · 𝑃data (𝑑𝑧) 𝑇 −𝑡 and 𝑊𝐷[ 𝑘 ] (𝑡, 𝑥) is the posterior mixture weight or Bayes classifier, i.e., the conditional probability of the → − label taking 𝑘 given the state 𝑋 𝑡 = 𝑥, specifically, −𝑝 [ 𝑘 ] (𝑥) 𝜋𝑘 · → [𝑘] 𝑇 −𝑡 𝑊𝐷 (𝑡, 𝑥) = Í . → − [ℓ ] 𝐾 ℓ=1 𝜋ℓ · 𝑝 𝑇 −𝑡 (𝑥)

(3)

Thus, the reverse process (2) evolves according to the posterior weights given the current denoising state ← − 𝑋 𝑡 , indicating that it implicitly performs a dynamic Bayesian classification within the clustered data setting. The insight presented above suggests that understanding the generative mechanism of diffusion models ← − for clustered data relies on studying the stochastic behaviour of the posterior weights 𝑊𝐷[ 𝑘 ] (𝑡, 𝑋 𝑡 ) for all 𝑡 ∈ [0, 𝑇). In this paper, we focus on Gaussian mixture models (GMMs) as the target data, as they provide a solid analytical framework for rigorously examining the dynamic Bayesian classification. Under this setting, the marginal distribution of the forward noise also follows a Gaussian mixture. Also, cluster-wise low-dimensional structure can be employed within each covariance. Thus, we define [𝑘] 𝑃data ( 𝐴) = N ( 𝐴; 𝜇 𝑘 , Σ 𝑘 ), 𝐴 ∈ B (R𝐷 ) and specify the target distribution as: 𝑃data ( 𝐴) =

𝐾 ∑︁

𝜋 𝑘 · N ( 𝐴; 𝜇 𝑘 , Σ 𝑘 ),

𝐾 ≥ 2, 𝐴 ∈ B (R𝐷 ),

(4)

𝑘=1

where 𝜇 𝑘 ∈ R𝐷 , Σ 𝑘 ∈ R𝐷×𝐷 with Σ 𝑘 ⪰ 0 and N ( 𝐴; 𝜇 𝑘 , Σ 𝑘 ) is a Gaussian measure on R𝐷 defined as
N ( 𝐴; 𝜇 𝑘 , Σ 𝑘 ) ≡ 𝛾 𝐷 {𝑧 ∈ R𝐷 | 𝜇 𝑘 + Σ1/2 𝑘 𝑧 ∈ 𝐴} . The measure admits a well-defined Lebesgue density if Σ 𝑘 is positive definite. Also, note that the −𝑝 [ 𝑘 ] (·), 𝑡 ∈ [0, 𝑇), 𝑘 ∈ [𝐾] follows a Gaussian with distribution of the cluster-wise forward noising → 𝑇 −𝑡 mean 𝑚 𝑇 −𝑡 𝜇 𝑘 and covariance 𝑆𝑇𝑘 −𝑡 given as 𝑆𝑇𝑘 −𝑡 = 𝑚 𝑇2 −𝑡 Σ 𝑘 + 𝜎𝑇2 −𝑡 𝐼 𝐷

(5)

which is positive definite due to the added noise for 𝑡 ∈ [0, 𝑇). We now introduce the following structure for the target measure (4):

Assumption 1 (Prior Weights). There exists a constant 𝐶1 ≥ 0 such that max log 𝜋𝜋𝑖𝑗 ≤ 𝐶1 .

𝑖, 𝑗 ∈ [𝐾 ]

Assumption 2 (Covariance).

(i) Σ𝑖 is symmetric positive semi-definite for any cluster 𝑖 ∈ [𝐾]. Also,

6

Y. Iguchi and P. Fearnhead

there exist a constant 𝐶2 > 0 and 𝛼 ∈ [0, 1) such that for all 𝐷 ∈ N, max ∥Σ𝑖 ∥ 2 ≤ 𝐶2 · 𝐷 𝛼 .

𝑖 ∈ [𝐾 ]

(ii) (Low-dimensional structure) For each cluster 𝑖 ∈ [𝐾], there exist a positive fixed integer 𝑀𝑖 < 𝐷 and 𝜈𝑖 ≥ 0 such that for all 𝐷 ∈ N, we have 𝐷 ∑︁ 𝑘=𝑀𝑖 +1

𝜆 𝑘 (Σ𝑖 ) ≤ 𝜈𝑖 ,

𝜈𝑖 ≤𝑅 𝑀𝑖

(6)

for some universal constant 𝑅 ∈ [0, 1). Assumption 3 (Mean). For any two clusters 𝑖, 𝑗 ∈ [𝐾], there exist constants 𝐶3 , 𝐶4 > 0 such that for all 𝐷 ∈ N, 𝐶3 · 𝐷 ≤ ∥𝜇𝑖 − 𝜇 𝑗 ∥ 2 ≤ 𝐶4 · 𝐷.

(7)

The first assumption is weak as it just bounds the ratio of the mixture weights. The second assumption imposes the structure of each component. Assumption 2(i) allows at most sub-linear dependency on 𝐷 for the maximum eigenvalue of each cluster’s covariance, which is weaker than assuming a uniform bound w.r.t. 𝐷. Sublinear growth is a technical condition that facilitates the phase transition analysis later in Section 3. Assumption 2(ii) together with (i) covers both exact or approximately low-rank structure; if data from component 𝑖 is confined to a manifold of dimension 𝑀𝑖 , then (6) would hold with 𝜈𝑖 = 0. Notably, Assumption 2(ii) is a weaker condition that only requires data from each component to lie close to a low-dimensional manifold, with 𝑅 being a global measure of how close data from each mixture component is to its manifold. Finally, Assumption 3 defines the separation of clusters as a function of the ambient dimension 𝐷. The square distance between the cluster means is assumed to increase linearly with 𝐷. It is also natural for image data where increasing 𝐷 corresponds to observing the same image at increasing resolution. We will also empirically assess these data geometries using high-dimensional datasets (images and RNA sequencing data) in Section 4.

3. Posterior Classification Analysis and Intrinsic Dimension Adaptivity 3.1. Overview This section provides an overview of the generative mechanism of the reverse process within the GMM framework (4) detailed in Section 2.2. We will also explain how the dynamic Bayesian classification perspective can be linked to the KL convergence analysis, thereby demonstrating its ability to adapt to the geometry of the low-dimensional structure within each cluster. ← − Sections 3.2–3.4 analyse the stochastic behaviour of the posterior weights 𝑊𝐷[ 𝑘 ] (𝑡, 𝑋 𝑡 ) given in (3), revealing that the denoising process decomposes into the following two phases: In the Phase I – the mixing phase – several posterior weights remain non-negligible, and the trajectory is affected by global inter-cluster geometry. Then, in Phase II – the commitment phase –, one posterior weight becomes dominant, that is the posterior weight concentrates on a single class, i.e., 𝑊𝐷[ 𝑘 ] ≈ 1 for some 𝑗 ∈ [𝐾] and 𝑊𝐷[𝑖 ] ≈ 0, 𝑖 ∈ [𝐾] \ { 𝑗 }, and the score is effectively reduced to the corresponding cluster-wise score ← − − −𝑝 → − [ 𝑗] ← ∇ log→ 𝑇 −𝑡 ( 𝑋 𝑡 ) ≈ ∇ log 𝑝 𝑇 −𝑡 ( 𝑋 𝑡 ). The analytic results presented in Sections 3.3–3.4 demonstrate a key effect of the ambient dimension 𝐷 on the posterior weight concentration. Roughly, the transition from Phase I to II or mode commitment

Diffusion Models for Clustered Data

7

can occur with probability at least 1 − 𝐷 −𝑞 , 𝑞 ≥ 1 after the signal-to-noise ratio (SNR) 𝑚2

SNR𝑇 −𝑡 ≡ 𝜎𝑇2 −𝑡 ,

𝑡 ∈ [0, 𝑇),

(8)

𝑇 −𝑡

reaches scale log(𝐾 𝐷)/𝐷, see Theorem 1 in Section 3.4. These observations imply that inter-class separation that scales with the ambient dimension 𝐷 facilitates earlier resolution of the cluster component identification, and the rest of the denoising process then focuses on building fine details from the single cluster. Table 1 summarises this two-phase picture. Table 1. Summary of the dynamical regimes of the denoising diffusion for high-dimensional 𝐾-clustered data. Phase I Phase Characterisation

Phase II

Mixing phase

Commitment to a single cluster 𝑗

Signal-to-noise ratio (SNR𝑇 −𝑡 )

O log(𝐾 𝐷)/𝐷

Posterior Weights

𝑊𝐷[ 𝑘 ] ≥ 0, 𝑘 ∈ [𝐾] (Onset of posterior collapse)

Effective Geometry

Inter-cluster inference + local covariances

O (1) [ 𝑗]

𝑊𝐷 ≈ 1, 𝑊𝐷[ 𝑘 ] ≈ 0, 𝑘 ≠ 𝑗 (Posterior concentration) Single local covariance (cluster 𝑗)

Given the above characterisation of the denoising process, we bound the KL error between the target distribution and its approximation realised by DDPM samplers in Section 3.5. The main analytic result therein is that, up to a logarithmic factor, the time-discretisation error component of the KL divergence scales linearly with the intrinsic dimension, i.e., the low-rank property within each cluster, rather than with the ambient dimension 𝐷. To show this we use different approaches to bound the KL error in Phase I and Phase II: For Phase I we use the fact that the KL error involves terms that depend on the SNR, and in this phase the SNR is O (log(𝐾 𝐷)/𝐷), which offsets the inter-cluster separation depending on 𝐷 in the first and second moments. For Phase II, due to cluster-commitment, the denoiser behaves as one targeting a single cluster, so we can bound the error in terms of the cluster’s intrinsic dimension. A more detailed overview of our argument is given in Section 3.6. Remark 1. Similar attempts to split the denoising process into distinct phases are considered in [Biroli et al., 2024, Achilli et al., 2026]. One of the focuses in these works is on characterising when the marginal density of the process transitions from a single-well structure (when the process is contractive) to a multi-well one. They study this in an asymptotic manner, i.e., consider a Taylor series expansion of the score function in very small signal 𝑚 𝑇 −𝑡 ≪ 1 and analyse the curvature of the approximate Hessian. In contrast, our analysis proceeds with a probabilistic characterisation of posterior weights concentration without relying on the Taylor expansion in 𝑚 𝑇 −𝑡 . 3.2. High Probability Bounds for Posterior Weights We begin by stating the following dynamic characterisation of the posterior weights, which serves as a basis to develop the phase transition arguments hereafter: Lemma 3.1. Let 𝑇 ≥ 1 and 𝛿 ∈ (0, 1). Then, for each 𝑖, 𝑗 ∈ [𝐾], it holds that:    [𝑖, 𝑗 ] ← − ← − [ 𝑗] P log 𝑊𝐷[𝑖 ] (𝑡, 𝑋 𝑡 ) ≤ min −𝐶𝐷 (𝑡, 𝛿), 0 | 𝑋 𝑇 ∼ 𝑃data ≥ 1 − 𝛿

8

Y. Iguchi and P. Fearnhead

for each 𝑡 ∈ [0, 𝑇), where we have set: [𝑖, 𝑗 ]

𝐶𝐷

𝜋 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
(𝑡, 𝛿) = log 𝜋𝑖𝑗 + KL → 𝑇 −𝑡 𝑇 −𝑡 √︃ √︃ [𝑖, 𝑗 ] [𝑖, 𝑗 ] [𝑖, 𝑗 ] 2 − 2𝒯𝐷 (𝑡) log 𝛿 − 2𝒰𝐷 (𝑡) log 2𝛿 − 2𝒱𝐷 (𝑡) log 2𝛿 ,

(9)

with 2

[𝑖, 𝑗 ]

(𝑡) = 𝑚 𝑇2 −𝑡 𝜇𝑖 − 𝜇 𝑗 (𝑆 𝑖 ) −1 (𝑆 𝑗 ) (𝑆 𝑖 ) −1 ; 𝑇 −𝑡 𝑇 −𝑡 𝑇 −𝑡  
2 [𝑖, 𝑗 ] 𝑗 𝒰𝐷 (𝑡) = 12 tr 𝐼 𝐷 − 𝑆𝑇 −𝑡 (𝑆𝑇𝑖 −𝑡 ) −1 ;

𝒯𝐷

[𝑖, 𝑗 ]

𝒱𝐷

(𝑡) = ∥𝐼 𝐷 − 𝑆𝑇 −𝑡 (𝑆𝑇𝑖 −𝑡 ) −1 ∥ 2 . 𝑗

The proof of Lemma 3.1 is postponed to Appendix B.1. We further introduce event sets n o  [𝑖, 𝑗 ] ← − 𝐴 [𝑖, 𝑗 ] (𝑡, 𝛿) = 𝜔 ∈ Ω | log 𝑊𝐷[𝑖 ] (𝑡, 𝑋 𝑡 ) ≤ min −𝐶𝐷 (𝑡, 𝛿), 0 , 𝑖, 𝑗 ∈ [𝐾].

(10)

Then, the following result is immediate: Proposition 1. Let 𝑇 ≥ 1, 𝛿 ∈ (0, 1) and 𝑗 ∈ [𝐾]. It holds that:   Ù
← − [ 𝑗] 𝐴 [𝑖, 𝑗 ] 𝑡, 𝐾𝛿 | 𝑋 𝑇 ∼ 𝑃data ≥ 1 − 𝛿, P 𝑖 ∈ [𝐾 ]

for each 𝑡 ∈ [0, 𝑇). ← − Proposition 1 tells us that for denoising trajectories { 𝑋 𝑡 } eventually sampling from the cluster 𝑗 at the ← − terminal time 𝑇, the posterior weights 𝑊𝐷[𝑖 ] (𝑡, 𝑋 𝑡 ) with any cluster index 𝑖 rather than 𝑗 are bounded by 
[𝑖, 𝑗 ] min exp −𝐶𝐷 (𝑡, 𝛿/𝐾) , 1 with probability at least 1 − 𝛿, 𝛿 ∈ (0, 1). Proof of Proposition 1. We have that:  Ù   Ø  
𝑐 ←
𝑐 ← − − [ 𝑗] [ 𝑗] P 𝐴 [𝑖, 𝑗 ] 𝑡, 𝐾𝛿 | 𝑋 𝑇 ∼ 𝑃data = P 𝐴 [𝑖, 𝑗 ] 𝑡, 𝐾𝛿 | 𝑋 𝑇 ∼ 𝑃data 𝑖 ∈ [𝐾 ]

≤

𝑖 ∈ [𝐾 ]

∑︁

 ∑︁
𝑐 ← − [ 𝑗] 𝛿 P 𝐴 [𝑖, 𝑗 ] 𝑡, 𝐾𝛿 | 𝑋 𝑇 ∼ 𝑃data ≤ 𝐾 = 𝛿, 

𝑖 ∈ [𝐾 ]

𝑖 ∈ [𝐾 ]

where we used Lemma 3.1 in the last line. The proof is complete. 3.3. Posterior Weights in High-Dimensional Limit Proposition 1 indicates that the decay of the posterior weights is dynamically controlled by how large the −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
is compared to the prior ratio and the fluctuation terms. The following result signal KL → 𝑇 −𝑡 𝑇 −𝑡 [𝑖, 𝑗 ] demonstrates how the signal component dominates in the term 𝐶𝐷 in the high-ambient-dimensional regime, 𝐷 → ∞. Proposition 2. Let 𝑇 ≥ 1 and 𝑗 ∈ [𝐾]. Also, let Assumptions 1, 2, 3 hold and assume that for each 𝑡 ∈ [0, 𝑇) and 𝑖, 𝑗 ∈ [𝐾], the following limit exists: −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
≡ 𝐶 e∞[𝑖, 𝑗 ] (𝑡) ≥ 0. lim 𝐷1 × KL → 𝑇 −𝑡 𝑇 −𝑡

𝐷→∞

Diffusion Models for Clustered Data

9

 ← − [ 𝑗] Then, we have that: conditionally on the event 𝜔 ∈ Ω | 𝑋 𝑇 ∼ 𝑃data , ← − e∞[𝑖, 𝑗 ] (𝑡), lim sup 𝐷1 log 𝑊𝐷[𝑖 ] (𝑡, 𝑋 𝑡 ) ≤ − 𝐶

∀𝑖 ∈ [𝐾],

𝐷→∞

holds true almost surely for each 𝑡 ∈ [0, 𝑇).

Proposition 2 implies that in the high ambient dimensional limit 𝐷 → ∞, the decay of the posterior weight is mainly governed by the KL divergence between pairs of noised cluster distributions. The data geometry stated in Assumptions 1, 2 and 3 ensure that the KL divergence (signal) eventually dominates the fluctuation terms (noise) specified in (9) as 𝐷 → ∞. If the clusters in the target distribution are well-separated, e.g., in the sense of the first or second moment, then the signal (KL divergence) is also expected to increase as 𝑡 approaches the terminal time 𝑇. That is, for small 𝑡 ∈ [0, 𝑇 − 𝜀], the noised clusters are nearly indistinguishable, but the gap gradually increases as the denoising proceeds, where those differences are measured by the KL divergence. In the early stage of denoising, the posterior weights are basically similar to the prior weights, but, for some classes, these posterior weights start to degenerate toward 0, followed by a rapid posterior concentration for a single cluster 𝑗. In particular, such −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
reaches a a rapid cluster commitment occurs with an exponential rate after the signal KL → 𝑇 −𝑡 𝑇 −𝑡 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
+ 𝑜(𝐷) takes a sufficiently large negative value. certain level so that −KL → 𝑇 −𝑡 𝑇 −𝑡 Proof of Proposition 2. Set 𝛿 = 𝐷 −𝑞 , 𝑞 ≥ 2. Recall the definition of the event 𝐴 [𝑖, 𝑗 ] in (10) and define
Ñ ← − [ 𝑗] [ 𝑗] 𝐸 𝐷 (𝑡, 𝛿) = 𝑖 ∈ [𝐾 ] 𝐴 [𝑖, 𝑗 ] 𝑡, 𝐾𝛿 and 𝐵 [ 𝑗 ] = {𝜔 ∈ Ω | 𝑋 𝑇 ∼ 𝑃data }. Then, Proposition 1 yields ∞ ∞   ∑︁ ∑︁
𝑐 [ 𝑗] P 𝐸 𝐷 (𝑡, 𝐷 −𝑞 ) | 𝐵 [ 𝑗 ] ≤ 𝐷 −𝑞 < ∞. 𝐷=1

𝐷=1

Thus, the first Borel-Cantelli lemma gives that:   [ 𝑗] P lim inf 𝐸 𝐷 (𝑡, 𝐷 −𝑞 )|𝐵 [ 𝑗 ] = 1, 𝐷→∞

that is, for P(·|𝐵 [ 𝑗 ] ) almost every 𝜔, there exists 𝐷 † ∈ N such that for all 𝐷 ≥ 𝐷 † we have ← − [𝑖, 𝑗 ] log 𝑊𝐷[𝑖 ] (𝑡, 𝑋 𝑡 ) + 𝐶𝐷 (𝑡, 𝐾𝛿 ) ≤ 0

(11)

for all 𝑖 ∈ [𝐾]. Dividing the both sides of (11) by 𝐷 and taking lim sup𝐷→∞ gives ← − [𝑖, 𝑗 ] lim sup 𝐷1 log 𝑊𝐷[𝑖 ] (𝑡, 𝑋 𝑡 ) ≤ − lim sup 𝐷1 𝐶𝐷 (𝑡, 𝐾𝛿 ). 𝐷→∞

𝐷→∞

The proof completes once we show that for 𝛿 = 𝐷 −𝑞 , 𝑞 ≥ 2, [𝑖, 𝑗 ]

lim sup 𝐷1 𝐶𝐷

[𝑖, 𝑗 ]

e∞ (𝑡, 𝐾𝛿 ) = 𝐶

(𝑡),

(12)

𝐷→∞

where all the fluctuation terms in (9) multiplied by 1/𝐷 converges to 0 as 𝐷 → ∞. To this end, we exploit the following upper bounds for the fluctuation terms in (9) whose proof is contained in Appendix D:

10

Y. Iguchi and P. Fearnhead

Lemma 3.2. It holds that for each 𝑡 ∈ [0, 𝑇) and all 𝑖, 𝑗 ∈ [𝐾], [𝑖, 𝑗 ]

𝒯𝐷

(𝑡) ≤ 1 + SNR𝑇 −𝑡 · ∥Σ𝑖 − Σ 𝑗 ∥ 2 · 𝑚 𝑇2 −𝑡 ∥𝜇𝑖 − 𝜇 𝑗 ∥ 2 𝑗

(𝑆𝑇 −𝑡 ) −1

;

(13)

2

[𝑖, 𝑗 ]

(𝑡) ≤ 21 SNR𝑇2 −𝑡 Σ𝑖 − Σ 𝑗 𝐹 ;

(14)

[𝑖, 𝑗 ]

(𝑡) ≤ SNR𝑇 −𝑡 · ∥Σ𝑖 − Σ 𝑗 ∥ 2 .

(15)

𝒰𝐷 𝒱𝐷

[𝑖, 𝑗 ]

Lemma 3.2 with Assumptions 1 – 3 yields that: for each 𝑡 ∈ [0, 𝑇) with a fixed 𝑇 ≥ 1, 𝒯𝐷 (𝑡) = [𝑖, 𝑗 ] [𝑖, 𝑗 ] O (𝐷 𝛼+1 ), 𝒰𝐷 (𝑡) = O (𝐷 2𝛼 ) and 𝒱𝐷 (𝑡) = O (𝐷 𝛼 ), leading to the desired equation (12), where we 2 2𝛼 used ∥Σ𝑖 − Σ 𝑗 ∥ 𝐹 = O (𝐷 ) shown in (75). The proof of Proposition 2 is now complete. 3.4. Posterior Concentration and the Scale of SNR We hereby examine the time scale at which the posterior weight decay, or concentration, occurs. Inspired by the high-dimensional asymptotics presented in Proposition 2, we analytically identify the point in time after which the main signal consistently outweighs other components in equation (9), such as the ratio of prior weights and fluctuation terms. Specifically, by setting a fixed early stopping criterion 𝜀 ∈ (0, 1), we [𝑖, 𝑗 ] −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
for all 𝑡 ≥ 𝑡 † , for a fixed consider the time 𝑡 † ∈ [0, 𝑇 − 𝜀] such that 𝐶𝐷 (𝑡, 𝛿/𝐾) ≥ 𝜂KL → 𝑇 −𝑡 𝑇 −𝑡 constant 𝜂 ∈ (0, 1). To analytically characterise the scale of 𝑡 † , we introduce an additional condition: Assumption 4. There exists a constant 𝐶5 > 0 such that for all 𝑖, 𝑗 ∈ [𝐾], max ⟨𝜇𝑖 − 𝜇 𝑗 , 𝑢 𝑘 ( 𝑗)⟩ ≤ 𝐶5 ,

(16)

𝑘 ∈ [𝐷 ]

where 𝑢 𝑘 ( 𝑗) is the eigenvector associated with the eigenvalue 𝜆 𝑘 (Σ 𝑗 ). Assumption 4 states that the projection of the mean difference vector between clusters onto the eigenvectors of one cluster does not concentrate on a few principal components. We then have the following result, whose proof is postponed to Section B.2: Theorem 1 (SNR Scale for Posterior Concentration). Let 𝑇 ≥ 1 and 𝜀 ∈ (0, 1) be a fixed early stopping. Also set 𝜂 ∈ (0, 1) and 𝑞 ≥ 2. Let Assumptions 1, 2, 3, 4 hold. Then, we have the following: There exists 𝐷 † = 𝐷 † (𝜂, 𝑞, 𝜀) such that for any integer 𝐷 ≥ 𝐷 † ,  Ù n   o ← ← − − [ 𝑗] [ 𝑗] [𝑖 ] [𝑖 ]
→ − → − P 𝑊𝐷 (𝑡, 𝑋 𝑡 ) ≤ exp −𝜂KL 𝑝 𝑇 −𝑡 || 𝑝 𝑇 −𝑡 | 𝑋 𝑇 ∼ 𝑃data ≥ 1 − 𝐷 −𝑞 , 𝑗 ∈ [𝐾], (17) 𝑖 ∈ [𝐾 ]

h i h i   𝑆 for each 𝑡 ∈ max 0, 𝑇 + 12 log 1+𝑆 , 𝑇 − 𝜀 , equivalently SNR𝑇 −𝑡 ∈ max SNR𝑇 , 𝑆 , SNR 𝜀 , where 𝑆 ≡ 𝑆(𝐷, 𝑞, 𝜂) satisfies:   log(𝐾 𝐷) 𝑆=Θ . (18) 𝐷 Besides, for any 𝛽 ∈ (0, 1),  Ù   [𝑖 ] ← − ← − [ 𝑗] P 𝑊𝐷 (𝑡, 𝑋 𝑡 ) ≤ 𝛽 | 𝑋 𝑇 ∼ 𝑃data ≥ 1 − 𝐷 −𝑞 ,

𝑗 ∈ [𝐾],

𝑖 ∈ [𝐾 ] 𝑖≠ 𝑗

whenever the reverse time 𝑡 satisfies SNR𝑇 −𝑡 ≥ 𝐿 ·

 max log(𝐾 𝐷), log 𝛽1 𝐷

(19)

Diffusion Models for Clustered Data

11

for some constant 𝐿 > 0 independent of 𝐷. Theorem 1 establishes an SNR scaling law for posterior concentration. That is, for trajectories eventually sampling from cluster 𝑗, the posterior weights assigned to all other classes are sufficiently
small with high probability once the SNR reaches a threshold Θ log(𝐾 𝐷)/𝐷 . In the proof, Assumption −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
, 3 (mean separation ) together with Assumption 4 plays a role in quantifying the signal KL → 𝑇 −𝑡 𝑇 −𝑡 leading to the obtained scale of the SNR threshold 𝑆. The above result also gives rise to the following key implications on the appropriate scaling of 𝑇 with respect to the ambient dimension 𝐷 and its effect on the length of the mixing phase:  𝑆 Corollary 1. Consider the same setting as the statement in Theorem 1 and let 𝑡 † = max 0, 𝑇 + 12 log 1+𝑆 . Then, we have the following: (i). If the time horizon 𝑇 > 0 is a fixed constant, then 𝑡 † → 0 as 𝐷 → ∞, i.e., (17) holds for each 𝑡 ∈ [0, 𝑇 − 𝜀]. (ii). If 𝑇 ≥ 12 log 𝐷, then 𝑡 † ≳ log log 𝐷. Thus, the threshold 𝑡 † does not collapse to 0 as 𝐷 → ∞, i.e., the mixing phase can remain non-trivial. We remark that for 𝑇 ≥ 12 log 𝐷, we have SNR𝑇 ≤ 2/𝐷, 𝐷 ≥ 2, and so SNR𝑇 ≤ S for any sufficiently −𝑝 and large 𝐷. The choice of time interval 𝑇 ≥ 12 log 𝐷 is also reasonable in reducing the gap between → 𝑇 N (0, 𝐼 𝐷 ) so that the forward noising distribution at time 𝑇 is sufficiently close to the normal distribution. See also Theorem 2 below. Proof of Corollary 1. The first statement (i) is immediate by noticing that 𝑇 + 12 log 1+𝑆 → −∞ as 𝑆

𝐷 → ∞ given (18). For (ii), 𝑇 ≥ 12 log 𝐷 with (18) gives 𝑇 + 12 log 1+𝑆 ≥ 12 log 1+𝑆 ≳ log log 𝐷 for any sufficiently large 𝐷. The proof is now complete. 𝑆

𝐷𝑆

3.5. KL Convergence Analysis – Intrinsic Dimension Adaptivity via Posterior Concentration Given the above arguments on the dynamical regimes, we now argue how the time-discretisation error scales with the ambient dimension 𝐷 and the intrinsic dimensionality 𝑀𝑖 involved in each cluster specified in Assumption 2. Let 𝜀 ∈ (0, 1) be a fixed early stopping. We write the time-discretisation schedule {𝑡 𝑖 }𝑖 ∈ [ 𝑁 ] satisfying 0 = 𝑡0 < 𝑡1 < · · · < 𝑡 𝑁 = 𝑇 − 𝜀 with 𝑇 ≥ 1 and the number of steps 𝑁 ∈ N. We define the time-discretised denoising sampler : ← ← − √ ← − − ← −  𝑇 −𝑡 𝜇(𝑇 ˆ − 𝑡 , 𝑌 ) 𝑑𝑡 + 𝑑 𝑌 𝑡 = 𝑌 𝑡 − 2 𝜎𝑌2 𝑡 + 2 𝑚 2𝑑𝐵𝑡 , 𝑖 𝑡 2 𝑖 𝜎 𝑇 −𝑡

𝑡 ∈ [𝑡 𝑖 , 𝑡𝑖+1 ],

𝑇 −𝑡

← − → − → − with 𝑌 0 ∼ N (0, 𝐼 𝐷 ), where 𝜇ˆ 𝑇 −𝑡𝑖 (𝑦) is an estimator of the true posterior mean 𝜇𝑇 −𝑡𝑖 (𝑦) = E[ 𝑋 0 | 𝑋 𝑇 −𝑡𝑖 = → − ← − → − ← − 𝑦]. We write P 𝑡 and Q 𝑡𝑖 as the law of 𝑋 𝑡 and 𝑌 𝑡𝑖 , respectively and study the upper bound for → − ← −
KL P 𝜀 || Q 𝑡𝑁 . We then need to decide the schedule for the time discretisation. In this work, we adopt the following clear and simple update rule for grid points: for a precision parameter 𝜏 ∈ (0, 1/4], 𝑡0 = 0;  𝑡 𝑗+1 = min 𝑡 𝑗 + 𝜏𝜎𝑇2 −𝑡 𝑗 , 𝑇 − 𝜀 ,

𝑗 ∈ {0, . . . , 𝑁 − 1}

(20)

where 𝑁 is the minimum integer s.t. 𝑡 𝑁 −1 + 𝜏𝜎𝑇2 −𝑡𝑁 −1 ≥ 𝑇 − 𝜀. The design of the schedule is determined by the local noise 𝜎𝑇2 −𝑡 𝑗 , but is independent of the geometry of multimodal distributions. Let ℎ 𝑗+1 = 𝑡 𝑗+1 − 𝑡 𝑗 , 𝑗 = 0, 1, . . . , 𝑁 − 1, then this schedule has two properties. First, we have ℎ 𝑗+1 ≤ 2𝜏 min{1, 𝑇 − 𝑡 𝑗 }, which aligns with the exponential-decreasing-type schedules often employed in convergence analyses; see, e.g., Benton et al. [2024], Potaptchik et al. [2025], Huang et al. [2026]. The

12

Y. Iguchi and P. Fearnhead

second property is that the number of time-steps 𝑁 increases as 1/𝜏, as shown in the following lemma, whose proof is postponed to Appendix A: Lemma 3.3. For a given precision parameter 𝜏 ∈ (0, 1/4], the number of steps 𝑁 specified in the time-discretisation schedule (20) satisfies 
 (21) 𝑁 = Θ 1 + 𝜏1 𝑇 − 1 + log 1𝜀 . To state our main result, we introduce the following assumption regarding the accuracy of score estimation. Assumption 5 (Score Estimate Accuracy). There exists 𝜀score > 0 such that 𝑁 −1   1 ∑︁ ← − ← − 2 ℎ 𝑗+1 × E ∥∇ log 𝑝 𝑇 −𝑡 𝑗 ( 𝑋 𝑡 𝑗 ) − 𝑠ˆ (𝑇 − 𝑡 𝑗 , 𝑋 𝑡 𝑗 ) ∥ 2 ≤ 𝜀score , 𝑇 𝑗=0

where 𝑠ˆ (𝑇 − 𝑡 𝑗 , 𝑥) ≡ −

𝑥 − 𝑚 𝑇 −𝑡 𝑗 𝜇(𝑇 ˆ − 𝑡 𝑗 , 𝑥) 𝜎𝑇2 −𝑡 𝑗

,

𝑥 ∈ R𝐷 .

(22)

Treating the accuracy as a black box, as per this assumption, is commonly adopted in diffusion model convergence analysis [Benton et al., 2024, Potaptchik et al., 2025, Huang et al., 2026]. We now have the following KL convergence guarantee:  Theorem 2 (KL Convergence under Dynamical Regimes). Let 𝑇 ≥ max 1, 12 log 𝐷 and 𝜀 ∈ (0, 1) be a fixed early stopping parameter. Consider the time-discretisation schedule (20) with 𝜏 ∈ (0, 1/4]. Let Assumptions 1, 2, 3, 4 and 5 hold. Also, assume E 𝑃data [∥ 𝑋 ∥ 2 ] < ∞, 𝐾 ≲ 𝐷 𝛾 for a fixed 𝛾 > 0 and 𝜀 ≳ 𝑅/𝑇, where 𝑅 ∈ [0, 1) is found in Assumption 2. Set 𝑀 † = max 𝑘 ∈ [𝐾 ] 𝑀𝑘 . Then, there exists 𝐷 † ∈ N → − ← − such that for all 𝐷 ≥ 𝐷 † , we have KL( P 𝜀 || Q 𝑡𝑁 ) ≲ E1 + E2 + E3 , where  
2 E1 ≡ 𝜏𝑀 † 𝑇 + log 1𝜀 , , E3 ≡ 𝐷 + E 𝑃data [∥ 𝑋 ∥ 2 ] 𝑒 −2𝑇 . (23) E2 ≡ 𝑇 𝜀score The proof of Theorem 2 is postponed to Section 5. The bound (23) shows the three sources of errors – E1 : time-discretisation error; E2 : score-estimation error; and E3 : prior gap – which aligns well with the standard theoretical results in the literature. The main highlight is that the discretisation error scales with the intrinsic dimension 𝑀 † , rather than the ambient dimension 𝐷, despite the O (𝐷) inter-class separation, while the number of mixture components 𝐾 is allowed to grow with 𝐷 polynomially. Moreover, combining (21) with Theorem 2 gives the following complexity bound: 
Corollary 2. Consider the setting of Theorem 2 and set 𝐴 ≡ 𝑀 † 𝑇 + log 1𝜀 . For any 𝜒 ∈ 0, 𝐴/4 , choosing 𝜏 = 𝜒/𝐴 ∈ (0, 1/4] leads to E1 = 𝜒, where the number of steps 𝑁 satisfies   2 𝑀† 1 𝑁 = O 𝜒 𝑇 + log 𝜀 . Remark 2. The requirement for the early stopping criterion 𝜀 ≳ 𝑅/𝑇 gives insight into the role of 𝜀 and its relation to the geometry of the data. When 𝜈 † = max 𝑘 ∈ [𝐾 ] 𝜈 𝑘 is exactly 0 (each cluster’s covariance is exactly low-rank) or nearly 0, i.e., data within clusters concentrates on each thin subspace, then the corresponding 𝑅 is exactly or nearly 0, and one can choose a very small 𝜀. If 𝑅 is large, i.e., clusters no longer have a low-dimensional structure, then one instead obtains an error bound depending on 𝐷 linearly without relying on such a requirement for 𝜀.

Diffusion Models for Clustered Data

13

3.6. Overview of the Proof of Theorem 2 We briefly outline the main arguments for the proof of Theorem 2 and highlight how the phase transition analysis is employed to obtain the intrinsic-dimension (𝑀 † ) scaling in the time-discretisation error, even when the number of clusters 𝐾 grows polynomially with 𝐷. The proof consists of five main steps. The first step in Section 5.1 uses standard arguments [Benton → − ← − et al., 2024, Conforti et al., 2025, Huang et al., 2026] to bound KL( P 𝜀 || Q 𝑡𝑁 ) as the sum of three terms that correspond to the error from approximating the score, the error in how we initiate the denoiser and the time-discretisation error. The main novelty in our proof is how we bound the time-discretisation error, which is covered in Steps 2 to 4 and outlined below. Step 5 then collects together all the error terms, and obtains the simplified bound presented in (23). In Step 2 (see Section 5.2) we provide an upper bound on the time-discretisation error involving the following term: 𝑁 −1   h ∑︁
i → − 𝜏 SNR𝑇 −𝑡 𝑗+1 − SNR𝑇 −𝑡 𝑗 · E tr Σ𝑇 −𝑡 𝑗 ( 𝑋 𝑇 −𝑡 𝑗 ) , 𝑗=1

→ − → − where Σ𝑇 −𝑡 (𝑦) = Cov[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦], (𝑡, 𝑦) ∈ [0, 𝑇) × R𝐷 , is the posterior covariance. Under the GMM framework, the posterior covariance can be decomposed as shown by the following lemma. Lemma 3.4. It holds that for all (𝑡, 𝑦) ∈ [0, 𝑇) × R𝐷 , Σ𝑇 −𝑡 (𝑦) = Σ𝑇𝐼 −𝑡 (𝑦) + Σ𝑇𝐼 𝐼−𝑡 (𝑦) with Σ𝑇𝐼 −𝑡 (𝑦) =

𝐾 ∑︁

𝑊𝐷[ 𝑘 ] (𝑡, 𝑦)𝜎𝑇2 −𝑡 Σ 𝑘 (𝑆𝑇𝑘 −𝑡 ) −1 ;

𝑘=1

1 Σ𝑇𝐼 𝐼−𝑡 (𝑦) = 2

(24) ∑︁

⊗2 [ 𝑗] 𝑊𝐷 (𝑡, 𝑦)𝑊𝐷[ 𝑘 ] (𝑡, 𝑦) 𝜇 [ 𝑗 ] (𝑡, 𝑦) − 𝜇 [ 𝑘 ] (𝑡, 𝑦) ,

𝑗,𝑘 ∈ [𝐾 ]

where we have set 𝜎2 → − → − → − [ 𝑗] −𝑝 [ 𝑗 ] (𝑦) + 𝑦 . 𝜇 [ 𝑗 ] (𝑡, 𝑦) = E[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦, 𝑋 0 ∼ 𝑃data ] = 𝑇 −𝑡 ∇ log→ 𝑇 −𝑡 𝑚 𝑇 −𝑡 𝑚 𝑇 −𝑡

(25)

The proof of Lemma 3.4 is provided in Appendix E.1. The term Σ𝑇𝐼 −𝑡 captures the local within-cluster geometry, whereas Σ𝑇𝐼 𝐼−𝑡 reflects the between-cluster geometry and may therefore scale with 𝐷 through differences between components. 
 → − In Section 5.3 (Step 3), we show the upper bound on the term E tr Σ𝑇𝐼 −𝑡 𝑗 ( 𝑋 𝑇 −𝑡 𝑗 ) depends linearly 
 → − on 𝑀 † but is independent of 𝐾. Section 5.4 (Step 4) deals with the inter-class term E tr Σ𝑇𝐼 𝐼−𝑡 𝑗 ( 𝑋 𝑇 −𝑡 𝑗 ) by exploiting the phase transition. Roughly, we divide the time interval [0, 𝑇 − 𝜀] into two regimes
according to whether SNR𝑇 −𝑡 𝑗 ≤ 𝐿 log(𝐾 𝐷) /𝐷 for an arbitrary fixed constant 𝐿 > 0. During the first regime (mixing phase), the small SNR offsets the mismatches in the first and second moments that can grow polynomially with 𝐷. For the latter regime (cluster commitment phase), Theorem 1 implies that, conditional on a trajectory terminating in a cluster 𝑘, at least one of the two posterior weights → − → − [ 𝑗] 𝑊𝐷[𝑖 ] (𝑡 𝑗 , 𝑋 𝑇 −𝑡 𝑗 ), 𝑊𝐷 (𝑡 𝑗 , 𝑋 𝑇 −𝑡 𝑗 ), 𝑖 ≠ 𝑗, is exponentially small in 𝐷 with high probability, thereby suppressing the inter-class difference. Furthermore, the dependence on 𝐾 is shown to be logarithmic via a finite exponential-moment bound for the maximum over the Gaussian components.

4. Numerical Experiments of Posterior Commitments in Real Data So far, we have studied the intrinsic dimension adaptivity of diffusion models for target distributions formulated as a union of distinct (linear) manifolds within a Gaussian mixture framework. The key ingredient for showing such adaptivity in the KL error analysis (Theorem 2) is that the scale of the SNR

14

Y. Iguchi and P. Fearnhead

for posterior weights concentration is approximately Θ(1/𝐷), as shown in Theorem 1 under cluster separation. This section provides empirical evidence that this SNR scale at posterior concentration can also be observed for high-dimensional data beyond the GMM setting. All technical details of the implementation can be found in Appendix F, and the code to reproduce the experiments is available at https://github.com/YugaIgu/diffusion_models_phase_transition 4.1. Set-Up We consider two types of high-dimensional data, which are typical examples of the application of diffusion models: (a) Single cell RNA sequencing (scRNA-seq) of peripheral blood mononuclear cells (PBMCs) with 32738 genes and (b) Image data (STL-10) consisting of a resolution of 96 by 96 pixels with 3 colour channels (RGB). For dataset (a), we vary the ambient dimension 𝐷 by retaining the top 𝐷 ∈ {500, 1000, 2000, 4000, 8000} highly variable genes (HVGs). For the STL-10 image data, we construct lower-resolution versions to obtain three resolutions: 𝐷 = 32 × 32 × 3, 𝐷 = 64 × 64 × 3, and 𝐷 = 96 × 96 × 3 (raw data), by using bilinear interpolation, implemented in PyTorch. Then, we select three classes (labels), specifically, CPBMC = {CD4 T, CD14+ Monocytes, CD19+ B} and Cimage = {dog, cat, airplane} for datasets (a) and (b). We simulate the denoising SDE to sample from the target clusters. We directly use the score of the noised empirical distribution rather than a trained neural-network score, in the spirit of the exact-empirical-score setting considered in Biroli et al. [2024]. This allows us to isolate the dynamical phase transition from the uncertainty due to score estimation error. We perform the following four steps to observe the empirical SNR at the posterior weight concentration: I. Simulating trajectories. We set the terminal time as 𝑇 = 6 and simulate 𝑀 = 500 independent backward trajectories { 𝑋¯ 𝑡[𝑗𝑘 ] }1≤ 𝑗 ≤𝑛,1≤ 𝑘 ≤ 𝑀 using the Euler-Maruyama discretisation with a uniform time-step ℎ = 𝑡 𝑗+1 − 𝑡 𝑗 = 0.003 and so the number of steps 𝑛 = 𝑇/ℎ = 2000. These are used as numerical approximations for the exact backward process. II. Labelling. For each trajectory, we assign a label of the class from which the path eventually samples, which is obtained as the argument of the maximum of the posterior weight (3): 𝐶 [ 𝑘 ] ≡ argmax𝑐∈ C 𝑊𝐷[𝑐] (𝑡 𝐽 , 𝑋¯ 𝑡[𝐽𝑘 ] ), 𝑘 = 1, . . . , 𝑀, where 𝐽 = 𝑛 − 10 (to avoid numerical instability) and C is CPBMC or Cimage . III. Sorting. Those paths are sorted within the three distinct classes as X𝑐 =



𝑋¯ 𝑡[𝑗𝑘 ]

𝑗=1,..., 𝐽

| 𝐶 [𝑘] =

𝑐, 𝑘 = 1, . . . , 𝑀 , 𝑐 ∈ C, that is, the set of paths eventually sampling from the class 𝑐. IV. Recording. We then record the values of three types of posterior weights 𝑊𝐷[𝑐] , 𝑐 ∈ C, evaluated with † the trajectories contained in X𝑐 for some target class 𝑐† ∈ C at each time-step 𝑗 = 1, . . . , 𝐽. In the above procedure, the prior weights are assumed to be equal across the classes, and the (intractable) posterior weights are estimated by plugging an empirical approximation of the forward noised density into the definition of the weight (3). 4.2. Results We first present the data geometry for the datasets (a) and (b) in Tables 2 and 3. Table 2 indicates that for the STL-10 datasets the 90% effective rank is stable for varying ambient dimension 𝐷, while the maximum eigenvalue 𝜆max seems to grow approximately linearly with 𝐷. On the other hand, 𝜆max for PBMC data exhibits O (𝐷 𝛼 ) growth with 𝛼 < 1. We observe from Table 3 that the mean separation of STL-10 datasets scales linearly with 𝐷 and that covariance separations are very clear. The pairwise geometry for PBMC data also exhibits a similar separation growth as 𝐷 increases – presented in Appendix F.

Diffusion Models for Clustered Data

15

Table 2. Empirical covariance geometry for PBMC and STL-10 across ambient dimensions. For each class, we report the maximum covariance eigenvalue (𝜆max ), the trace (tr(b Σ)), and the 90% effective rank (𝑟 eff,90 ), computed from the empirical data. PBMC Class

𝐷

𝜆 max

tr(b Σ)

𝑟 eff,90

CD4 T

500 1000 2000 4000 8000

3.42 5.41 9.85 16.30 25.72

469.88 943.59 1807.21 3600.29 7465.22

313 507 686 817 903

CD14+ Mono.

500 1000 2000 4000 8000

8.41 14.14 23.32 36.64 57.85

485.49 985.57 2057.65 4001.79 7781.14

209 289 339 372 395

CD19+ B

500 1000 2000 4000 8000

9.03 13.03 22.68 37.03 55.90

554.19 1050.05 2177.58 4723.49 9230.85

182 229 254 272 285

STL-10 𝜆 max

tr(b Σ)

𝑟 eff,90

3072 12288 27648

160.69 638.36 1437.94

783.73 3022.54 7061.82

145 130 154

Cat

3072 12288 27648

132.58 531.07 1195.06

700.49 2679.13 6308.94

150 134 163

Airplane

3072 12288 27648

193.60 778.14 1750.48

756.55 2920.58 6823.12

105 97 119

Class

𝐷

Dog

Table 3. Pairwise data geometry of the STL-10 clusters across image resolutions. For each pair and ambient dimension 𝐷, we report the squared mean separation, the maximum mean difference projection defined in (16), and the squared Frobenius covariance separation, computed from empirical data. ∥𝜇 b𝑖 − 𝜇 b𝑗 ∥ 2

𝑀𝜇

∥b Σ𝑖 − b Σ 𝑗 ∥ 2𝐹

Dog–Cat

3072 12288 27648

9.20 35.10 80.46

2.09 4.10 6.17

4.03×103 5.98×104 3.19×105

Cat–Airplane

3072 12288 27648

232.68 928.48 2091.26

12.81 25.59 38.39

1.35×104 2.12×105 1.09×106

Dog–Airplane

3072 12288 27648

189.10 753.39 1697.45

10.74 21.46 32.19

1.18×104 1.85×105 9.55×105

Pair

𝐷

Figure 1 demonstrates the dynamical concentration or decay of the posterior weights, evaluated for trajectories committing to a particular class, in the case of (a) PBMC data. Note that the reverse process proceeds from the left (low SNR) to the right (high SNR). The lines labelled with a target class 𝑐† ∈ C † (‘class name’ (target) in the legend) indicate the lower envelope of the posterior probability 𝑊𝐷[𝑐 ] achieved † by 99% of the paths in X𝑐 , while the lines with competing classes 𝑐 ∈ C \ {𝑐† } indicate the upper † envelope of 𝑊𝐷[𝑐] by 99% of the paths in X𝑐 . As expected, for trajectories labelled with the chosen target †

cluster 𝑐† (e.g., ‘CD4 T’), the posterior weight 𝑊𝐷[𝑐 ] eventually concentrates on 1 as the SNR increases (as the physical time 𝑡 approaches the terminal 𝑇) while the others 𝑊𝐷[𝑐] (e.g., ‘CD14+ Monocyte’ and ‘CD19+ B’) degenerate to 0. More importantly, these plots showcase the rapid posterior concentration, or the phase transition, for each ambient dimension 𝐷, and its timing and corresponding SNR shifts earlier as 𝐷 increases. In particular, the sudden decay of non-target posterior weights, followed by the sharp rise of the target posterior weight, is interpreted as a sign that the KL signal between the corresponding noisy clusters begins to dominate the fluctuation terms in the log-posterior ratio.

16

Y. Iguchi and P. Fearnhead

Figure 1. Posterior probability dynamics in the PBMC experiment. We construct increasing-dimensional representations by retaining the top-𝐷 highly variable genes, with 𝐷 ∈ {500, 1000, 2000, 4000, 8000}. The posterior concentration timing shifts to lower SNR as 𝐷 increases.

In Figure 2, we show the dynamical behaviour of the posterior probability for the image datasets with a fixed ambient dimension 𝐷 = 96 × 96 × 3, but now with plots of the KL divergence between the noised −𝑝 [𝑐† ] || → −𝑝 [𝑐]
, 𝑗 = 1, . . . , 𝑛. target cluster 𝑐† and the competing class 𝑐 ∈ C \ {𝑐† }, precisely, KL → 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 The KL divergence is not available in closed form and is therefore estimated from empirical data. The left panel of Figure 2 shows that, for trajectories eventually committing to the Dog class, the posterior weight of Airplane decays earlier than that of Cat. This behaviour is also reflected in the noised KL curves: the noised KL divergence between Dog and Airplane is already larger than that between Dog and Cat in the early-SNR phase. This also agrees with the pairwise class separation presented in Table 3. Thus, the reverse trajectories committing to Dog first rule out the Airplane class, while the finer ambiguity between Dog and Cat persists until a larger SNR. By contrast, in the right panel of Figure 2, for trajectories committing to Airplane, the KL curves against Dog and Cat are much closer to each other. Correspondingly, the Dog and Cat posterior weights decay in a more symmetric manner. These observations are consistent with the phase-transition analysis in Section 3, where the decay of non-target posterior weights is shown to be governed by the KL divergence between noised clusters. Finally, Figure 3 shows the mean and standard deviation of the first SNR level after which the lower † 1% quantile of the target posterior weight 𝑊𝐷[𝑐 ] , evaluated among trajectories eventually assigned to the target cluster 𝑐† , remains above a threshold 0.99 for all subsequent time steps. This indicates that for at

Diffusion Models for Clustered Data

17

Figure 2. Posterior commitment in the STL-10 experiment. Posterior probabilities are evaluated using the reverse trajectories, grouped by their final committed class. Solid curves show the posterior probabilities of the three classes, while dashed curves show the corresponding noised KL divergences from the target class to competing classes.

(a) PBMC single-cell data.

(b) STL-10 image data.

Figure 3. Average empirical SNR at the posterior probability concentration over five repetitions of experiments, with the error bar showing their standard deviations. The plots show the SNR at which the lower 1 % quantile of the posterior weight among trajectories assigned to a given cluster permanently exceeds 0.99. In both STL-10 and PBMC experiments, the transition SNR exhibits an approximately O (1/𝐷) decay.

least 99% of the trajectories, the posterior weights permanently exceed the threshold after that time. The mean and standard deviation are computed over five independent repetitions of the four-step procedure explained in Section 4.1. The metric provides an empirical proxy for the phase transition timing analysed in Section 3. In both the STL-10 and PBMC experiments, the SNR at the phase transition decreases as the ambient dimension 𝐷 increases, exhibiting an empirical O (1/𝐷) scaling, in agreement with the posterior-commitment scaling predicted by the theory for GMMs (Theorem 1) up to logarithmic factors.

5. Proof of Theorem 2 ← − ← − ← − ← − We write the measures of path 𝑋 and 𝑌 over time interval T = [𝑎, 𝑏], 0 < 𝑎 < 𝑏, as P T and Q T , respectively, and SNR 𝑗 ≡ 𝑚 𝑇2 −𝑡 𝑗 /𝜎𝑇2 −𝑡 𝑗 , 𝑗 = 0, 1, . . . , 𝑁, for ease of notation.

18

Y. Iguchi and P. Fearnhead

5.1. Step 1: Error Decomposition of KL Divergence We first decompose the KL divergence into three sources of error by following the arguments in convergence analysis for diffusion models [Benton et al., 2024, Conforti et al., 2025, Huang et al., 2026]. → − ← − First, the data processing inequality and the equivalence of measures P 𝜀 ≡ P 𝑇 − 𝜀 lead to → − ← − ← − ← − KL( P 𝜀 || Q 𝑡𝑁 ) ≤ KL( P [0,𝑇 − 𝜀 ] || Q [0,𝑇 − 𝜀 ] ). Further, employing the arguments in [Benton et al., 2024, Section 3.3], together with the use of Girsanov’s theorem (see e.g., [Benton et al., 2024, Appendix F] or [Chen et al., 2023b, Section 5.2]), we get
← − ← − KL P [0,𝑇 − 𝜀 ] || Q [0,𝑇 − 𝜀 ] = 𝑅1 + 𝑅2 with
→ − 𝑅1 ≡ KL P 𝑇 || N (0, 𝐼 𝐷 ) ,

𝑅2 ≡

𝑁 −1 ∫ 𝑡 𝑗+1 ∑︁ 𝑗=0

𝑡𝑗

2  ← − ← − 2 𝑚𝑇 −𝑡 × E 𝜇(𝑇 − 𝑡, 𝑋 𝑡 ) − 𝜇(𝑇 ˆ − 𝑡 𝑗 , 𝑋 𝑡 𝑗 ) 𝑑𝑡, 𝜎𝑇4 −𝑡

→ − → − where we recall 𝜇(𝑡, 𝑥) = E[ 𝑋 0 | 𝑋 𝑡 = 𝑥]. For the first term 𝑅1 , following the proof arguments of [Chen et al., 2023a, Lemma 9] or [Benton et al., 2024, Proposition 9], we have:   𝑅1 ≲ (𝐷 + E 𝑃data ∥ 𝑋 ∥ 2 )𝑒 −2𝑇 .

(26)

The second term 𝑅2 is further decomposed as 𝑅2 ≲ 𝑅2𝐼 + 𝑅2𝐼 𝐼 , with 𝑅2𝐼 =

𝑁 −1 ∫ 𝑡 𝑗+1 ∑︁ 𝑗=0

𝑅2𝐼 𝐼 =

𝑡𝑗

𝑁 −1 ∫ 𝑡 𝑗+1 ∑︁ 𝑗=0

𝑡𝑗

2  ← − ← − 2 𝑚𝑇 −𝑡 𝜇(𝑇 ˆ − 𝑡 , 𝑋 × E ) − 𝜇(𝑇 − 𝑡 , 𝑋 𝑑𝑡; 𝑗 𝑡 𝑗 𝑡𝑗 ) 4 𝑗 𝜎𝑇 −𝑡

2  ← − ← − 2 𝑚𝑇 −𝑡 × E 𝜇(𝑇 − 𝑡 𝑗 , 𝑋 𝑡 𝑗 ) − 𝜇(𝑇 − 𝑡, 𝑋 𝑡 ) 𝑑𝑡. 𝜎𝑇4 −𝑡

(27)

For the term 𝑅2𝐼 , we have from (22), Tweedie’s formula and Assumption 5 that 𝑅2𝐼 ≲

𝑁 −1 ∑︁

 ← − ← − 2 −𝑝 2 ≤ 𝑇 𝜀score . ℎ 𝑗+1 × E 𝑠ˆ (𝑇 − 𝑡 𝑗 , 𝑋 𝑡 𝑗 ) − ∇ log→ 𝑇 −𝑡 𝑗 ( 𝑋 𝑡 𝑗 )

(28)

𝑗=0 2 𝜎𝑇4 −𝑡 𝑚𝑇 −𝑡 where we also used 𝜎 4 · 𝑚2 𝑗 𝑇 −𝑡 𝑇 −𝑡

𝑗

2 𝑚𝑇 −𝑡

≤ 𝜎4

𝑗+1

𝑇 −𝑡 𝑗+1

𝜎𝑇4 −𝑡

· 𝑚2

𝑗

≤ 4 exp(2𝜏) from Lemma A.1.

𝑇 −𝑡 𝑗

5.2. Step 2. Decomposition of the Time-Discretisation Error Our main focus is thus on bounding the term 𝑅2𝐼 𝐼 in (27). We exploit the following result established in the literature. The proof can be found, e.g., in [Huang et al., 2026, Lemma 9]. Lemma 5.1. Let 𝑗 ∈ [𝑁]. It then holds that for any 𝑡 ∈ [𝑡 𝑗 −1 , 𝑡 𝑗 ], h h h
i ← − ← − 2i ← − ← −
i E 𝜇(𝑇 − 𝑡 𝑗 −1 , 𝑋 𝑡 𝑗 −1 ) − 𝜇(𝑇 − 𝑡, 𝑋 𝑡 ) = E tr Σ𝑇 −𝑡 𝑗 −1 ( 𝑋 𝑡 𝑗 −1 ) − E tr Σ𝑇 −𝑡 ( 𝑋 𝑡 ) . h ← −
i Lemma 5.1 implies that 𝑡 ↦→ E tr Σ𝑇 −𝑡 ( 𝑋 𝑡 ) is a non-increasing function, thus h h h
i ← − ← − 2i ← − ← −
i E 𝜇(𝑇 − 𝑡 𝑗 −1 , 𝑋 𝑡 𝑗 −1 ) − 𝜇(𝑇 − 𝑡, 𝑋 𝑡 ) ≤ E tr Σ𝑇 −𝑡 𝑗 −1 ( 𝑋 𝑡 𝑗 −1 ) − E tr Σ𝑇 −𝑡 𝑗 ( 𝑋 𝑡 𝑗 ) , for any 𝑡 ∈ [𝑡 𝑗 −1 , 𝑡 𝑗 ]. Recall ℎ 𝑗+1 ≤ 𝜏𝜎𝑇2 −𝑡 𝑗 where the equality holds for 𝑗 = 0, 1, . . . , 𝑁 − 2. Then, we

Diffusion Models for Clustered Data

19

apply Lemma 5.1 to the term 𝑅2𝐼 𝐼 and obtain 𝑅2𝐼 𝐼 ≤

𝑁 −1 ∑︁

2 𝑚𝑇 −𝑡

ℎ 𝑗+1 𝜎 4

𝑇 −𝑡 𝑗+1

𝑗=0

≤ 2𝜏

𝑗+1

𝑁 −1 ∑︁

  
 ← − ← −
 · E tr Σ𝑇 −𝑡 𝑗 ( 𝑋 𝑡 𝑗 ) − E tr Σ𝑇 −𝑡 𝑗+1 ( 𝑋 𝑡 𝑗+1 )

  
 ← − ← −
 SNR 𝑗+1 E tr Σ𝑇 −𝑡 𝑗 ( 𝑋 𝑡 𝑗 ) − E tr Σ𝑇 −𝑡 𝑗+1 ( 𝑋 𝑡 𝑗+1 )

𝑗=0 𝑁 −1 ∑︁ 

 → − → −
 SNR 𝑗+1 − SNR 𝑗 · E tr Σ𝑇 −𝑡 𝑗 ( 𝑋 𝑇 −𝑡 𝑗 ) ≡ T , ≲ 𝜏 SNR1 · E tr Σ𝑇 ( 𝑋 𝑇 ) + 𝜏

(29)

𝑗=1

where in the second line we used the bound 𝜎𝑇2 −𝑡 𝑗 /𝜎𝑇2 −𝑡 𝑗+1 ≤ 2, which is presented in Lemma A.1. For the ease of notation, we write: Δ 𝑗 = (SNR 𝑗+1 − SNR 𝑗 ),

→ − 𝑍 𝑗 = (𝑡 𝑗 , 𝑋 𝑇 −𝑡 𝑗 ),

𝑗 = 0, . . . , 𝑁 − 1.

(30)

Using the closed-form expression of posterior covariance (24) in Lemma 3.4, we decompose the term T as T = T1 + T2 , where T1 = 𝜏

T2 = 𝜏2

𝑁 −1  ∑︁

𝐾  ∑︁   
−1  SNR1 1 𝑗=0 + Δ 𝑗 1 𝑗 ≥1 E 𝑊𝐷[ 𝑘 ] (𝑍 𝑗 ) · 𝜎𝑇2 −𝑡 𝑗 tr Σ 𝑘 𝑚 𝑇2 −𝑡 𝑗 Σ 𝑘 + 𝜎𝑇2 −𝑡 𝑗 𝐼 𝐷 ;

𝑗=0

𝑘=1

𝑁 −1  ∑︁

𝐾  ∑︁

SNR1 1 𝑗=0 + Δ 𝑗 1 𝑗 ≥1

h i 2 E 𝑊𝐷[ 𝑘1 ] (𝑍 𝑗 )𝑊𝐷[ 𝑘2 ] (𝑍 𝑗 ) 𝜇 [ 𝑘1 ] (𝑍 𝑗 ) − 𝜇 [ 𝑘2 ] (𝑍 𝑗 )

(31)

(32)

𝑘1 ,𝑘2 =1 𝑘1 ≠𝑘2

𝑗=0

→ − → − → − [𝑖 ]  where we recall that for 𝑧 ≡ (𝑡, 𝑥) ∈ [0, 𝑇) × R𝐷 , 𝜇 [𝑖 ] (𝑧) = E 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑥, 𝑋 0 ∼ 𝑃data , 𝑖 ∈ [𝐾].

5.3. Step 3. Bounding the Term T1 The eigen-decomposition Σ 𝑘 ≡ 𝑈Λ 𝑘 𝑈 ⊤ together with Assumption 2 gives that: for 𝑗 = 0, 1, . . . , 𝑁 − 1,      −1   −1  2 2 2 2 2 2 𝜎𝑇 −𝑡 𝑗 tr Σ 𝑘 𝑚 𝑇 −𝑡 𝑗 Σ 𝑘 + 𝜎𝑇 −𝑡 𝑗 𝐼 𝐷 = 𝜎𝑇 −𝑡 𝑗 tr Λ 𝑘 𝑚 𝑇 −𝑡 𝑗 Λ 𝑘 + 𝜎𝑇 −𝑡 𝑗 𝐼 𝐷 = 𝜎𝑇2 −𝑡 𝑗

𝐷 ∑︁ ℓ=1

𝜆ℓ (Σ 𝑘 ) 2 2 𝑚𝑇 𝜆 −𝑡 ℓ (Σ 𝑘 )+𝜎𝑇 −𝑡 𝑗

𝜎𝑇2 −𝑡

≤ 𝑚2 𝑗

𝑗

𝑀𝑘 + 𝜈 𝑘 .

(33)

𝑇 −𝑡 𝑗

Thus, we get 𝐾 𝑁 −1 𝐾   ∑︁ ∑︁ ∑︁    𝜎𝑇2 −𝑡  [𝑘]   𝜎𝑇2 Δ𝑗 E 𝑊𝐷[ 𝑘 ] (𝑍 𝑗 ) · 𝑚2 𝑗 𝑀𝑘 + 𝜈 𝑘 T1 ≤ 𝜏SNR1 E 𝑊𝐷 (𝑍0 ) · 𝑚2 𝑀𝑘 + 𝜈 𝑘 + 𝜏 𝑇

𝑘=1

𝑗=1

𝑘=1

𝑁 −1   ∑︁
† Δ𝑗 SNR1 † ≤ 𝜏 SNR + SNR 𝑅 𝑀 + 𝜏 1 + SNR 𝑅 𝑀 1 𝑗 SNR 𝑗 0 𝑗=1



1 = 𝜏 SNR SNR0 +

𝑁 −1 ∑︁ 𝑗=1

Δ𝑗 SNR 𝑗



𝑀 † + 𝜏SNR 𝑁 𝑅,

𝑇 −𝑡 𝑗

20

Y. Iguchi and P. Fearnhead

where 𝑀 † = max 𝑘 ∈ [𝐾 ] 𝑀𝑘 , 𝑅 ∈ [0, 1) is a universal constant defined in Assumption 2 and we used Í𝐾 [𝑘] 𝐷 𝑘=1 𝑊𝐷 (𝑡, 𝑥) = 1 for all (𝑡, 𝑥) ∈ [0, 𝑇) × R . We then get under 𝜀 ≳ 𝑅/𝑇 that     † 𝑅 1 𝑁 𝑀 + 𝜏 ≲ 𝜏 log + 𝑇 𝑀 †, T1 ≲ 𝜏 1 + log SNR SNR1 𝜀 𝜀

(34)

since it follows that for 𝑗 ∈ [𝑁 − 1], SNR 𝑗+1 Δ𝑗 SNR 𝑗 = SNR 𝑗

∫ SNR 𝑗+1 SNR 𝑗

SNR 𝑗+1 1 SNR 𝑗+1 𝑑𝑦 ≤ SNR 𝑗

∫ SNR 𝑗+1 SNR 𝑗

SNR 𝑗+1 1 𝑦 𝑑𝑦 ≲ log SNR 𝑗

(35)

with SNR 𝑘+1 /SNR 𝑘 ≤ 4, 𝑘 = 0, 1, . . . , 𝑁 − 1 from Lemma A.1 and 1 ≤ − log 𝑚 𝑇2 −𝑡1 = 2(𝑇 − 𝑡 1 ) ≲ 𝑇, log SNR 1  
1 1 log SNR 𝑁 = log exp(2𝜀) ∵ exp(𝑥) − 1 ≥ 𝑥, ∀𝑥 > 0 . −1 ≤ log 𝜀 ,

5.4. Step 4. Bounding the Term T2 We hereby focus on bounding the term involving the inter-class difference between clusters, i.e., 2 𝜇 [ 𝑘1 ] (𝑍 𝑗 ) − 𝜇 [ 𝑘2 ] (𝑍 𝑗 ) that can depend on ∥𝜇𝑖 − 𝜇 𝑗 ∥ 2 = O (𝐷) under Assumption 3. To avoid the polynomial dependence on 𝐷 in the bound, we make use of the phase transition from mixing to single-cluster commitment, as argued in Section 3. We thus separate the time interval [0, 𝑇 − 𝜀] into two parts and get T2 = T2𝐼 + T2𝐼 𝐼 with 𝐼

T2 = 𝜏2

𝑁∑︁ 1 −1

SNR1 1 𝑗=0 + Δ 𝑗 1 𝑗 ≥1

h i 2 E 𝑊𝐷[ 𝑘1 ] (𝑍 𝑗 )𝑊𝐷[ 𝑘2 ] (𝑍 𝑗 ) 𝜇 [ 𝑘1 ] (𝑍 𝑗 ) − 𝜇 [ 𝑘2 ] (𝑍 𝑗 ) ;

(36)

𝑘1 ,𝑘2 =1

𝑗=0

T2𝐼 𝐼 = 𝜏2

𝐾
∑︁

𝑁 −1 ∑︁

Δ𝑗

𝑗=𝑁1

𝐾 ∑︁

i h 2 E 𝑊𝐷[ 𝑘1 ] (𝑍 𝑗 )𝑊𝐷[ 𝑘2 ] (𝑍 𝑗 ) 𝜇 [ 𝑘1 ] (𝑍 𝑗 ) − 𝜇 [ 𝑘2 ] (𝑍 𝑗 ) ,

(37)

𝑘1 ,𝑘2 =1

where in the above we have set: o n log(𝐾 𝐷) 𝑁1 ≡ inf 𝑗 ∈ [𝑁] | SNR 𝑗 > 𝐿 𝐷

(38)

for an arbitrary fixed constant 𝐿 > 0 satisfying 𝐿 log(𝐾 𝐷)/𝐷 ≥ 𝑆, with 𝑆 specified in the statement log(𝐾 𝐷) of Theorem 1. We note that 𝑁1 ≥ 1 under 𝑇 ≥ 12 log 𝐷 due to Corollary 1 and SNR 𝑗 ≲ for 𝐷 → − [𝑖 ] [𝑖 ] 𝑗 ∈ {1, . . . , 𝑁1 − 1}. In what follows, we write 𝑍 𝑗 = (𝑡 𝑗 , 𝑋 𝑇 −𝑡 𝑗 ), 𝑗 ∈ {0, 1, . . . , 𝑁 − 1}, 𝑖 ∈ [𝐾] where → − [𝑖 ] −𝑝 [𝑖 ] . 𝑋 follows the distribution → 𝑇 −𝑡 𝑗

𝑇 −𝑡 𝑗

5.4.1. Bounding the term T2𝐼 We begin with the term T2𝐼 , that is, the error during the mixing phase. We then get that T2𝐼 = 𝜏2

𝑁∑︁ 1 −1 𝑗=0

SNR1 1 𝑗=0 + Δ 𝑗 1 𝑗 ≥0

∑︁

∑︁

𝜋𝑖 · 𝑄 𝑘[𝑖1],𝑘2 ( 𝑗)

𝑖 ∈ [𝐾 ] 𝑘1 ,𝑘2 =1

with i h [ 𝑘2 ] [𝑖 ] [𝑖 ] [𝑖 ] 2 ] [ 𝑘1 ] [ 𝑘2 ] ) − 𝜇 (𝑍 ) , 𝑄 𝑘[𝑖1],𝑘2 ( 𝑗) = E 𝑊𝐷[ 𝑘1 ] (𝑍 [𝑖 )𝑊 (𝑍 ) 𝜇 (𝑍 𝑗 𝑗 𝑗 𝑗 𝐷

Diffusion Models for Clustered Data

21

→ − [𝑖 ] ] where the expectation operator acts on the second component of the variable 𝑍 [𝑖 𝑗 under the law of 𝑋 𝑇 −𝑡 𝑗 . We have that  ∑︁  ∑︁  ] [𝑖 ] 2  [𝑘] [𝑖 ] [ 𝑘2 ] 𝑄 𝑘[𝑖1],𝑘2 ( 𝑗) ≤ E max 𝜇 [ 𝑘1 ] (𝑍 [𝑖 ) − 𝜇 (𝑍 ) ∵ 𝑊 (𝑍 ) = 1 𝑗 𝑗 𝑗 𝐷 𝑘1 ,𝑘2 ∈ [𝐾 ]

𝑘1 ,𝑘2 ∈ [𝐾 ]

  = 𝜉1𝑗 log exp E 𝜉 𝑗 ·  h  ≤ 𝜉1𝑗 log E exp 𝜉 𝑗  h = 𝜉1𝑗 log E max

max

] ] [ 𝑘2 ] 𝜇 [ 𝑘1 ] (𝑍 [𝑖 (𝑍 [𝑖 𝑗 )−𝜇 𝑗 )

2

max

] ] [ 𝑘2 ] 𝜇 [ 𝑘1 ] (𝑍 [𝑖 (𝑍 [𝑖 𝑗 )−𝜇 𝑗 )

2

𝑘1 ,𝑘2 ∈ [𝐾 ]

𝑘1 ,𝑘2 ∈ [𝐾 ]

𝑘1 ,𝑘2 ∈ [𝐾 ]

𝑘 ∈ [𝐾 ]



i

∵ Jensen′ s inequality

i  ] ] 2 [ 𝑘2 ] (𝑍 [𝑖 exp 𝜉 𝑗 𝜇 [ 𝑘1 ] (𝑍 [𝑖 𝑗 )−𝜇 𝑗 )

(39)

for any 𝜉 𝑗 > 0 and for each fixed 𝑖 ∈ [𝐾] and 𝑗 ∈ {0, 1, . . . , 𝑁1 − 1}. To further bound the above term, we exploit the following result whose proof is postponed to Appendix E.2:

Lemma 5.2. Under Assumptions 2, 3, we have the following: for any 𝑖 ∈ [𝐾], 𝑗 ∈ {0, 1, . . . , 𝑁 − 1} and  any 𝜌 ∈ 0, 𝑚 𝑇2 −𝑡 𝑗 /(16𝜎𝑇4 −𝑡 𝑗 max𝑖,𝑘 ∈ [𝐾 ] ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ) , it holds  h log E max

𝑘1 ,𝑘2 ∈ [𝐾 ]

 i
] [𝑖 ] 2 [ 𝑘2 ] exp 𝜌 𝜇 [ 𝑘1 ] (𝑍 [𝑖 ) − 𝜇 (𝑍 ) ≲ log 𝐾 + 𝜌 𝐷 + SNR 𝑗 𝑀 † 𝐷 2𝛼 𝑗 𝑗

where we have set  𝐻 𝑘,𝑖 ( 𝑗) = 𝑚 𝑇4 −𝑡 𝑗 (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 (Σ 𝑘 − Σ𝑖 ) (𝑆𝑇𝑖 −𝑡 𝑗 ) −1 (Σ 𝑘 − Σ𝑖 ) (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 .

(40)

Setting 𝜉 𝑗 = 𝑚 𝑇2 −𝑡 𝑗 /(16𝜎𝑇4 −𝑡 𝑗 max𝑖,𝑘 ∈ [𝐾 ] ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ) in (39), we apply Lemma 5.2 to get 𝑁∑︁ 1 −1 Δ𝑗 † 2𝛼
† 2𝛼
1 𝐷 ) + 𝜏 log 𝐾 + 𝜉 (𝐷 + SNR 𝑀 T2𝐼 ≲ 𝜏 SNR 0 𝑗 𝜉0 𝜉 𝑗 log 𝐾 + 𝜉 𝑗 (𝐷 + SNR 𝑗 𝑀 𝐷 ) 𝑗=1 𝑁∑︁ 1 −1    ≲ 𝜏 SNR1 + Δ 𝑗 (log 𝐾)SNR 𝑁1 −1 𝐷 2𝛼 + 𝐷 + SNR 𝑁1 −1 𝑀 † 𝐷 2𝛼 𝑗=1




= 𝜏SNR 𝑁1 SNR 𝑁1 −1 log 𝐾 + 𝑀 † 𝐷 2𝛼 + 𝐷  

≲ 𝜏SNR 𝑁1 −1 SNR 𝑁1 −1 log 𝐾 + 𝑀 † 𝐷 2𝛼 + 𝐷 ∵ Lemma A.1 2 i h log(𝐾 𝐷)
log(𝐾 𝐷) †
log 𝐾 + 𝑀 + log(𝐾 𝐷) ∵ SNR 𝑁1 −1 ≲ ≲𝜏 1− 𝛼 𝐷 𝐷

(41)

where in the second line we used 𝜎𝑇4 −𝑡 𝑗 1 max ∥𝐻 𝑘,𝑖 ∥ 2 ≲ SNR 𝑗 𝐷 2𝛼 ≲ SNR 𝑁1 −1 𝐷 2𝛼 . 𝜉 𝑗 ≲ 𝑚2 𝑖,𝑘 ∈ [𝐾 ]

(42)

𝑇 −𝑡 𝑗

5.4.2. Bounding the term T2𝐼 𝐼 We then consider the error during the cluster commitment phase, that is, T2𝐼 𝐼 defined in (37). In this phase, we can exploit the high-probability concentration of the posterior weight stated in Theorem 1.

22

Y. Iguchi and P. Fearnhead

] Introducing an event set 𝐸 [𝑖 𝑗 (𝜂) ⊆ Ω, 𝜂 ∈ (0, 1) as

n ] 𝐸 [𝑖 (𝜂) = 𝜔∈Ω| 𝑗

o  [ 𝑘 ] [𝑖 ] −𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]

, 𝑊𝐷 (𝑍 𝑗 ) ≤ exp −𝜂KL → 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗

Ù 𝑘 ∈ [𝐾 ]\{𝑖 }

𝑐

we express the term T2𝐼 𝐼 as T2𝐼 𝐼 = T2𝐼 𝐼,𝐸 + T2𝐼 𝐼,𝐸 with T2𝐼 𝐼,𝐸 = 𝜏2

𝑁 −1 ∑︁

Δ𝑗

T2

= 𝜏2

𝑁 −1 ∑︁ 𝑗=𝑁1

∑︁

] 𝜋𝑖 · 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 𝐸 [𝑖 𝑗 (𝜂) ;

(43)

] 𝑐
𝜋𝑖 · 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 𝐸 [𝑖 𝑗 (𝜂) ,

(44)

𝑖 ∈ [𝐾 ] 𝑘1 ,𝑘2 ∈ [𝐾 ]

𝑗=𝑁1 𝐼 𝐼,𝐸 𝑐

∑︁

Δ𝑗

∑︁

∑︁

𝑖 ∈ [𝐾 ] 𝑘1 ,𝑘2 ∈ [𝐾 ]

where we have set: for an event 𝐴 ⊆ Ω, i h [𝑖 ] [𝑖 ] 2 ] [ 𝑘2 ] [𝑖 ] [ 𝑘1 ] [ 𝑘2 ] 1 𝜇 (𝑍 ) − 𝜇 (𝑍 ) 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 ( 𝐴) = E 𝑊𝐷[ 𝑘1 ] (𝑍 [𝑖 )𝑊 (𝑍 ) 𝐴 . 𝑗 𝑗 𝑗 𝑗 𝐷 For notational simplicity, we write: h i [𝑖 ] 2 ] [𝑖 ] 𝐷 𝑘[𝑖 ] ( 𝑗) = E 𝜇 [ 𝑘 ] (𝑍 [𝑖 ) − 𝜇 (𝑍 ) , 𝑗 𝑗

𝑘 ∈ {𝑘 1 , 𝑘 2 }.

] Then, for the term T2𝐼 𝐼,𝐸 , we have the following bound on the event set 𝐸 [𝑖 𝑗 (𝜂):

   ] → −𝑝 [𝑖 ] || → −𝑝 [ 𝑘1 ]
+ KL → −𝑝 [𝑖 ] || → −𝑝 [ 𝑘2 ]
·  𝐷 [𝑖 ] ( 𝑗) + 𝐷 [𝑖 ] ( 𝑗) 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 (𝐸 [𝑖 (𝜂)) ≲ exp −𝜂 KL 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑘1 𝑘2   
−𝑝 [𝑖 ] || → −𝑝 [ 𝑘1 ]
+ KL → −𝑝 [𝑖 ] || → −𝑝 [ 𝑘2 ]
≲ ∥Σ𝑖 ∥ 2 + SNR−1 exp −𝜂 KL → 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 o n [𝑖 ] [ 𝑘1 ]
[𝑖 ] → − → − → − −𝑝 [ 𝑘2 ]
× KL 𝑝 𝑇 −𝑡 𝑗 || 𝑝 𝑇 −𝑡 𝑗 + KL 𝑝 𝑇 −𝑡 𝑗 || → 𝑇 −𝑡 𝑗   


→ −𝑝 [𝑖 ] || → −𝑝 [ 𝑘1 ] + KL → −𝑝 [𝑖 ] || → −𝑝 [ 𝑘2 ] ≲ ∥Σ𝑖 ∥ 2 + SNR−1 exp −e 𝜂 · KL 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗    2 2

e 𝐶 SNR 𝐷 ≤ ∥Σ𝑖 ∥ 2 + SNR−1 exp −e 𝜂 · SNR 𝑗 max𝑘 ∈ [𝐾 𝑗] tr(Σ𝑘 )+𝐷 ∵ Lemma 𝐶.2 𝑗 e = 𝐶3 /𝐶5 , where in the second line above we used the following result whose for some 𝜂e ∈ (0, 𝜂) and 𝐶 proof is postponed to Appendix E.3:

Lemma 5.3. It holds that:
−𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
. 𝐷 𝑘[𝑖 ] ( 𝑗) ≤ 4 ∥Σ𝑖 ∥ 2 + SNR−1 · KL → 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗

(45)

] To further bound the term 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 (𝐸 [𝑖 𝑗 (𝜂)) under the case 𝑗 ∈ {𝑁1 , . . . , 𝑁 − 1}, we define an integer 𝑁2 as

    inf 𝑗 = 𝑁1 , 𝑁1 + 1, . . . , 𝑁 | SNR 𝑗 > 𝐷/max 𝑘 ∈ [𝐾 ] tr(Σ 𝑘 ) ,  𝑁2 =   𝑁, 

if the set is not empty;

(46)

otherwise.

We then consider the following two cases for 𝑗: (a). 𝑗 = 𝑁1 , . . . , 𝑁2 − 1 and (b). 𝑗 = 𝑁2 , . . . , 𝑁 − 1, where case (b) is understood to be empty when 𝑁2 = 𝑁. For case (a), we have SNR 𝑗 ∥Σ𝑖 ∥ 2 ≤

Diffusion Models for Clustered Data

23

SNR 𝑗 max 𝑘 ∈ [𝐾 ] tr(Σ 𝑘 ) ≤ 𝐷 and thus   e2 ] 𝜂 e𝐶 −1 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 (𝐸 [𝑖 (𝜂)) ≲ SNR SNR 𝐷 (𝐷 + 1) exp − 𝑗 𝑗 𝑗 2

−𝑐𝐿 ≤ SNR−1 , ∵ SNR 𝑗 ≥ 𝐿log(𝐾 𝐷)/𝐷 , 𝑗 𝐷 + 1 · (𝐾 𝐷)

(47)

e2 /2. On the other hand, under case (b), we have where 𝑐 = 𝜂e𝐶 ] 𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 (𝐸 [𝑖 𝑗 (𝜂))

  e2

𝜂 e𝐶 𝛼 𝐷2 ∵ Assumption 2 ≲ SNR−1 𝑗 SNR 𝑗 𝐷 + 1 exp − 2 · max 𝑘 ∈ [𝐾 ] tr(Σ 𝑘 )  e2 
𝜂 e𝐶 𝐷2 𝐷𝛼 ≲ SNR−1 ∵ Assumption 2 & SNR 𝑗 ≤ SNR 𝜀 ≲ 1𝜀 𝑗 · 𝜀 · exp − 2 · 𝑀 † (𝐶2 𝐷 𝛼 +𝑅)     e2 𝜂 e𝐶 2− 𝛼 𝐷𝛼 𝐷 2𝛼−2 ≤ SNR−1 ≲ SNR−1 · exp −e 𝑐 · 𝐷 2− 𝛼 𝑗 · 𝜀 · exp − 4𝑀 † 𝐶 · 𝐷 𝑗 · 𝜀 2

(48)

e2 /4𝑀 †𝐶2 ). Substituting the bounds (47) and (48) into (43), we obtain for some e 𝑐 ∈ (0, 𝜂e𝐶 " # 𝑁∑︁ 𝑁 −1 2 −1 ∑︁ 2 Δ Δ 𝐼 𝐼,𝐸 𝑗 𝑗 𝐾 T2 ≲ 𝜏 𝐾 2 𝐷 (𝐾 𝐷) −𝑐𝐿 · 𝑐 𝐷 2− 𝛼 ) · SNR 𝑗 + 𝜀 exp(−e SNR 𝑗 𝑗=𝑁1

𝑗=𝑁2

oi
n 2 𝑐 𝐷 2− 𝛼 ) , ≲ 𝜏 log 1𝜀 + log 𝐷 · 𝐾 2 𝐷 (𝐾 𝐷) −𝑐𝐿 + 𝐾𝜀 exp(−e h

(49)

where we used max

2 −1 n 𝑁∑︁

Δ𝑗 SNR 𝑗 ,

𝑗=𝑁1

𝑁 −1 ∑︁

Δ𝑗 SNR 𝑗

o

≤

𝑗=𝑁2

𝑁 −1 ∑︁

Δ𝑗 1 SNR 𝑗 ≤ log 𝜀 + log 𝐷.

(50)

𝑗=𝑁1

] 𝑐 We now consider the term T2𝐼 𝐼,𝐸 defined in (44). For ease of notation, we write 𝐵 = (𝐸 [𝑖 𝑗 (𝜂)) and −𝑝 [𝑖 ] , and introduce a probability measure P defined as dP = 1𝐵 dP . We then get P =→ 𝑐

1

2

𝑇 −𝑡 𝑗

𝐾 ∑︁

𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 (𝐵) ≤ EP1

h

𝑘1 ,𝑘2 =1 P1 𝐵 = 𝜉𝑗

h · E P2 𝜉 𝑗

max

𝑘1 ,𝑘2 ∈ [𝐾 ]

max

𝑘1 ,𝑘2 ∈ [𝐾 ]

( h  P1 𝐵 ≤ 𝜉 𝑗 · log EP1 exp 𝜉 𝑗

𝑘1 ,𝑘2 ∈ [𝐾 ]

2

1

P1 ( 𝐵)

] ] [ 𝑘2 ] 𝜇 [ 𝑘1 ] (𝑍 [𝑖 (𝑍 [𝑖 𝑗 )−𝜇 𝑗 )

] ] [ 𝑘2 ] 𝜇 [ 𝑘1 ] (𝑍 [𝑖 (𝑍 [𝑖 𝑗 )−𝜇 𝑗 )

max

2

2

· 1𝐵

i

i

] ] [ 𝑘2 ] 𝜇 [ 𝑘1 ] (𝑍 [𝑖 (𝑍 [𝑖 𝑗 )−𝜇 𝑗 )

2

i

) + KL P2 || P1

,

h  i ] [ 𝑘2 ] (𝑍 [𝑖 ] ) 2 for 𝜉 𝑗 > 0 such that EP1 exp 𝜉 𝑗 max 𝑘1 ,𝑘2 ∈ [𝐾 ] 𝜇 [ 𝑘1 ] (𝑍 [𝑖 ) − 𝜇 < ∞, where the second 𝑗 𝑗 bound is obtained via Donsker-Varadhan’s variational formula [Dupuis and Ellis, 1997]. We im
mediately have KL P2 || P1 = − log P1 (𝐵). Also, Theorem 1 gives that for any 𝑞 ≥ 1, 𝜂 ∈ (0, 1), there exists 𝐷 † ∈ N such that for any 𝐷 ≥ 𝐷 † , we have P1 (𝐵) ≤ 𝐷 −𝑞 . Thus, by setting 𝜉 𝑗 =

24

Y. Iguchi and P. Fearnhead

𝑚 𝑇2 −𝑡 𝑗 /(16𝜎𝑇4 −𝑡 𝑗 max𝑖,𝑘 ∈ [𝐾 ] ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ), we obtain: 𝐾 ∑︁

𝑄 [𝑖𝑗 ],𝑘1 ,𝑘2 (𝐵) ≲ P1 (𝐵) ·

n

† 2𝛼 1 + 𝜉1𝑗 log P1 (1𝐵) 𝜉 𝑗 log 𝐾 + 𝐷 + SNR 𝑗 𝑀 𝐷

o

∵ Lemma 5.2

𝑘1 ,𝑘2 =1

≲ 𝐷 2𝛼−𝑞 · SNR 𝑗 · log 𝐾 + 𝑀 † + 𝐷 1−𝑞 + SNR 𝑗 𝐷 2𝛼 P1 (𝐵) log P1 (1𝐵)
≲ 𝐷 2𝛼−𝑞 · SNR 𝑗 · log 𝐾 + 𝑀 † + 𝐷 1−𝑞 + 𝐷 2𝛼−𝑞 · SNR 𝑗 · log 𝐷
2𝛼−𝑞 ≲ SNR−1 · 𝜀 −2 · (log 𝐾 + 𝑀 † + log 𝐷) + 𝐷 1−𝑞 · 𝜀 −1 , 𝑗 𝐷

∵ (42)

(51)

where we used for 𝑥 log(1/𝑥) ≤ 𝑟 log(1/𝑟) 𝑥 ∈ [0, 𝑟], 𝑟 ≤ 𝑒 −1 , and SNR 𝑗 ≲ 𝜀 −1 in the third and fourth inequality, respectively. Thus, the bound (51) together with (50) leads to T2

𝐼 𝐼,𝐸 𝑐

h

·𝜀

−1

≲ 𝜏 𝐷 2𝛼−𝑞 · 𝜀 −2 · (log 𝐾 + 𝑀 † + log 𝐷) + 𝐷 1−𝑞 · 𝜀

−1

≲𝜏 𝐷

2𝛼−𝑞

·𝜀

−2

†

· (log 𝐾 + 𝑀 + log 𝐷) + 𝐷

1−𝑞

𝑁 ∑︁

Δ𝑗 SNR 𝑗

i

𝑗=𝑁1

h

i log 1𝜀 + log 𝐷 .

(52)

5.5. Step 5: Final Bound Combining the upper bounds (34), (41), (49) and (52) into the term 𝑅2𝐼 𝐼 defined in (29), the timediscretisation error is bounded as follows: 𝑐

𝑅2𝐼 𝐼 ≲ T1 + T2𝐼 + T2𝐼 𝐼,𝐸 + T2𝐼 𝐼,𝐸 h    2
log(𝐾 𝐷) ≲ 𝜏 𝑀 † log 1𝜀 + 𝑇 + 𝐷 1− 𝛼 log 𝐾 + 𝑀 † + log(𝐾 𝐷) o
n 2 𝑐 𝐷 2− 𝛼 ) + log 1𝜀 + log 𝐷 · 𝐾 2 𝐷 (𝐾 𝐷) −𝑐𝐿 + 𝐾𝜀 exp(−e

i + 𝐷 2𝛼−𝑞 · 𝜀 −2 · (log 𝐾 + 𝑀 † + log 𝐷) + 𝐷 1−𝑞 · 𝜀 −1 log 1𝜀 + log 𝐷 .

(53)

For the number mixture of 𝐾 that grows polynomially in 𝐷, by choosing sufficiently large 𝑞 ≥ 2 and 𝐿 > 0, we have under 𝑇 ≥ 12 log 𝐷 that for any sufficiently large 𝐷,   𝑅2𝐼 𝐼 ≲ 𝜏𝑀 † log 1𝜀 + 𝑇 .

(54)

We thus conclude Theorem 2 from (26), (28) and (54).

6. Conclusion and Discussion This work has analytically studied the geometric adaptivity of denoising diffusion models when the target distributions to sample from are characterised as a union of well-separated classes with their own low-dimensional structure, motivated by the union of manifolds hypothesis [Brown et al., 2022]. By focusing on 𝐾-Gaussian mixture models, our analysis has focused on the dynamical behaviour of the score function and, in particular, the behaviour of the posterior weights for sampling from each cluster component. Under the conditions, particularly the Θ(𝐷) inter-class separation in cluster means, we have shown that the concentration of the posterior weights occurs with high probability after the SNR
reaches a threshold Θ log(𝐾 𝐷)/𝐷 . We then use this result to obtain a KL divergence error bound, where the time-discretisation error is shown to scale linearly with the intrinsic dimensionality rather than the ambient dimensionality 𝐷. The main difference between our analysis and the literature (e.g., Benton et al. [2024], Conforti et al. [2025]) is that our analysis accounts for the dynamical regimes of the denoising process. Specifically, during the mixing phase (before the posterior commitment), the small

Diffusion Models for Clustered Data

25

SNR suppresses the effect of the inter-class separation of means and covariances in terms of 𝐷; whereas after the cluster commitment, when the SNR is no longer sufficiently small, the posterior weight decay neutralises the large cluster separations. We also note that the literature on diffusion models targeting GMMs [Li et al., 2026] establishes a dimension-free bound under shared isotropic covariances, while our work treats GMMs with heterogeneous anisotropic covariances and derives a desirable error bound by employing the phase transition. Our analysis exploits Θ(𝐷) separation in the mean vectors to induce the phase transition, with the key SNR scaling with 𝐷, but it could be extended to cases where covariance separation plays a central role in the transition. The numerical experiments in Section 4 demonstrate the SNR scaling Θ(log 𝐷/𝐷) in the phase transition for real high-dimensional datasets beyond our geometric framework (GMMs). Given this, our phase transition-driven error analysis can be applied to a wider class of mixture models where we believe similar intrinsic-dimension adaptivity would hold, provided we can still prove concentration results on the posterior weights and the score function for each mixture component is sufficiently well-behaved. The former is likely to hold if the mixture components are sub-Gaussian, while the latter would apply under assumptions such as strong log-concavity of the cluster densities. Another interesting extension would be to hierarchical clusters. For a mixture model with ‘supercomponents’, each of which is a Gaussian mixture model, with O (𝐷) separation between the supercomponent means, and smaller separation, say O (1), between the means of cluster components within each super-component. For this setting, we would still get concentration of the posterior weights to a super-component with high probability. Then, the KL divergence error is controlled by the intrinsic dimensionality in a similar manner, so that the relatively small separation within each super-component does not degrade the dimensional scaling in the error bound. Acknowledgement We thank Prof Kenji Fukumizu and Prof Paul Jenkins for detailed comments and feedback on this work. The authors acknowledge support from EPSRC grant EP/Y028783/1.

References Beatrice Achilli, Marco Benedetti, Giulio Biroli, and Marc Mézard. Theory of speciation transitions in diffusion models with general class structure. arXiv preprint arXiv:2602.04404, 2026. Brian D.O. Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313–326, May 1982. ISSN 0304-4149. doi: 10.1016/0304-4149(82)90051-5. Vahan Arsenyan, Elen Vardanyan, and Arnak Dalalyan. Assessing the quality of denoising diffusion models in Wasserstein distance: noisy score and optimal bounds. Advances in Neural Information Processing Systems, 38:19548–19591, 2026. Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly 𝑑-linear convergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learning Representations, 2024. Giulio Biroli, Tony Bonnaire, Valentin De Bortoli, and Marc Mézard. Dynamical regimes of diffusion models. Nature Communications, 15(1):9957, 2024. Tony Bonnaire, Raphaël Urfin, Giulio Biroli, and Marc Mezard. Why diffusion models don’t memorize: The role of implicit dynamical regularization in training. In D. Belgrave, C. Zhang, H. Lin, R. Pascanu, P. Koniusz, M. Ghassemi, and N. Chen, editors, Advances in Neural Information Processing Systems, volume 38, pages 141266–141286. Curran Associates, Inc., 2025. Miha Brešar and Aleksandar Mijatović. Nonasymptotic bounds for forward processes in denoising diffusions: Ornstein–uhlenbeck is hard to beat. The Annals of Applied Probability, 35(6):4439–4463, 2025.

26

Y. Iguchi and P. Fearnhead

Bradley CA Brown, Anthony L Caterini, Brendan Leigh Ross, Jesse C Cresswell, and Gabriel LoaizaGanem. Verifying the union of manifolds hypothesis for image data. arXiv preprint arXiv:2207.02862, 2022. Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning, pages 4735–4763. PMLR, 2023a. Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. International Conference on Learning Representations, 2023b. Giovanni Conforti, Alain Durmus, and Marta Gentiloni Silveri. KL convergence guarantees for score diffusion models under minimal data assumptions. SIAM Journal on Mathematics of Data Science, 7 (1):86–109, 2025. Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. Transactions on Machine Learning Research, 2022. Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780–8794, 2021. Benjamin Dupuis, Dario Shariatian, Maxime Haddouche, Alain Durmus, and Umut Simsekli. Algorithmand data-dependent generalization bounds for diffusion models. In D. Belgrave, C. Zhang, H. Lin, R. Pascanu, P. Koniusz, M. Ghassemi, and N. Chen, editors, Advances in Neural Information Processing Systems, volume 38, pages 101897–101936. Curran Associates, Inc., 2025. Paul Dupuis and Richard S. Ellis. A Weak Convergence Approach to the Theory of Large Deviations. Wiley Series in Probability and Statistics. John Wiley & Sons, New York, 1997. doi: 10.1002/9781118165904. Xuefeng Gao, Hoang M Nguyen, and Lingjiong Zhu. Wasserstein convergence guarantees for a general class of score-based generative models. Journal of Machine Learning Research, 26(43):1–54, 2025. U. G. Haussmann and E. Pardoux. Time reversal of diffusions. The Annals of Probability, 14(4), October 1986. ISSN 0091-1798. doi: 10.1214/aop/1176992362. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020. Daniel Hsu, Sham Kakade, and Tong Zhang. A tail inequality for quadratic forms of subGaussian random vectors. Electronic Communications in Probability, 17(none), January 2012. ISSN 1083-589X. doi: 10.1214/ecp.v17-2079. URL http://dx.doi.org/10.1214/ecp.v17-2079. Zhihan Huang, Yuting Wei, and Yuxin Chen. Denoising diffusion probabilistic models are optimally adaptive to unknown low dimensionality. Mathematics of Operations Research, March 2026. ISSN 1526-5471. doi: 10.1287/moor.2024.0769. Yuta Koike. Wasserstein bounds for denoising diffusion probabilistic models via the Föllmer process. arXiv preprint arXiv:2605.18069, 2026. Gen Li and Yuling Yan. Adapting to unknown low-dimensional structures in score-based diffusion models. Advances in Neural Information Processing Systems, 37:126297–126331, 2024. Gen Li and Yuling Yan. 𝑂 (𝑑/𝑇) convergence theory for diffusion probabilistic models under minimal assumptions. Journal of Machine Learning Research, 26(292):1–55, 2025. Gen Li, Changxiao Cai, and Yuting Wei. Dimension-free convergence of diffusion models for approximate Gaussian mixtures. 2026. Puheng Li, Zhong Li, Huishuai Zhang, and Jiang Bian. On the generalization properties of diffusion models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 2097–2127. Curran Associates, Inc., 2023.

Diffusion Models for Clustered Data

27

Jiadong Liang, Zhihan Huang, and Yuxin Chen. Low-dimensional adaptation of diffusion models: Convergence in total variation (extended abstract). In Nika Haghtalab and Ankur Moitra, editors, Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pages 3723–3729. PMLR, 30 Jun–04 Jul 2025. Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11461–11471, 2022. Antoine Maillard and Sebastian Goldt. Memorisation, convergence and generalisation in generative models, 2026. Samet Oymak and Talha Cihad Gulcu. A theoretical characterization of semi-supervised learning with self-training for Gaussian mixture models. In International Conference on Artificial Intelligence and Statistics, pages 3601–3609. PMLR, 2021. Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008. Vadim Popov, Ivan Vovk, Vladimir Gogoryan, Tasnima Sadekova, and Mikhail Kudinov. Grad-tts: A diffusion probabilistic model for text-to-speech. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8599–8608. PMLR, 18–24 Jul 2021. Peter Potaptchik, Iskander Azangulov, and George Deligiannidis. Linear convergence of diffusion models under the manifold hypothesis. In The Thirty Eighth Annual Conference on Learning Theory, pages 4668–4685. PMLR, 2025. Rahul Satija, Jeffrey A. Farrell, David Gennert, Alexander F. Schier, and Aviv Regev. Spatial reconstruction of single-cell gene expression data. Nature Biotechnology, 33(5):495–502, 2015. Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of Gaussians using the DDPM objective. Advances in Neural Information Processing Systems, 36:19636–19649, 2023. Marta Gentiloni Silveri and Antonio Ocello. Beyond log-concavity and score regularity: Improved convergence bounds for score-based generative models in W2-distance. In Forty-second International Conference on Machine Learning, 2025. Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. Tim Stuart, Andrew Butler, Paul Hoffman, Christoph Hafemeister, Efthymia Papalexi, William M. Mauck, Yuhan Hao, Marlon Stoeckius, Peter Smibert, and Rahul Satija. Comprehensive integration of single-cell data. Cell, 177(7):1888–1902.e21, 2019. Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi S. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. In The Eleventh International Conference on Learning Representations, 2023. Peng Wang, Huijie Zhang, Zekai Zhang, Siyi Chen, Yi Ma, and Qing Qu. Diffusion Models Learn Low-Dimensional Distributions VIA Subspace Clustering. In 2025 IEEE 10th International Workshop on Computational Advances in Multi-Sensor Adaptive Processing (CAMSAP), page 211–215, 2025. F. Alexander Wolf, Philipp Angerer, and Fabian J. Theis. SCANPY: large-scale single-cell gene expression data analysis. Genome Biology, 19(1):15, 2018.

28

Y. Iguchi and P. Fearnhead

Yuchen Wu, Minshuo Chen, Zihao Li, Mengdi Wang, and Yuting Wei. Theoretical insights for diffusion guidance: A case study for Gaussian mixture models. In International Conference on Machine Learning, pages 53291–53327. PMLR, 2024. Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1–39, 2023.

Appendix A. Property of Time-Discretisation Schedule To show Lemma 3.3, we first prove the following result, which is also key in the proof of Theorem 2: Lemma A.1. For the time-discretisation schedule (20), it holds that: 𝜎𝑇2 −𝑡

𝑗 𝜎𝑇2 −𝑡 𝑗+1

SNR𝑇 −𝑡 𝑗+1 SNR𝑇 −𝑡 𝑗 ≤ 4,

≤ 2,

𝑗 = 0, . . . , 𝑁 − 1.

Proof. Noticing that ℎ𝑖+1 = 𝜏𝜎𝑇2 −𝑡𝑖 , 𝑖 = 0, 1, . . . , 𝑁 − 2 and ℎ 𝑁 ≤ 𝜏𝜎𝑇2 −𝑡𝑁 −1 , we have that for all 𝑗 = 0, 1, . . . , 𝑁 − 1, 𝜎𝑇2 −𝑡 𝑗 − 𝜎𝑇2 −𝑡 𝑗+1 =

∫ 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗+1

2 𝑑 2 𝑑𝑠 𝜎𝑠 𝑑𝑠 ≤ 2ℎ 𝑗+1 ≤ 2𝜏𝜎𝑇 −𝑡 𝑗 ,

which leads to 𝜎𝑇2 −𝑡 𝑗 /𝜎𝑇2 −𝑡 𝑗+1 ≤ 1/(1 − 2𝜏) ≤ 2 for 𝜏 ∈ (0, 1/4]. For the latter inequality, using
𝑚 𝑇2 −𝑡 𝑗+1 /𝑚 𝑇2 −𝑡 𝑗 ≤ exp 2ℎ 𝑗+1 ≤ exp(2𝜏), we get 2 𝑚𝑇 𝜎𝑇2 −𝑡 SNR𝑇 −𝑡 𝑗+1 −𝑡 𝑗+1 𝑗 · = SNR𝑇 −𝑡 𝑗 𝑚2 𝜎2 𝑇 −𝑡 𝑗

≤ 2 exp(2𝜏) ≤ 4,

𝑇 −𝑡 𝑗+1

which concludes the proof.

□

∫𝑡 (Proof of Lemma 3.3) We consider the term 𝑇 𝑗 ≡ 𝑡 𝑗+1 𝜎𝑑𝑠 , 𝑗 = 0, 1, . . . , 𝑁 − 1. Because 𝑠 ↦→ 1/𝜎𝑇2 −𝑠 2 𝑗

𝑇 −𝑠

is an increasing function, (20) and Lemma A.1 give that 𝑇 𝑗 ≤ 𝜏𝜎𝑇2 −𝑡 𝑗 /𝜎𝑇2 −𝑡 𝑗+1 ≤ 2𝜏, 𝑗 = 0, 1, . . . , 𝑁 − 1 where for 𝑗 = 𝑁 − 1 we considered ℎ 𝑁 ≤ 𝜏𝜎𝑇2 −𝑡𝑁 −1 . On the other hand, we have (

𝜏

𝑗 = 0, 1, . . . , 𝑁 − 2;

0

𝑗 = 𝑁 − 1.

Í 𝑁 −1

𝑇 𝑗 ≤ 2𝑁 · 𝜏 and

𝑇𝑗 ≥ Taking summation, we get 𝜏(𝑁 − 1) ≤ 𝑁 −1 ∑︁ 𝑗=0

∫ 𝑇−𝜀 𝑇𝑗 = 0

which concludes Lemma 3.3.

𝑗=0

1 − exp(−2𝑇) 𝑑𝑠 = 𝑇 − 𝜀 + 21 log = Θ 𝑇 − 1 + log 1𝜀 𝜎𝑇2 −𝑠 1 − exp(−2𝜀)

,

Diffusion Models for Clustered Data

29

B. Proofs for Posterior Classification Analysis

B.1. Proof of Lemma 3.1

The definition of posterior weight yields log 𝑊𝐷[𝑖 ] (𝑡, 𝑦) = log Í

−𝑝  𝜋𝑖 ·→ (𝑖, 𝑗 ) 𝑇 −𝑡 ( 𝑦) ≤ min log 𝜋𝜋𝑖𝑗 + 𝑍𝑡 (𝑦), 0 , [ℓ ] → − 𝜋 ·𝑝 ( 𝑦) [𝑖 ]

ℓ ∈ [𝐾 ]

(𝑖, 𝑗 )

(𝑦) = log

ℓ

𝑦 ∈ R𝐷 , 𝑗 ∈ [𝐾],

𝑇 −𝑡

N ( 𝑦; 𝑚𝑇 −𝑡 𝜇𝑖 ,𝑆𝑇𝑖 −𝑡 )

[ 𝑗]

. The conditional random variable can be expressed as 𝑌e𝑇 −𝑡 ∼ → − ← − 𝑗 𝑗 N (𝑚 𝑇 −𝑡 𝜇 𝑗 , 𝑆𝑇 −𝑡 ) due to Law( 𝑋 𝑡 ) = Law( 𝑋 𝑇 −𝑡 ), where we recall 𝑆𝑇 −𝑡 in (5). We then decompose (𝑖, 𝑗 ) [ 𝑗] 𝑍𝑡 (𝑌e𝑇 −𝑡 ) = Ψ1 + Ψ2 + Ψ3 where

with 𝑍𝑡

𝑗

N ( 𝑦; 𝑚𝑇 −𝑡 𝜇 𝑗 ,𝑆𝑇 −𝑡 )

[ 𝑗] Ψ1 ≡ −𝑚 𝑇 −𝑡 (𝑌e𝑇 −𝑡 − 𝑚 𝑇 −𝑡 𝜇 𝑗 ) ⊤ (𝑆𝑇𝑖 −𝑡 ) −1 (𝜇 𝑗 − 𝜇𝑖 );
[ 𝑗] [ 𝑗] 𝑗 Ψ2 ≡ 12 (𝑌e𝑇 −𝑡 − 𝑚 𝑇 −𝑡 𝜇 𝑗 ) ⊤ (𝑆𝑇 −𝑡 ) −1 − (𝑆𝑇𝑖 −𝑡 ) −1 (𝑌e𝑇 −𝑡 − 𝑚 𝑇 −𝑡 𝜇 𝑗 ); det 𝑆

𝑗

Ψ3 ≡ 21 log det 𝑆𝑇𝑖 −𝑡 − 𝑇 −𝑡

2 𝑚𝑇 ⊤ 𝑖 −1 −𝑡 2 (𝜇 𝑗 − 𝜇𝑖 ) (𝑆𝑇 −𝑡 ) (𝜇 𝑗 − 𝜇𝑖 ).

The first term follows a Gaussian distribution with mean 0 and variance given as [𝑖, 𝑗 ]

Var[Ψ1 ] = 𝑚 𝑇2 −𝑡 (𝜇 𝑗 − 𝜇𝑖 ) ⊤ (𝑆𝑇𝑖 −𝑡 ) −1 (𝑆𝑇 −𝑡 ) (𝑆𝑇𝑖 −𝑡 ) −1 (𝜇 𝑗 − 𝜇𝑖 ) ≡ 𝒯𝐷 𝑗

(𝑡).

Thus, the standard high probability argument for the Gaussian variable gives: for 𝛿 ∈ (0, 1), √︃   P |Ψ1 | ≥ 2Var[Ψ1 ] log 2𝛿 ≤ 2𝛿 . [ 𝑗] 𝑗 The second term Ψ2 , quadratic w.r.t. the Gaussian 𝑌e𝑇 −𝑡 − 𝑚 𝑇 −𝑡 𝜇 𝑗 ∼ N (0, 𝑆𝑇 −𝑡 ) is expressed as ⊤ Ψ2 = 𝑧 𝐴𝑧 with

𝑗 𝑗 𝐴 ≡ 21 𝐼 𝐷 − (𝑆𝑇 −𝑡 ) 1/2 (𝑆𝑇𝑖 −𝑡 ) −1 (𝑆𝑇 −𝑡 ) 1/2 ,

𝑧 ∼ N (0, 𝐼 𝐷 ).

Then, application of the concentration inequality [Hsu et al., 2012, Proposition 1] to Ψ2 yields: for 𝛿 ∈ (0, 1), √︃   P Ψ2 ≥ tr( 𝐴) + 4tr( 𝐴2 ) log 2𝛿 + 2∥ 𝐴∥ 2 log 2𝛿 ≤ 2𝛿 . Thus with probability at least 1 − 𝛿, we have: √︃ √︃ (𝑖, 𝑗 ) [ 𝑗] 𝑍𝑡 (𝑌e𝑇 −𝑡 ) ≤ Ψ3 + tr( 𝐴) + 2Var[Ψ1 ] log 2𝛿 + 2Var[Ψ2 ] log 2𝛿 + 2∥ 𝐴∥ 2 log 2𝛿 √︃ √︃ [𝑖, 𝑗 ] [𝑖, 𝑗 ] [ 𝑗] [𝑖, 𝑗 ] [𝑖 ]
→ − → − 2 = −KL 𝑝 𝑇 −𝑡 || 𝑝 𝑇 −𝑡 + 2𝒯𝐷 (𝑡) log 𝛿 + 2𝒰𝐷 (𝑡) log 2𝛿 + 2𝒱𝐷 (𝑡) log 2𝛿 , [𝑖, 𝑗 ]

where we used: Var[Ψ2 ] = 2tr( 𝐴2 ) ≡ 𝒰𝐷

(𝑡) and

−𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
, Ψ3 + tr( 𝐴) = −KL → 𝑇 −𝑡 𝑇 −𝑡 which completes the proof of Lemma 3.1.

30

Y. Iguchi and P. Fearnhead

B.2. Proof of Theorem 1

[𝑖, 𝑗 ] Recall the definition of 𝐶𝐷 in (9). We set 𝛿 = 𝐷 −𝑞 , 𝑞 ≥ 2 and write 𝑟 ≡ log 2𝐾 𝛿 = Θ log(𝐾 𝐷) , and [𝑖, 𝑗 ] [𝑖, 𝑗 ] [𝑖, 𝑗 ] 𝑆 = SNR𝑇 −𝑡 for convenience. We also write 𝒯(𝑡) = 𝒯𝐷 (𝑡), 𝒰(𝑡) = 𝒰𝐷 (𝑡) and 𝒱(𝑡) = 𝒱𝐷 (𝑡) whose definitions are found in (9). Then, the inequality [𝑖, 𝑗 ]

𝐶𝐷

−𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
(𝑡, 𝛿/𝐾) ≥ 𝜂KL → 𝑇 −𝑡 𝑇 −𝑡

is equivalent to (1 − 𝜂) ≥ 𝐺 (𝑡) with 𝐺 (𝑡) = 𝑁 (𝑡)/𝐷 (𝑡), where √︁ √︁ 𝑁 (𝑡) ≡ log 𝜋𝜋𝑖𝑗 + 2𝒯(𝑡)𝑟 + 2𝒰(𝑡)𝑟 + 2𝒱(𝑡)𝑟; −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
. 𝐷 (𝑡) ≡ KL → 𝑇 −𝑡 𝑇 −𝑡

(55)

(56) (57)

We will derive a sufficient condition for (55). To this end, we consider a simplified upper bound for 𝐺 (𝑡), e(𝑡), that involves only 𝑆 = SNR𝑇 −𝑡 as a function of time 𝑡, and then solve the inequality denoted by 𝐺 e(𝑡), which is sufficient for the original condition (55). First, we consider an upper bound of the 1−𝜂 ≥𝐺 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
= KL → − [ 𝑗] → − [𝑖 ]
→ − [ 𝑗] → − [𝑖 ]
term 𝑁 (𝑡). Noticing that KL → 𝑚 𝑝 𝑇 −𝑡 || 𝑝 𝑇 −𝑡 + KL𝑐 𝑝 𝑇 −𝑡 || 𝑝 𝑇 −𝑡 with 𝑇 −𝑡 𝑇 −𝑡 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
:= 𝑚𝑇2 −𝑡 ∥𝜇 − 𝜇 ∥ 2 KL𝑚 → 𝑖 𝑗 (𝑆 𝑖 2 𝑇 −𝑡 𝑇 −𝑡

𝑇 −𝑡 )

−1

;

  𝑗 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
:= 1 tr 𝑆 𝑗 (𝑆 𝑖 ) −1
− 𝐷 − log det 𝑆𝑇 −𝑡 , KL𝑐 → 𝑖 𝑇 −𝑡 𝑇 −𝑡 𝑇 −𝑡 𝑇 −𝑡 det 𝑆𝑇 −𝑡 2 we deduce from the bound (13) in Lemma 3.2 that √︃ √︁
−𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
2𝒯(𝑡)𝑟 ≤ 4𝑟 1 + 𝑆 · ∥Σ𝑖 − Σ 𝑗 ∥ 2 KL𝑚 → 𝑇 −𝑡 𝑇 −𝑡
2𝑟 1+𝑆· ∥Σ𝑖 −Σ 𝑗 ∥ 2 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
, ≤ + 1−2 𝜂 KL𝑚 → (1− 𝜂) 𝑇 −𝑡 𝑇 −𝑡

(58) (59)

(60)

√ (1− 𝜂) 𝑏 𝑎 where we used Arithmetic/Geometric Mean inequality: 𝑎𝑏 ≤ 2(1− , 𝑎, 𝑏 > 0. Similarly, we 𝜂) + 2 obtain from Lemma C.1 in Appendix C that
√︁ 2𝑟 1+𝑆· ∥Σ 𝑗 ∥ 2 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
. + 1−2 𝜂 KL𝑐 → 2𝒰(𝑡)𝑟 ≤ (61) (1− 𝜂) 𝑇 −𝑡 𝑇 −𝑡 Substitution of the bounds (60–61) and (15) into (56) with Assumptions 1, 2, 3 leads to   1− 𝜂 → − [ 𝑗] → − [𝑖 ]
𝛼 10−4𝜂 4𝑟 + 𝐶 𝐷 𝑁 (𝑠) ≤ 𝐶1 + 1− 2 𝜂 1− 𝜂 𝑟 · 𝑆 + 2 KL 𝑝 𝑇 −𝑡 || 𝑝 𝑇 −𝑡 , where the constants 𝐶1 , 𝐶2 are collected in Assumptions 1 and 2. Thus, we obtain    

10−4𝜂 10−4𝜂 4𝑟 4𝑟 𝐶1 + 1− 𝜂 +𝐶2 𝐷 𝛼 1− 𝜂 𝑟 ·𝑆 𝐶1 + 1− 𝜂 +𝐶2 𝐷 𝛼 1− 𝜂 𝑟 ·𝑆 1− 𝜂 e(𝑡),
𝐺 (𝑡) ≤ + 2 ≤ + 1−2 𝜂 ≡ 𝐺  2 −𝑝 [ 𝑗 ] | | → −𝑝 [𝑖 ] KL → 𝑆 𝐶3 𝐷2 𝑇 −𝑡 𝑇 −𝑡 2 𝐶5 · 𝑆·tr(Σ𝑖 )+𝐷

where in the second inequality we used Lemma C.2 under Assumptions 3 and 4, provided in Appendix C. e(𝑡) holds, which is simplified as: Then, 1 − 𝜂 ≥ 𝐺 (𝑡) is satisfied if 1 − 𝜂 ≥ 𝐺 e2 𝐷 2 1− 𝜂 𝐶 𝑆 2 · 𝑆·tr(Σ𝑖 )+𝐷 · 2 ≥ ℬ + 𝒞𝑆,

Diffusion Models for Clustered Data

31

where we have set: 4𝑟 ℬ ≡ 𝐶1 + 1− 𝜂 > 0,

𝒞 ≡ 𝐶2 𝐷 𝛼



10−4𝜂 1− 𝜂



𝑟 > 0,

e ≡ 𝐶3 > 0. 𝐶 𝐶5

Rearranging the terms with respect to 𝑆, we obtain the following quadratic inequality: 𝒞 tr(Σ𝑖 ) 𝑆 2 − 𝒜𝑆 + ℬ𝐷 ≤ 0,

where

e2 𝐷 2 − 𝒞𝐷 − ℬ tr(Σ𝑖 ). 𝒜 ≡ 1−4 𝜂 𝐶

(62)

Under Assumption 2, we have tr(Σ𝑖 ) =

𝐷 ∑︁

𝜆 𝑘 (Σ𝑖 ) ≤ 𝐶2 𝑀𝑖 𝐷 𝛼 + 𝜈𝑖 ≤ 𝐶2 (𝐷 𝛼 + 𝑅) 𝑀 †

(63)

𝑘=1

where the constants 𝐶2 > 0, 𝜈𝑖 ≥ 0, 𝑅 ∈ [0, 1) are collected in Assumption 2. Thus, there exists 𝐷 1 ∈ N such that for all 𝐷 ≥ 𝐷 1 , we have the discriminant strictly positive, i.e., Δ ≡ 𝒜 2 − 4𝒞tr(Σ𝑖 )ℬ𝐷 > 0. This is because 𝒜 2 = Θ(𝐷 4 ) and 𝒞tr(Σ𝑖 )ℬ𝐷 = O (𝐷 1+2𝛼 (log 𝐷) 2 ), and thus the leading term dominates the lower-order   term for any 𝛼 ∈ [0, 1). Under Δ > 0, solving (62) under 𝑆 ≥ SNR𝑇 leads to 𝑆 ∈ max{SNR𝑇 , 𝑆}, 𝑆 , where √ 𝒜− Δ , 𝑆 = 2𝒞tr(Σ 𝑖)

√ 𝒜+ Δ . 𝑆 = 2𝒞tr(Σ 𝑖)

First, we note from (63) that for all sufficiently large 𝐷, 2

2−2𝛼

𝒜 𝑆 > 2𝒞tr(Σ ≳ 𝐷 2𝛼𝐷log 𝐷 = 𝐷log 𝐷 , 𝑖)

where the right-hand side is an increasing function of 𝐷 since 𝛼 < 1. Therefore, we have: for any fixed early stopping criterion 𝜀 ∈ (0, 1), there exists 𝐷 2 such that for all 𝐷 ≥ 𝐷 2 , the upper root 𝑆 is strictly larger than the early stopping threshold 𝑆 𝜀 = SNR 𝜀 = 𝑚 2𝜀 /𝜎𝜀2 , which is independent of 𝐷. Notice that 𝑆 = SNR𝑇 −𝑡 is a non-decreasing function of the physical time 𝑡. Thus, we have: for all 𝐷 ≥ 𝐷 † = max{𝐷 1 , 𝐷 2 }, the target condition (55) holds for all 𝑡 ∈ [0, 𝑇 − 𝜀] satisfying 𝑆 SNR𝑇 −𝑡 ∈ max{SNR𝑇 , 𝑆}, SNR 𝜀 , equivalently to 𝑡 ∈ max{0, 𝑇 + 12 log 1+𝑆 }, 𝑇 − 𝜀 . We now study √ the scale of the threshold 𝑆. We multiply the numerator and denominator by the conjugate 𝒜 + Δ and have from Δ > 0 and the definition of terms ℬ and 𝒜 that     𝐷 log(𝐾 𝐷) log(𝐾 𝐷) √ ≤ 2ℬ𝐷 = O = O , (64) 𝑆 = 2ℬ𝐷 𝒜 𝐷 𝐷2 𝒜+ Δ

log(𝐾 𝐷) where we used 𝑟 = Θ log(𝐾 𝐷) . Similarly, we have 𝑆 ≥ ℬ𝐷 . The first part of the statement 𝒜 ≳ 𝐷 has now been proved. Finally, we prove the second part of the statement. For any 𝛽 ∈ (0, 1), directly solving −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]

≤ 𝛽 gives exp −𝜂KL → 𝑇 −𝑡 𝑇 −𝑡 −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
≥ 1 log 1 KL → 𝜂 𝛽 𝑇 −𝑡 𝑇 −𝑡

(65)

   𝑆 under 𝑡 ∈ max 0, 𝑇 + 21 log 1+𝑆 , 𝑇 − 𝜀 . We see that the inequality (65) is satisfied if the lower bound  2 2 3 of left-hand-side given in Lemma C.2 is no lower than 𝜂1 log 𝛽1 , that is, SNR2𝑇 −𝑡 𝐶 · SNR𝑇 −𝑡𝐷·tr(Σ𝑖 )+𝐷 ≥ 𝐶5 1 1 𝜂 log 𝛽 . Solving the latter inequality with respect to SNR𝑇 −𝑡 under the range SNR𝑇 −𝑡 ∈ [𝑆, SNR 𝜀 ]

32

Y. Iguchi and P. Fearnhead

leads to (19), and then the second part of the statement has been shown due to the high probability bound (17). Proof of Theorem 1 is now complete.

C. Auxiliary Results for Theorem 1 C.1. First Result Lemma C.1. Let 𝑖, 𝑘 ∈ [𝐾] and 𝑌 = (𝑆𝑇𝑖 −𝑡 ) 1/2 (𝑆𝑇𝑘 −𝑡 ) −1 (𝑆𝑇𝑖 −𝑡 ) 1/2 − 𝐼 𝐷 . We have:   −𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
, tr(𝑌 2 ) ≤ 4 SNR𝑇 −𝑡 ∥Σ𝑖 ∥ 2 + 1 · KL𝑐 → 𝑇 −𝑡 𝑇 −𝑡 −𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
is defined in (59). where KL𝑐 → 𝑇 −𝑡 𝑇 −𝑡 Proof. Let {𝜂𝑙 }𝑙=1,...,𝐷 be the eigenvalues of (𝑆𝑇𝑖 −𝑡 ) 1/2 (𝑆𝑇𝑘 −𝑡 ) −1 (𝑆𝑇𝑖 −𝑡 ) 1/2 . Then, the eigenÍ𝐷 2 2 𝑖 values of 𝑌 are {𝜂𝑙 − 1}𝑙 , yielding tr(𝑌 ) = Because 𝑆𝑇 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 is sim𝑙=1 (𝜂𝑙 − 1) . 𝑖 1/2 𝑘 −1 𝑖 1/2 ilar to (𝑆𝑇 −𝑡 ) (𝑆𝑇 −𝑡 ) (𝑆𝑇 −𝑡 ) , they share the identical eigenvalues 𝜂𝑙 . We thus get [𝑖 ] [𝑘]
→ − → − 1 Í𝐷 KL𝑐 𝑝 𝑇 −𝑡 || 𝑝 𝑇 −𝑡 = 2 𝑙=1 (𝜂𝑙 − 1 − log 𝜂𝑙 ). The maximum eigenvalue 𝜂1 is bounded by 𝑆𝑇𝑖 −𝑡 2 (𝑆𝑇𝑘 −𝑡 ) −1 2 ≤ 𝐶 𝜂 (𝑡) with 𝐶 𝜂 (𝑡) ≡ (SNR𝑇 −𝑡 ∥Σ𝑖 ∥ 2 + 1). By elementary calculus, the algebraic inequality (𝜂 − 1) 2 ≤ 2𝐶 𝜂 (𝑡) (𝜂 − 1 − log 𝜂) holds strictly for all 𝜂 ∈ (0, 𝐶 𝜂 (𝑡)]. Applying this inequality to each eigenvalue 𝜂𝑙 and summing them completes the proof. □ C.2. Second Result Lemma C.2 (Upper and lower bounds for KL divergence). It holds under Assumptions 3 and 4 that: for any 𝑖, 𝑗 ∈ [𝐾], SNR𝑇 −𝑡 · 2



𝐶3 𝐶5

2

𝐷2 SNR𝑇 −𝑡 tr(Σ𝑖 )+𝐷

(66)

−𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
≤ SNR𝑇 −𝑡 · 𝐶 𝐷 + SNR𝑇2 −𝑡 · ∥Σ − Σ ∥ 2 , ≤ KL → 4 𝑖 𝑗 𝐹 2 4 𝑇 −𝑡 𝑇 −𝑡 where the constants 𝐶3 , 𝐶4 , 𝐶5 > 0 are collected in Assumptions 3 and 4. Proof. We begin with proving the lower bound of (66). Since the covariance matrix Σ𝑖 is symmetric due to Assumption 2, its eigen-decomposition gives Σ𝑖 = 𝑈 (𝑖) Λ (𝑖) 𝑈 (𝑖) ,⊤ with the orthogonal matrix
𝑈 (𝑖) ≡ [𝑢 1 (𝑖), . . . , 𝑢 𝐷 (𝑖)] and Λ (𝑖) = diag 𝜆1 (Σ𝑖 ), . . . , 𝜆 𝐷 (Σ𝑖 ) . For notational simplicity, we write
𝑣 𝑘 = ⟨𝜇 𝑗 − 𝜇𝑖 , 𝑢 𝑘 (𝑖)⟩. Then, 𝑆𝑇𝑖 −𝑡 = 𝑈 (𝑖) 𝑚 𝑇2 −𝑡 Λ (𝑖) + 𝜎𝑇2 −𝑡 𝐼 𝐷 𝑈 (𝑖) ,⊤ and (58) yield that −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
≥ KL → −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
= 𝑚𝑇2 −𝑡 KL → 𝑚 2 𝑇 −𝑡 𝑇 −𝑡 𝑇 −𝑡 𝑇 −𝑡

𝐷 ∑︁

𝑣𝑘2 2 𝑚𝑇 −𝑡 𝜆 𝑘 (Σ𝑖 )+𝜎𝑇2 −𝑡

.

(67)

𝑘=1

Further, the Cauchy-Schwarz inequality gives 𝐷 ∑︁



Í𝐷

𝑘=1 |𝑣𝑘 |

2



Í𝐷

𝑘=1 |𝑣𝑘 |

2

𝑣𝑘2 ≥ Í𝐷 (𝑚2 𝜆 (Σ )+𝜎 2 ) = 𝑚2 tr(Σ )+𝜎 2 𝐷 . 2 𝜆 (Σ )+𝜎 2 𝑚𝑇 𝑖 𝑖 𝑖 −𝑡 𝑘 𝑇 −𝑡 𝑘=1 𝑇 −𝑡 𝑇 −𝑡 𝑇 −𝑡 𝑘 𝑇 −𝑡

(68)

𝑘=1 𝐶3 𝑘=1 |𝜈 𝑘 | ≥ 𝐶5 𝐷 under Assumption 3, which

For the numerator in the above lower bound, we have comes from

Í𝐷

𝐷 ∑︁

𝐷 ∑︁

𝑘=1

𝑣 2𝑘 = ∥𝜇𝑖 − 𝜇 𝑗 ∥ 2 ≥ 𝐶3 𝐷,

𝑘=1

𝑣 2𝑘 ≤ 𝐶5

𝐷 ∑︁ 𝑘=1

|𝑣 𝑘 |.

Diffusion Models for Clustered Data

33

We thus obtain the desired lower bound by recalling the definition of SNR: SNR𝑇 −𝑡 = 𝑚 𝑇2 −𝑡 /𝜎𝑇2 −𝑡 . We now show the upper bound of (66). The upper bound for the mean part of the KL (58) is immediately obtained from (67) with Assumption 3 as: −𝑝 [ 𝑗 ] || → −𝑝 [𝑖 ]
≤ 𝑚𝑇2 −𝑡 KL𝑚 → 𝑇 −𝑡 𝑇 −𝑡 2𝜎 2

𝐷 ∑︁

𝑣 2𝑘 = SNR2𝑇 −𝑡 ∥𝜇𝑖 − 𝜇 𝑗 ∥ 2 ≤ SNR2𝑇 −𝑡 · 𝐶4 𝐷.

𝑇 −𝑡

𝑘=1

We now consider the covariance part of the KL divergence, given in (59). Introducing 𝑆e𝜆 = 𝐼 𝐷 + 𝜆(𝑆𝑇𝑖 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 − 𝐼 𝐷 ), 𝜆 ∈ [0, 1], we use standard matrix calculus [Petersen et al., 2008] to get:
log det 𝑆𝑇𝑖 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 =

∫ 1 0

𝑑 e 𝑑𝜆 log det 𝑆 𝜆 𝑑𝜆 =

∫ 1 h
i tr ( 𝑆e𝜆 ) −1 𝑆𝑇𝑖 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 − 𝐼 𝐷 𝑑𝜆. 0

Thus, (59) is expressed as: ∫ 1 h

i tr 𝐼 𝐷 − ( 𝑆e𝜆 ) −1 𝑆𝑇𝑖 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 − 𝐼 𝐷 𝑑𝜆 0 ∫ 1 h
2i 𝜆 · tr ( 𝑆e𝜆 ) −1 𝑆𝑇𝑖 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 − 𝐼 𝐷 𝑑𝜆. = 21

−𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
= 1 KL𝑐 → 2 𝑇 −𝑡 𝑇 −𝑡

(69)

0

The expression 𝑆e𝜆 = 𝜆𝑆𝑇𝑖 −𝑡 + (1 − 𝜆)𝑆𝑇𝑘 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 and the cycliciy of the trace operator yield h
2i tr ( 𝑆e𝜆 ) −1 𝑆𝑇𝑖 −𝑡 (𝑆𝑇𝑘 −𝑡 ) −1 − 𝐼 𝐷 i h
−1 = tr 𝜆𝑆𝑇𝑖 −𝑡 + (1 − 𝜆)𝑆𝑇𝑘 −𝑡 (𝑆𝑇𝑖 −𝑡 − 𝑆𝑇𝑘 −𝑡 ) (𝑆𝑇𝑘 −𝑡 ) −1 (𝑆𝑇𝑖 −𝑡 − 𝑆𝑇𝑘 −𝑡 ) ≤ ∥𝑆𝑇𝑖 −𝑡 − 𝑆𝑇𝑘 −𝑡 ∥ 2𝐹 · ∥ (𝑆𝑇𝑘 −𝑡 ) −1 ∥ 2 · ∥ (𝜆𝑆𝑇𝑖 −𝑡 + (1 − 𝜆)𝑆𝑇𝑘 −𝑡 ) −1 ∥ 2 ≤ ∥𝑆𝑇𝑖 −𝑡 − 𝑆𝑇𝑘 −𝑡 ∥ 2𝐹 · 𝜎 41 ,

(70)

𝑇 −𝑡

where we used (𝑆𝑇𝑘 −𝑡 ) −1 ⪯ 𝜎 21 𝐼 𝐷 and the operator convexity of 𝑧 ↦→ 𝑧 −1 for positive definite matrices 𝑇 −𝑡 to get (𝜆𝑆𝑇𝑖 −𝑡 + (1 − 𝜆)𝑆𝑇𝑘 −𝑡 ) −1 ⪯ 𝜆(𝑆𝑇𝑖 −𝑡 ) −1 + (1 − 𝜆) (𝑆𝑇𝑘 −𝑡 ) −1 ⪯ 𝜎 21 𝐼 𝐷 . 𝑇 −𝑡

Substitution of (70) into (69) together with 𝑆𝑇𝑖 −𝑡 − 𝑆𝑇𝑘 −𝑡 = 𝑚 𝑇2 −𝑡 (Σ𝑖 − Σ 𝑘 ) yields −𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
≤ SNR𝑇2 −𝑡 ∥Σ − Σ ∥ 2 , KL𝑐 → 𝑖 𝑘 𝐹 4 𝑇 −𝑡 𝑇 −𝑡 which completes the proof.

□

D. Proof of Lemma 3.2 [𝑖, 𝑗 ]

Proof of (15). We start with the easiest term 𝒱𝐷

(𝑡) = ∥𝐼 𝐷 − 𝑆𝑇𝑖 −𝑡 (𝑆𝑇 −𝑡 ) −1 ∥ 2 . We have 𝑗

𝑗
𝑗 𝑗 𝑗 𝐼 𝐷 − 𝑆𝑇𝑖 −𝑡 (𝑆𝑇 −𝑡 ) −1 = 𝑆𝑇 −𝑡 − 𝑆𝑇𝑖 −𝑡 (𝑆𝑇 −𝑡 ) −1 = 𝑚 𝑇2 −𝑡 Σ 𝑗 − Σ𝑖 (𝑆𝑇 −𝑡 ) −1 and immediately get the desired bound from (𝑆𝑇 −𝑡 ) −1 ⪯ 𝜎 21 𝐼 𝐷 . 𝑗

𝑇 −𝑡

(71)

34

Y. Iguchi and P. Fearnhead

Proof of (13). We have [𝑖, 𝑗 ]

𝒯𝐷

2

(𝑡) := 𝑚 𝑇2 −𝑡 𝜇𝑖 − 𝜇 𝑗 (𝑆 𝑖

𝑇 −𝑡 )

𝑖 −1 −1 (𝑆 𝑗 𝑇 −𝑡 ) (𝑆𝑇 −𝑡 )

2

2

= 𝑚 𝑇2 −𝑡 𝜇𝑖 − 𝜇 𝑗 (𝑆 𝑖

𝑇 −𝑡 )

−1

+ 𝑚 𝑇2 −𝑡 𝜇𝑖 − 𝜇 𝑗 (𝑆 𝑖

𝑇 −𝑡 )

𝑖 𝑖 −1 −1 −1 (𝑆 𝑗 𝑇 −𝑡 ) (𝑆𝑇 −𝑡 ) − (𝑆𝑇 −𝑡 )

.

(72)

and (𝑆𝑇𝑖 −𝑡 ) −1 (𝑆𝑇 −𝑡 ) (𝑆𝑇𝑖 −𝑡 ) −1 − (𝑆𝑇𝑖 −𝑡 ) −1
= 𝑚 𝑇2 −𝑡 (𝑆𝑇𝑖 −𝑡 ) −1 Σ 𝑗 − Σ𝑖 (𝑆𝑇𝑖 −𝑡 ) −1 ⪯ SNR𝑇 −𝑡 ∥Σ𝑖 − Σ 𝑗 ∥ 2 (𝑆𝑇𝑖 −𝑡 ) −1 . 𝑗

Substitution of this into (72) yields the desired bound. Proof of (14). The formula (71) gives [𝑖, 𝑗 ]

𝒰𝐷

(𝑡) := 21 tr 𝐼 𝐷 − 𝑆𝑇 −𝑡 (𝑆𝑇𝑖 −𝑡 ) −1 𝑗

2

=

4 𝑚𝑇 −𝑡 2 tr




2 SNR2 Σ𝑖 − Σ 𝑗 (𝑆𝑇𝑖 −𝑡 ) −2 ≤ 2𝑇 −𝑡 ∥Σ𝑖 − Σ 𝑗 ∥ 2𝐹 ,

which completes the proof of Lemma 3.2.

E. Proofs of the Auxiliary Results for Theorem 2 E.1. Proof of Lemma 3.4 → − → − From the definition of posterior covariance, we have Σ𝑇 −𝑡 (𝑦) ≡ Cov[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦] =
⊗2 → − → − → − → − E[( 𝑋 0 ) ⊗2 | 𝑋 𝑇 −𝑡 = 𝑦] − E[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦] . Using the posterior weights and (25), we express the conditional means as: 𝐾 n ∑︁
⊗2 o → − → − → − → − → − [𝑘] E[( 𝑋 0 ) ⊗2 | 𝑋 𝑇 −𝑡 = 𝑦] = 𝑊𝐷[ 𝑘 ] (𝑡, 𝑦) · Cov[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦, 𝑋 0 ∼ 𝑃data ] + 𝜇 [ 𝑘 ] (𝑡, 𝑦) ; 𝑘=1 𝐾

∑︁ → − → − E[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦] = 𝑊𝐷[ 𝑘 ] (𝑡, 𝑦)𝜇 [ 𝑘 ] (𝑡, 𝑦). 𝑘=1

Substituting these into the formula of Σ𝑇 −𝑡 (𝑦) gives Σ𝑇 −𝑡 (𝑦) =

𝐾 ∑︁

𝐾

∑︁
⊗2 → − → − → − [𝑘] 𝑊𝐷[ 𝑘 ] (𝑡, 𝑦) Cov[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦, 𝑋 0 ∼ 𝑃data ]+ 𝑊𝐷[ 𝑘 ] (𝑡, 𝑦) 𝜇 [ 𝑘 ] (𝑡, 𝑦)

𝑘=1

𝑘=1

−

𝐾 ∑︁

𝑊𝐷[ 𝑘1 ] (𝑡, 𝑦)𝑊𝐷[ 𝑘2 ] (𝑡, 𝑦)𝜇 [ 𝑘1 ] (𝑡, 𝑦)𝜇 [ 𝑘2 ] (𝑡, 𝑦) ⊤

𝑘1 ,𝑘2 =1

=

𝐾 ∑︁

→ − → − → − [𝑘] 𝑊𝐷[ 𝑘 ] (𝑡, 𝑦) Cov[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦, 𝑋 0 ∼ 𝑃data ]

𝑘=1

+

𝐾
⊗2 1 ∑︁ 𝑊𝐷[ 𝑘1 ] (𝑡, 𝑦)𝑊𝐷[ 𝑘2 ] (𝑡, 𝑦) 𝜇 [ 𝑘1 ] (𝑡, 𝑦) − 𝜇 [ 𝑘2 ] (𝑡, 𝑦) . 2 𝑘 ,𝑘 =1 1

2

Finally, the conditional covariance formula for joint Gaussian distributions deduces that → − → − → − [𝑘] Cov[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑦, 𝑋 0 ∼ 𝑃data ] = Σ 𝑘 − 𝑚 𝑇2 −𝑡 Σ 𝑘 (𝑆𝑇𝑘 −𝑡 ) −1 Σ 𝑘 = 𝜎𝑇2 −𝑡 Σ 𝑘 (𝑆𝑇𝑘 −𝑡 ) −1 , which concludes Lemma 3.4.

Diffusion Models for Clustered Data

35

E.2. Proof of Lemma 5.2 First, we have for 𝜌 > 0 i  h  ] [𝑖 ] 2 [ 𝑘2 ] ) − 𝜇 (𝑍 ) log E max exp 𝜌 · 𝜇 [ 𝑘1 ] (𝑍 [𝑖 𝑗 𝑗 𝑘1 ,𝑘2 ∈ [𝐾 ]   ∑︁  ] [𝑖 ] 2
 [ 𝑘2 ] ) − 𝜇 (𝑍 ) ≤ log E exp 𝜌 · 𝜇 [ 𝑘1 ] (𝑍 [𝑖 𝑗 𝑗 𝑘1 ,𝑘2 ∈ [𝐾 ]

≤ log



∑︁

Ö

√︃ 
 ] [𝑖 ] (𝑍 [𝑖 ] ) 2 . ) − 𝜇 E exp 4𝜌 · 𝜇 [ 𝑘 ] (𝑍 [𝑖 𝑗 𝑗

𝑘1 ,𝑘2 ∈ [𝐾 ] 𝑘 ∈ { 𝑘1 ,𝑘2 }

 ] [ 𝑘2 ] (𝑍 [𝑖 ] ) ∥ 2 ≤ 2 ∥𝜇 [ 𝑘1 ] (𝑍 [𝑖 ] ) − 𝜇 [𝑖 ] (𝑍 [𝑖 ] ) ∥ 2 + where in the last line we used ∥𝜇 [ 𝑘1 ] (𝑍 [𝑖 𝑗 ) − 𝜇 𝑗 𝑗 𝑗 ] [𝑖 ] 2 [𝑖 ] ∥𝜇 [ 𝑘2 ] (𝑍 [𝑖 and Cauchy-Schwartz inequality. For notational simplicity, we write 𝑗 ) − 𝜇 (𝑍 𝑗 ) ∥ → − → − → − [𝑘] 2 𝑟 𝑗 = 𝜎𝑇 −𝑡 𝑗 /𝑚 𝑇 −𝑡 𝑗 and recall 𝜇 [ 𝑘 ] (𝑡, 𝑥) = E[ 𝑋 0 | 𝑋 𝑇 −𝑡 = 𝑥, 𝑋 0 ∼ 𝑃data ]. Then, Tweedie’s formula for −𝑝 [ 𝑘 ] (𝑥), (𝑡, 𝑥) ∈ [0, 𝑇) × R𝐷 gives ∇ log→ 𝑇 −𝑡

−𝑝 [ 𝑘 ] (𝑥)
𝜇 [ 𝑘 ] (𝑡, 𝑥) = 𝑚𝑇1−𝑡 · 𝑥 + 𝜎𝑇2 −𝑡 ∇ log→ 𝑇 −𝑡 and then  − [𝑖 ] − [𝑖 ]  ] [𝑖 ] → −𝑝 [ 𝑘 ] (→ → −𝑝 [𝑖 ] (→ [𝑖 ] 𝜇 [ 𝑘 ] (𝑍 [𝑖 ) − 𝜇 (𝑍 ) = 𝑟 ∇ log 𝑋 ) − ∇ log 𝑗 𝑗 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑋 𝑇 −𝑡 𝑗 )
𝑑 = 𝑟 𝑗 · 𝐹𝑘,𝑖 ( 𝑗) + 𝐺 𝑘,𝑖 ( 𝑗)

(73)

where 𝐹𝑘,𝑖 ( 𝑗) = 𝑚 𝑇 −𝑡 𝑗 · (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 (𝜇 𝑘 − 𝜇𝑖 ),

𝐺 𝑘,𝑖 ( 𝑗) ∼ N 0, 𝐻 𝑘,𝑖 ( 𝑗)

with 𝐻 𝑘,𝑖 ( 𝑗) defined in (40). The moment generating function for the quadratic of Gaussians under the expression (73) gives h  i ] [𝑖 ] 2 [𝑖 ] log E exp 4𝜌 · 𝜇 [ 𝑘 ] (𝑍 [𝑖 ) − 𝜇 (𝑍 ) 𝑗 𝑗

−1 2 1 = − 2 log det 𝐼 𝐷 − 8𝜌𝑟 𝑗 𝐻 𝑘,𝑖 ( 𝑗) + 4𝜌𝑟 2𝑗 (𝐹𝑘,𝑖 ( 𝑗)) ⊤ 𝐼 𝐷 − 8𝜌𝑟 2𝑗 𝐻 𝑘,𝑖 ( 𝑗) 𝐹𝑘,𝑖 ( 𝑗) (74) for any 𝜌 > 0 such that 𝐼 𝐷 − 8𝜌𝑟 2𝑗 𝐻 𝑘,𝑖 ( 𝑗) is positive definite, which is satisfied if 𝜌 < 1/(8𝑟 2𝑗 ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ). Thus, under 𝜌 ≤ 1/(16𝑟 2𝑗 max 𝑘,𝑖 ∈ [𝐾 ] ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ), we have 𝐷 ∑︁


log 1 − 8𝜌𝑟 2𝑗 𝜆ℓ (𝐻 𝑘,𝑖 ( 𝑗)) ≤ 8𝜌𝑟 2𝑗 · tr 𝐻 𝑘,𝑖 ( 𝑗) , − 21 log det 𝐼 𝐷 − 8𝜌𝑟 2𝑗 𝐻 𝑘,𝑖 ( 𝑗) = − 12 ℓ=1

because − log(1 − 𝑥) ≤ 2𝑥 for 𝑥 ∈ [0, 1/2]. We have tr(𝐻 𝑘,𝑖 ( 𝑗)) = 𝑚 𝑇4 −𝑡 tr (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 (Σ 𝑘 − Σ𝑖 ) (𝑆𝑇𝑖 −𝑡 𝑗 ) −1 (Σ 𝑘 − Σ𝑖 ) (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 4 𝑚𝑇 −𝑡

≤ 𝜎6

𝑗

𝑇 −𝑡 𝑗

4 𝑚𝑇 −𝑡

∥Σ 𝑘 − Σ𝑖 ∥ 2𝐹 ≲ 𝜎 6

𝑗

𝑀 † 𝐷 2𝛼

𝑇 −𝑡 𝑗

where we used (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 ⪯ 𝜎 21 𝐼 𝐷 in the first inequality and the following bound in the second: 𝑇 −𝑡

∥Σ 𝑘 − Σ𝑖 ∥ 2𝐹 ≲ ∥Σ 𝑘 ∥ 2𝐹 + ∥Σ𝑖 ∥ 2𝐹 ≤ 2𝐶2 𝐷 𝛼 (𝐶2 𝐷 𝛼 + 𝑅) 𝑀 † ≲ 𝐷 2𝛼 𝑀 † .

(75)

36

Y. Iguchi and P. Fearnhead

where the constants 𝐶2 > 0, 𝑅 ∈ [0, 1) are collected in Assumption 2. In the above, we used

⊤ 𝛼 𝛼 𝛼 𝛼 ∥Σ 𝑘 ∥ 2𝐹 = tr Σ⊤ 𝑘 Σ 𝑘 = tr Λ 𝑘 Λ 𝑘 ≤ 𝜆 1 (Σ 𝑘 ) · tr(Λ 𝑘 ) ≤ 𝐶2 𝐷 (𝐶2 𝐷 𝑀 𝑘 + 𝜈 𝑘 ) ≤ 𝐶2 𝐷 (𝐶2 𝐷 + 𝑅)𝑀 † under Assumption 2 with decomposition Σ 𝑘 = 𝑈Λ 𝑘 𝑈 ⊤ . We thus have that for 𝜌 ≤ 1/(16𝑟 2𝑗 max 𝑘,𝑖 ∈ [𝐾 ] ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ),
− 21 log det 𝐼 𝐷 − 8𝜌𝑟 2𝑗 𝐻 𝑘,𝑖 ( 𝑗) ≤ 𝜌 SNR 𝑗 𝑀 † 𝐷 2𝛼 .

(76)

Similarly, we also have that for 𝜌 ≤ 1/(16𝑟 2𝑗 max 𝑘,𝑖 ∈ [𝐾 ] ∥𝐻 𝑘,𝑖 ( 𝑗) ∥ 2 ), 𝜌𝑟 2𝑗 (𝐹𝑘,𝑖 ( 𝑗)) ⊤ 𝐼 𝐷 − 8𝜌𝑟 2𝑗 𝐻 𝑘,𝑖 ( 𝑗)

−1

𝐹𝑘,𝑖 ( 𝑗) ≤ 2𝜌𝑟 2𝑗 ∥𝐹𝑘,𝑖 ( 𝑗) ∥ 2 ≲ 𝜌𝐷.

(77)

 𝑘
o −1 𝑖 −1 2 𝑖 + tr (𝑆 ) − (𝑆 ) ) 𝑆 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 𝑇 −𝑡 𝑗 , ) −2

(78)

∵ Assumption 3

Substituting (76) and (77) into (74), we conclude Lemma 5.2. E.3. Proof of Lemma 5.3 → − ] Writing 𝑋 𝑇[𝑖−𝑡 as 𝑚 𝑇 −𝑡 𝑗 𝜇𝑖 + 𝑍 with 𝑍 ∼ N (0, 𝑆𝑇𝑖 −𝑡 𝑗 ), we have 𝑗 𝜎𝑇4 −𝑡

𝐷 𝑘[𝑖 ] ( 𝑗) = 𝑚2

𝑗

𝑇 −𝑡 𝑗

i h  2 × E 𝑚 𝑇 −𝑡 𝑗 (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 (𝜇𝑖 − 𝜇 𝑘 ) + (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 − (𝑆𝑇𝑖 −𝑡 𝑗 ) −1 𝑍

n 𝜎𝑇4 −𝑡 = 𝑚2 𝑗 × 𝑚 𝑇2 −𝑡 𝑗 ∥𝜇𝑖 − 𝜇 𝑘 ∥ 2(𝑆 𝑘

𝑇 −𝑡 𝑗

𝑇 −𝑡

where in the first equation, we used: for any (𝑡, 𝑥) ∈ [0, 𝑇) × R𝐷 ,
𝜎2 ] ∇ log 𝑝 𝑇[𝑖−𝑡 (𝑥) − ∇ log 𝑝 𝑇[ 𝑘−𝑡] (𝑥) , 𝜇 [ 𝑘 ] (𝑇 − 𝑡, 𝑥) − 𝜇 [𝑖 ] (𝑇 − 𝑡, 𝑥) = 𝑚𝑇𝑇 −𝑡 −𝑡
with ∇ log 𝑝 𝑇[ 𝑘−𝑡] (𝑥) = −(𝑆𝑇𝑘 −𝑡 ) −1 𝑥 − 𝑚 𝑇 −𝑡 𝜇 𝑘 . The matrix inequality (𝑆𝑇𝑖 −𝑡 𝑗 ) −1 ⪯ 𝜎 21 𝐼 𝐷 together 𝑇 −𝑡

with (58) yields 𝑚 𝑇2 −𝑡 𝑗 ∥𝜇𝑖 − 𝜇 𝑘 ∥ 2(𝑆 𝑘

𝑇 −𝑡 𝑗

2 𝑚𝑇 −𝑡

≤ 𝜎2 ) −2

𝑗

𝑇 −𝑡 𝑗

∥𝜇𝑖 − 𝜇 𝑘 ∥ 2(𝑆 𝑘

𝑇 −𝑡 𝑗

) −1

−𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
. = 𝜎 22 KLm → 𝑇 −𝑡 𝑇 −𝑡

(79)

𝑇 −𝑡 𝑗

By setting 𝑌 := (𝑆𝑇𝑖 −𝑡 𝑗 ) 1/2 (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 (𝑆𝑇𝑖 −𝑡 𝑗 ) 1/2 − 𝐼 𝐷 , we have 

2 tr (𝑆𝑇𝑘 −𝑡 𝑗 ) −1 − (𝑆𝑇𝑖 −𝑡 𝑗 ) −1 𝑆𝑇𝑖 −𝑡 𝑗 = tr (𝑆𝑇𝑖 −𝑡 𝑗 ) −1𝑌 2 ≤ 𝜎 21 tr(𝑌 2 ) 𝑇 −𝑡 𝑗

≤ 𝜎 24

𝑇 −𝑡 𝑗

 𝑚2

𝑇 −𝑡 𝜎𝑇2 −𝑡

 −𝑝 [𝑖 ] || → −𝑝 [ 𝑘 ]
∥Σ𝑖 ∥ 2 + 1 · KLc → 𝑇 −𝑡 𝑇 −𝑡

(80)

where in the last line we used Lemma C.1 provided in Appendix C. Substituting the inequalities (79) and (80) into (78), we conclude.

F. Details on Numerical Experiments This section provides details of the numerical experiments presented in Section 4, and additional information to support the arguments of the section. F.1. Implementation details Data processing For the scRNA-seq experiments on PBMC, we use the PBMC3K cell-by-gene count matrix as the raw data and apply a standard single-cell preprocessing workflow; see e.g., Satija et al. [2015] and Stuart et al. [2019]. First, we remove genes expressed in fewer than three cells and apply total-count normalisation with a target sum of 104 , followed by log-transformation. Then, for each

Diffusion Models for Clustered Data

37

ambient dimension 𝐷 ∈ {500, 1000, 2000, 4000, 8000}, we select the top 𝐷 highly variable genes using the Seurat-flavour procedure implemented in the Python package Scanpy [Wolf et al., 2018]. We retain the CD4 T, CD14+ Monocyte, and CD19+ B populations with the number of samples 1144, 480 and 342 respectively. We then apply gene-wise standardisation, and the resulting 𝐷-dimensional gene-expression vectors are used as the empirical target distribution for the reverse trajectories. For the image experiment, we use the training split of the STL-10 dataset, whose images contain 96 × 96 pixels with 3 RGB channels. We retain three classes: Dog, Cat, and Airplane, and use 500 images from each class. To vary the ambient dimension, we construct three image resolutions 𝑠 ∈ {32, 64, 96}, corresponding to 𝐷 = 3𝑠2 . For 𝑠 = 32 and 𝑠 = 64, the original images are downsampled using bilinear interpolation implemented by torch.nn.functional.interpolate in PyTorch. For 𝑠 = 96, we use the original STL-10 resolution. Finally, pixel intensities are rescaled from [0, 255] to [−1, 1].

Calculation of posterior weights and KL divergence We approximate the true intractable KL divergence between two noised clusters, using a Monte Carlo approximation based on the empirical target distribution. That is, for two distinct clusters 𝐴, 𝐵 ∈ CPBMC or Cimage and 𝑡 ∈ [0, 𝑇), −𝑝 [ 𝐴] || → −𝑝 [ 𝐵]
= KL → 𝑇 −𝑡 𝑇 −𝑡

∫

→ −𝑝 [ 𝐴] ( 𝑥 ) → − [ 𝐴] 𝑇 −𝑡 log → −𝑝 [𝐵] ( 𝑥 ) · 𝑝 𝑇 −𝑡 (𝑥)𝑑𝑥 ≈

→ −𝑝 [ 𝐴], 𝑁 𝐴 ( 𝑥 )

∫

𝑇 −𝑡 log → − [𝐵], 𝑁 𝐵

𝑝 𝑇 −𝑡

𝑇 −𝑡

(𝑥)

−𝑝 [ 𝐴], 𝑁 𝐴 (𝑥)𝑑𝑥, ·→ 𝑇 −𝑡

where 𝑁 𝐴 and 𝑁 𝐵 denote the numbers of empirical samples in classes 𝐴 and 𝐵, respectively, and we define → −𝑝 [𝑖 ], 𝑁𝑖 (𝑥) ≡ 1 𝑁𝑖 𝑇 −𝑡

𝑁𝑖 ∑︁

→ −𝑝 𝑇 −𝑡 |0 (𝑥|𝑍 𝑗 ),

i.i.d.

{𝑍 𝑗 } 𝑗 ∼ 𝑝 0[𝑖 ] ,

𝑖 ∈ { 𝐴, 𝐵},

(81)

𝑗=1

where 𝑝 0[𝑖 ] is the population distribution for the class ‘𝑖’. Then, the estimate of the KL is finally given as: −𝑝 [ 𝐴] || → −𝑝 [ 𝐵]
≈ 1 KL → 𝑁𝐴 𝑀 𝑇 −𝑡 𝑇 −𝑡

𝑁 𝐴 ∑︁ 𝑀 ∑︁

→ −𝑝 [ 𝐴], 𝑁 𝐴 (𝑋 𝑘,𝑙 )

𝑇 −𝑡 log → − [𝐵], 𝑁 𝐵

𝑙=1 𝑘=1 i.i.d.

𝑝 𝑇 −𝑡

(𝑋 𝑘,𝑙 ) i.i.d.

where 𝑋 𝑘,𝑙 = 𝑚 𝑇 −𝑡 𝑍𝑙 + 𝜎𝑇 −𝑡 𝜉 𝑘 , {𝜉 𝑘 } 𝑘=1,..., 𝑀 ∼ 𝑁 (0, 𝐼 𝐷 ), {𝑍 𝑗 } 𝑗=1,..., 𝑁 𝐴 ∼ 𝑝 0[ 𝐴] and 𝑀 is the number of MC samples for the white noise. Since the true posterior weight for general target distributions is intractable, we estimate it using the empirical forward density (81) again as follows: for the set of classes C, b [𝑖 ], 𝑁 (𝑡, 𝑥) = Í 𝑊 𝐷

→ −𝑝 [𝑖 ], 𝑁𝑖 (𝑥) 𝑇 −𝑡 , → −𝑝 [𝑐], 𝑁𝑐 (𝑥)

𝑐∈ C

𝑖 ∈ C,

(𝑡, 𝑥) ∈ [0, 𝑇) × R𝐷 ,

𝑇 −𝑡

where 𝑁 𝑐 denotes the number of samples in the class 𝑐 ∈ C, and we assume equal priors across all the classes.

Reverse SDE with Empirical Score As mentioned in Section 4, we generate reverse trajectories using the empirical score function in the drift. Following the notation above, we write the reverse SDE as: √ ← − ← − − −𝑝 𝑁 ( ← 𝑑 𝑋 𝑡 = ( 𝑋 𝑡 + 2∇ log→ 𝑇 −𝑡 𝑋 𝑡 ))𝑑𝑡 + 2𝑑𝐵𝑡 ,

← − 𝑋 0 ∼ 𝑁 (0, 𝐼 𝐷 ),

38

Y. Iguchi and P. Fearnhead

where 𝑁 is the total number of empirical data, including all classes. In the above, the empirical score is given as ∑︁ −𝑝 𝑁 (𝑥) = −𝑝 [𝑐], 𝑁𝑐 (𝑥), b [𝑐], 𝑁 (𝑡, 𝑥) · ∇ log→ ∇ log→ 𝑊 𝑥 ∈ R𝐷 , 𝑇 −𝑡 𝐷 𝑇 −𝑡 𝑐∈ C

−𝑝 [𝑐], 𝑁𝑐 (𝑥) is the cluster-wise empirical score with the number of samples 𝑁 written as where ∇ log→ 𝑐 𝑇 −𝑡 −𝑝 [𝑐], 𝑁𝑐 (𝑥) = − 1 ∇ log→ 𝑇 −𝑡 𝜎2

𝑇 −𝑡

→ − → − 𝑥 − 𝑚 𝑇 −𝑡 E 𝑁𝑐 [ 𝑋 0[𝑐] | 𝑋 𝑇[𝑐]−𝑡 = 𝑥] ,

with 𝑁

𝑐 ∑︁ → −𝑝 → − → − 𝑇 −𝑡 |0 ( 𝑥 | 𝑍 𝑗 ) E 𝑁𝑐 [ 𝑋 0[𝑐] | 𝑋 𝑇[𝑐]−𝑡 = 𝑥] = 𝑁1𝑐 , 𝑍𝑗 · → − [𝑐], 𝑁𝑐

𝑝 𝑇 −𝑡

𝑗=1

(𝑥)

i.i.d.

{𝑍 𝑗 } 𝑗=1,..., 𝑁𝑐 ∼ 𝑝 0[𝑐] .

F.2. Additional Information We present the empirical pairwise geometry for the PBMC dataset. We observe that the mean separation increases with the ambient dimension 𝐷 approximately linearly, while the maximum mean difference projection 𝑀 𝜇 is relatively stable across varying 𝐷. Table 4. Pairwise empirical geometry of the PBMC clusters under variance-ordered feature truncation. The labels are the same as in Table 3. ∥𝜇 b𝑖 − 𝜇 b𝑗 ∥ 2

𝑀𝜇

∥b Σ𝑖 − b Σ 𝑗 ∥ 2𝐹

CD4 T–CD14+ Monocyte

500 1000 2000 4000 8000

31.55 91.95 240.40 352.70 509.33

1.48 2.67 4.37 4.62 5.68

1.34×103 4.16×103 1.55×104 5.35×104 1.90×105

CD14+ Monocyte–CD19+ B

500 1000 2000 4000 8000

42.77 111.03 296.89 440.36 614.00

1.24 2.39 3.41 4.01 5.29

2.40×103 7.31×103 2.94×104 1.16×105 4.04×105

CD4 T–CD19+ B

500 1000 2000 4000 8000

20.84 53.11 133.63 205.03 271.85

1.96 1.49 1.73 1.60 1.99

1.70×103 5.21×103 2.09×104 8.85×104 3.19×105

Pair

𝐷

Related documents

• Statistical Test for Diffusion-Based Anomaly Localization via Selective Inference#174024
• Statistical Test for Diffusion-Based Anomaly Localization via Selective Inference#173446
• Statistical Test for Diffusion-Based Anomaly Localization via Selective Inference#173196
• Statistical Test for Diffusion-Based Anomaly Localization via Selective Inference#172180
• Unsupervised Brain Anomaly Detection as a Bayesian Inverse Problem with Diffusion Prior#1050447