Source-linked AI summary
Two-Stage Score Estimation for Group-Specific Diffusion Models
Huitao Yang, Oscar Hernan Madrid Padilla
TL;DR
How can diffusion models generate accurately for a target group with limited observations while using structure shared across related distributions? The paper estimates a pooled score and learns a target-specific offset, deriving rate bounds and small-time regularity conditions for that approach. The resulting framework gives conditions for faster upper-bound rates than separate fitting and total-variation guarantees for generated distributions.
Problem
Diffusion transfer for related groups must use shared structure while preserving group differences, especially when target-group observations are limited.
Source: §1 · p. 1Method
The paper estimates a pooled score, freezes that fit, and learns a target-specific offset under a hierarchical composition model.
Source: §1 · pp. 1–2Results
The framework identifies conditions under which a simpler offset yields improved convergence-rate bounds over separate fitting.
Source: §1 · p. 2, §6 · p. 11Takeaways & Limitations
Under density and boundary conditions, small-time offset regularity supports polynomially vanishing cutoffs and total-variation guarantees for diffusion models as target-density estimators.
Source: §1 · p. 2The analysis assumes fixed structural classes and decoupled training stages, and its rate comparisons rely on upper bounds whose corresponding lower bounds remain open.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Generative modeling across related distributions raises a question: how can shared structure support accurate generation for a target distribution with limited observations? We develop a statistical framework for diffusion transfer that analyzes a shared pooled score together with a target-specific offset. Under Besov regularity of the pooled density and a stable hierarchical composition model for the offset, we derive convergence rates that separate the costs of learning these components and identify when a simpler offset yields a faster risk-bound rate than separate fitting. Under suitable boundary and regularity assumptions, we prove that the offset remains bounded and jointly regular as diffusion time approaches zero. These results allow the lower time cutoff to vanish polynomially and provide total-variation guarantees for diffusion models as density estimators. Synthetic simulations and image benchmarks provide a proof of concept for the framework.
1 Introduction
The paper studies diffusion transfer across related groups by learning a pooled score and a target-specific score offset. Its framework separates their estimation costs and analyzes when the offset’s structure supports faster rate bounds and stable generation as diffusion time approaches zero.
1 Introduction
- The target score is estimated as a pooled score plus a group-specific offset, with the pooled fit frozen before learning the offset from an independent target sample.The offset is the gradient of the log density ratio between the group and pooled diffused densities.Source: §1 · p. 1
- The conditional oracle inequality separates inherited pooled-score error from the offset’s approximation and estimation costs.The target-stage complexity depends only on the trainable offset class.Source: §1 · pp. 1–2
- Under pooled-density Besov regularity and a stable hierarchical offset model, the rate bounds identify when a simpler offset yields a faster upper-bound rate than separate fitting.The pooled-density smoothness and offset-constituent smoothness and arity determine the rate components.Source: §1 · p. 2
- Under density and boundary conditions, the offset remains bounded and jointly regular as diffusion time approaches zero, supporting polynomially vanishing cutoffs and total-variation guarantees.The guarantees concern diffusion models as estimators of the target data distribution.Source: §1 · p. 2
2 Related work
The paper positions its analysis at the intersection of nonparametric transfer learning and diffusion theory. It studies pooled-plus-offset score transfer while accounting for vanishing noise, offset complexity, and recovery of the target distribution.
Nonparametric transfer learning
- The paper adapts pooled estimation plus group-specific hierarchical offsets, studied for grouped regression by Ren et al. [2026], to diffusion scores and distributional recovery.The analysis accounts for vanishing noise.Source: §2 · p. 2
Transfer and guidance in diffusion models
- Unlike transfer approaches based on shared latent structure, conditioning, or density-ratio guidance, the analysis tracks pooled-fit error and compositional offset complexity in a joint bound.It quantifies when a simpler offset improves the bound over separate fitting.Source: §2 · p. 2
Statistical theory for diffusion models
- The pooled-stage rate draws on joint space-time Besov approximation, while hierarchical ReLU approximation supplies constituent-based offset rates.The compositional representation remains an assumption, and small-time cancellation supports regularity.Source: §2 · p. 2
3 Methodology
The method first learns an unconditional pooled score from observations across groups, then freezes that fit and learns a target-group offset. The training stages use pooled and independent target-group samples, and denoising losses target integrated score risk.
3.1 Grouped data and diffusion scores
- The pooled density mixes the group densities according to their prevalences, and Stage 1 uses pooled observations from groups without requiring labels for the pooled score fit.Stage 2 uses an independent sample from the target group; sample splitting can supply the independent samples from one grouped dataset.Source: §3.1 · p. 3, §3 · p. 3
3.2 Two-stage offset estimation
- The empirical denoising loss averages over forward-corrupted observations, and Vincent’s identity makes its population minimizer target integrated score risk.The pooled fit uses the corresponding loss with unconditional pooled observations.Source: §3.2 · p. 3
- Stage 1 estimates the pooled score, then freezes it while Stage 2 estimates the target offset from group-specific data.The target score is represented as the pooled score plus the group-specific offset.Source: §3.2 · pp. 3–4, §3 · p. 3
4 Theory
The theory separates pooled-score learning from target-specific offset learning, then derives conditions for faster transfer rates and stable offsets at small diffusion times. It also translates score accuracy into distributional accuracy for generated samples.
4.1 Error decomposition for two-stage estimation
- Theorem 1 separates inherited pooled-score error from offset approximation and estimation costs, while polynomial early stopping contributes only logarithmically to the empirical-loss scale.The pooled-score bound transfers to the target distribution under the stated overlap condition.Source: §4.1 · pp. 4–5
4.2 Convergence rates and gains from pooling
- Under pooled Besov regularity and a stable hierarchical offset model, Theorem 2 separates pooled and offset score-risk rates and identifies when two-stage fitting has a faster polynomial rate than separate fitting.The rate comparison depends on pooled smoothness and the offset constituents’ smoothness-to-arity ratios.Source: §4.2 · pp. 5–6
4.3 Small-time regularity of the score offset
- Under the stated density, ratio, collar, and schedule conditions, the score offset remains bounded and jointly regular as diffusion time approaches zero.The resulting offset admits a depth-one global-Lipschitz representation after polylogarithmic normalization.Source: §4.3 · pp. 7–8
4.4 Accuracy of the generated distribution
- Corollary 1 converts the score-estimation rate into a total-variation bound for the generated distribution under reverse-SDE localization and bounded variation of the zero-extended target density.The proof controls reverse-path relative entropy using integrated score error, then applies Pinsker’s inequality and bounds endpoint errors.Source: §4.4 · pp. 8–9
5 Experiments
The experiments test pooled-score transfer with target-specific offsets on synthetic shifts, a rough-density construction, and image benchmarks. Offset performs best on the reported MNIST/Fashion-MNIST configurations and selected larger-target CIFAR-10 conditions, while results also show limits when shifts are absent or target samples are small.
5 Experiments
- The experiments use shared data splits, diffusion schedules, and evaluation draws; target-stage update counts match for Separate, Top FT, and Offset.Pooled, Separate, and the narrower Offset use the same backbone family.Source: §5 · p. 9
5.1 Shared modes with shifted weights
- In the eight-mode simulation, increasing target sample size broadens Offset's advantage over Pooled, but Top FT is best at nℓ=20 and Pooled beats Offset in the no-shift control.At nℓ=40 in the no-shift control, Pooled's exact-score MSE is 0.066 versus 0.087 for Offset.Source: §5.1 · p. 9, Table 1 · p. 9
5.2 A rough density with a stable offset
- Offset outperforms Separate in all five seeds at each of three target sample sizes, while effective-offset error falls from 0.0257 to 0.0055 as nℓ grows from 64 to 1,024.Pooled-fit error remains fixed; finite Adam runs illustrate the construction, whereas the theorem concerns empirical minimizers over specified network classes.Source: §5.2 · p. 9
5.3 Image benchmarks
- Offset has the lowest mean DSM in all eight MNIST/Fashion-MNIST configurations and at nℓ∈{1024, 2048, 4096} for the selected CIFAR-10 conditions.Table 3's six larger-target conditions were selected after inspecting initial results.Source: Table 2 · p. 10, §5.3 · p. 10, Table 3 · p. 10
6 Conclusion
The conclusion presents a framework for separating pooled-score error from target-offset learning costs, with conditions for faster upper-bound rates than separate fitting. The section’s supplied material also covers notation, simulation and benchmark setup, and a TGDP adaptation.
6 Conclusion
- The framework separates pooled-score error from offset-learning costs and identifies conditions for faster convergence-rate bounds than separate fitting.The rate comparisons are upper bounds; corresponding lower bounds remain open, and the theory assumes fixed structural classes and decoupled stages.Source: §6 · pp. 10–11
- Practical sampling introduces discretization error, and the finite-budget empirical evaluations focus on score losses rather than systematic sample-based evaluation.The paper also leaves optimization analyses for future work.Source: §6 · p. 11
A Notation and function classes
- The section defines Besov smoothness notation and states that restriction smoothness can permit a jump in the zero extension.It also defines TV and KL distances and bounded variation.Source: A Notation and function classes · pp. 14–15
B.1 Eight-mode Gaussian mixture
- At nℓ=40, the no-shift control has mean exact-score MSE 0.066 for Pooled and 0.087 for Offset.The control sets all mixture weights to 1/8.Source: §B.1 · p. 15
B.2 MNIST and Fashion-MNIST
- The MNIST and Fashion-MNIST evaluation uses contrast-transformed images or class-specific targets, with pools containing 6,000 images and 2% targets.Nested target sets of 64 and 256 use base-image identities disjoint from the pool.Source: §B.2 · p. 15
B.3 TGDP adaptation
- The TGDP adaptation estimates target-to-pooled density ratios with clean and time-dependent classifiers, then adds a guidance-potential gradient to the frozen pooled score.The potential is fit using a scalar network; the resulting score is the pooled score plus the gradient of log hψ.Source: §B.3 · pp. 15–16
B.4 CIFAR-10 data, training, and evaluation
The supplied results report a rough-density comparison favoring Offset over Separate, while the CIFAR-10 passages specify the independent evaluation protocol. The section also describes model comparisons, loss normalization, and the rough-density simulation setup.
B.4 CIFAR-10 data, training, and evaluation
- CIFAR-10 uses contrast-transformed or truck targets, nested adaptation sets disjoint from the pool, and 1,000 official-test images per target.The fixed-grid comparison used five additional runs with Pooled, Separate, Offset, and TGDP.Source: §B.4 · p. 16
B.5 Architectures, optimization, and loss normalization
- Separate fits target data, Top FT updates Pooled’s output layer, and Offset and TGDP freeze Pooled while fitting an offset or guidance potential.Paired DSM differences compare score accuracy; the image benchmarks use direct score outputs.Source: §B.5 · pp. 16–17
B.6 Rough-density simulation
- The rough-density simulation evaluates ordinary-time integrated squared score error, using five runs with nested target samples independent of the pooled sample.Evaluation uses 256 independent target points and shared Gaussian noise; the time grid and quadrature settings are reported separately.Source: §B.6 · p. 17
- Offset beat Separate in all 15 paired fits, while all intervals against Pooled included zero; effective-offset errors were 0.025729, 0.008807, and 0.005504 at nℓ=64, 256, and 1,024.The table reports integrated target-score risk; each pooled fit is reused across sizes.Source: §B.6 · pp. 17–18, Table 8 · p. 18
C Proof of the two-stage oracle inequality
The proof strategy first controls clipped empirical loss and establishes an oracle inequality, then inserts a pooled-plus-offset comparator. The supplied boundary passage states explicit support, boundary, and density conditions for the boundary analysis.
C Proof of the two-stage oracle inequality
- The proof controls clipped loss and tails, derives an empirical-risk oracle inequality, and inserts the pooled-plus-offset comparator while conditioning on the pooled fit.The loss-difference identity equates differences in D with differences in population risk.Source: C Proof of the two-stage oracle… · p. 18
C.2 A finite-net oracle argument
- The finite-net argument gives a high-probability clipped empirical-risk oracle bound conditional on the pooled fit and target sample size.The comparator step transfers the first-stage fit and offset into the target-stage bound.Source: §C.2 · pp. 19, 21
C.1 Integrated envelopes, tails, and clipping
- The clipping argument uses Gaussian-tail bounds, and the proof tools also apply to any fixed clipped score class, including the pooled class.The cited extension requires clipping, an entropy bound, and i.i.d. clean observations.Source: §C.1 · p. 19, §C.2 · p. 20
D Details of boundary cancellation
- The boundary analysis assumes compact support with a C1 boundary near the point considered and a positive limiting density there.These are stated as conditions for the posterior boundary-growth analysis.Source: D Details of boundary cancellation · p. 21
D.1 Gaussian localization and quotient differentiation
The supplied proofs use Gaussian localization and posterior identities to control derivatives near boundaries, then establish offset regularity and network approximation bounds. The section also gives a numerical construction check and fitted-class complexity results.
D.1 Gaussian localization and quotient differentiation
- Gaussian convolution derivatives are polynomial multiples of the kernel, while contributions supported a fixed distance away decay exponentially in inverse variance.The localization bound applies to each fixed mixed derivative.Source: §D.1 · p. 22
D.2 Posterior identities and the boundary limit
- At a smooth boundary with positive limiting density, rescaled posterior limits are the Gaussian mass and first moment on the inward half-space.The proof uses local boundary rescaling and Gaussian kernel bounds away from the boundary neighborhood.Source: §D.2 · p. 23
D.3 Cancellation in the boundary and exterior zones
- A density ratio constant on a boundary collar keeps G bounded, while positive-order derivatives of log G decay exponentially in boundary and exterior zones.The decay applies to fixed spatial and variance derivatives under the stated conditions.Source: §D.3 · p. 23
D.4 Construction of Figure 1
- Figure 1 samples spatial and cumulative time maxima of the offset and its derivatives; doubling quadrature order changed fields by less than 10^-8 of their maximum magnitude.Doubling spatial resolution changed sampled maxima by less than 1.2% of the respective field maxima.Source: §D.4 · p. 23
E Construction of the stable offset network
- The fitted offset class is constructed with input and output clamps, and its network budgets are chosen to meet the approximation and entropy bounds.The class can be set to {0} when the offset is known to vanish.Source: E Construction of the stable offset… · p. 23, §E.2 · pp. 24–25
E.1 Imported approximation bounds
- For smooth hierarchical constituents, approximation rates depend on smoothness and arity; fixed compositions preserve these rates under Lipschitz error propagation.The cited construction has depth O(log M) and width O(M^K) for constituent arity K.Source: §E.1 · p. 24
E.2 Fitted network classes and coefficient bounds
- The fitted class’s covering-number bound scales with sparse parameter count and coefficient range, and output clipping preserves that bound.The offset approximation and complexity proposition applies under compact initial support and Assumption 2.Source: §E.2 · p. 25
F Pooled approximation and combination of the stages
The section combines a pooled score fit with a target-specific offset, then establishes conditions under which this two-stage approach can improve the rate over separate fitting. It also develops uniform small-time regularity bounds for the offset.
F Pooled approximation and combination of the stages
- The pooled estimator is fitted over a ReLU sieve, with its approximation and estimation costs combined under Assumption 3.The proof imports Besov score approximation and oracle, clipping, and entropy bounds.Source: F Pooled approximation and combination… · pp. 25–26
- Positivity of the OU transition kernel transfers pooled-density overlap to the diffused densities, yielding the group-weighted squared-error bound.A prevalence-uniform overlap constant imposes stronger overlap for rare groups.Source: F Pooled approximation and combination… · pp. 26–27
- If the offset vanishes, the fitted pooled score is already clipped, so the pooled bound alone applies.Source: F Pooled approximation and combination… · p. 27
- When nℓ ≍ n^ρ with 0 < ρ < 1 and b > a, both two-stage terms decay faster than the separate-fitting bound.The example attributes the gain to learning the shared component from n observations and using the target sample for the simpler offset.Source: F Pooled approximation and combination… · p. 27
G.1 Uniform boundedness
- The uniform-boundedness argument uses a variance inequality for a Gaussian conditioned on a convex cube, followed by bounded reweighting.The cited inequality bounds variance by v times the squared Lipschitz constant.Source: §G.1 · p. 27
G.2 Joint regularity under a smooth collar-flat ratio
- The offset has uniform parabolic Hölder bounds over T ≥ n^-κ0, and its extended space-time norm is polylogarithmic in n.For smoothness s > 1, the bound also controls the global Lipschitz constant.Source: §G.2 · pp. 28–29
H A rough density with a stable smooth offset
A constructed example separates a rough pooled-density cusp from a smooth group-specific tilt. The example verifies offset regularity and gives a rate-bound comparison in which the smoother offset, rather than lower dimension, accounts for the gain.
H A rough density with a stable smooth offset
- For the densities in (58), with d=2, s0=1/2, p*0=θ*0=∞, and κ̄=3/2, the pooled density is not locally in B^s_{∞,∞} at the cusp for any s>1/2; both group offsets satisfy Assumption 2 with depth one.Source: H A rough density with a stable smooth… · p. 30
- All space-time derivatives of the offset through total order two are bounded at small times, using support separation and localized heat estimates.Source: H A rough density with a stable smooth… · p. 30
- The bounds extend to remaining times and yield globally bounded, Lipschitz C1,1 constituents after normalization, uniformly for |c_g| ≤ 1/2.Source: H A rough density with a stable smooth… · p. 30
- n^-1/3 + m^-4/7 versus m^-1/3 gives the two-stage upper-bound terms and separate Besov-estimation bound, respectively, up to logarithms.Here m = nℓ ≍ n^3/4, and the gain comes from a smoother offset, not reduced spatial dimension.Source: H A rough density with a stable smooth… · p. 30
I Proofs for total-variation recovery
The total-variation argument decomposes recovery error into early-, middle-, and terminal-time contributions. Each regime is bounded under its corresponding regularity, localization, or entropy conditions.
H A rough density with a stable smooth offset
- The early-time contribution is at most C√T for compactly supported densities with uniformly bounded zero-extension BV norm, while the terminal contribution is at most Ce^-βT under a uniform Gaussian KL bound.The BV condition permits jumps at the support boundary.Source: I Proofs for total-variation recovery · p. 31
- A stopped-drift coupling and endpoint data processing establish the middle-time total-variation bound for the learned reverse process.The argument assumes measurable, globally Lipschitz reverse drifts with linear growth on the learning interval.Source: I Proofs for total-variation recovery · pp. 31–32
- The three bounds combine into the stated total-variation decomposition between the target density and the learned endpoint law.Source: I Proofs for total-variation recovery · p. 32
J Proof of the balanced TV rate
The section establishes total-variation recovery by combining the two-stage score bound with early-, middle-, and terminal-time controls. It also gives space-time composition and Besov conditions for the two fitted stages, noting when the simpler-offset gain is unavailable.
J Proof of the balanced TV rate
- Together with the stated BV and localization conditions, this verifies Proposition 8’s hypotheses; applying Theorem 2 with δ = n−b/2 yields, with probability at least 1 −n−b, TV(P0,ℓ, bPℓ) ≤ CεTV.Source: J Proof of the balanced TV rate · p. 33
K.1 Hierarchical composition in the joint space–time input
- Joint space-time hierarchical composition for both stages gives a high-probability rate bound with an inverse-polylogarithmic lower cutoff and nℓ ≳ log n.The theorem assumes uniform overlap and independent training samples.Source: §K.1 · p. 33
K.2 Besov regularity at both stages
- With Besov regularity at both stages, the target stage pays for two full-score approximations; the simpler-offset gain requires additional structure.Source: §K.2 · pp. 33–34
Related papers
This paper cites
- Score-Based Generative Modeling through Stochastic Differential Equations2020
- Classifier-Free Diffusion Guidance2022
- Diffusion Models Beat GANs on Image Synthesis2021
- Deep Unsupervised Learning using Nonequilibrium Thermodynamics2015
- Generative Modeling by Estimating Gradients of the Data Distribution2019
- Elucidating the Design Space of Diffusion-Based Generative Models2022