PaperlayerView exact revision on arXiv

Source-linked AI summary

Weighting Schedules Govern What and When Score-Based Generative Models Learn from Multimodal Data

Jérémie Klinger, Raphaël Urfin, Giulio Biroli, Marylou Gabrié

arXiv:2609.35322v1cs.LGcond-mat.dis-nn

TL;DR

The paper asks how training choices determine which features of multimodal data a score model learns and when. It analyzes feature learning across noise levels in high-dimensional Gaussian mixtures and tests the predictions in more complex data settings. Learning depends on effective weighting around speciation: far from it, directions can be learned without relative weights, while near it, weights and directions emerge jointly.

  • Problem

    Prior analyses locate mode commitment during generation but do not explain how a trained score acquires multimodal features or whether speciation matters for training.

    Source: §1 · p. 1, §2 · p. 4

  • Method

    The paper analyzes training dynamics at fixed signal-to-noise ratios in unbalanced and hierarchical Gaussian mixtures, then examines effective weighting in standard stochastic-interpolant pipelines and empirical data settings.

    Source: §1 · pp. 1–2

  • Results

    At high signal-to-noise ratios, mode directions are acquired together while relative weights remain unlearned; near speciation, weights and directions are learned jointly, with direction rates set by relative amplitudes.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Effective weighting near the speciation scale predicts which features are learned and when, and could be tailored to target structure or a desired compute budget.

    Source: §5 · p. 9

    Tailoring weighting schedules to acquire chosen features or meet a fixed compute budget is left for future work.

    Source: §5 · p. 9

Abstract

from arXiv · show

Score-based generative models generate new samples by integrating a time-dependent drift that carries Gaussian noise onto the target distribution. In practice this drift is modeled by a neural network, trained on a loss integrated over time $t$ with a weighting schedule $w(t)$. Along the backward dynamics, and for multi-modal distributions, trajectories commit to modes of the target within a narrow time window, the \textit{speciation time}. In this work, focusing on high-dimensional data, we decompose the integrated loss into its single-time contributions and analyze each at fixed signal-to-noise ratio $\Lambda(t)$: we show that $\Lambda(t)$ sets the rate at which each feature of a multimodal target - the mode directions and their relative weights - is acquired during training. Crucially, at high $\Lambda(t)$ all mode directions are acquired together, on a single timescale insensitive to their amplitudes, while the relative weights are not learned at all. Only near the speciation time, where $\Lambda(t)$ becomes of order one, do all features become learnable, each on its own timescale: the weights are acquired jointly with the directions, and the directions at rates set by their relative amplitudes. For models trained on time-integrated objectives, the learning dynamics is then governed by how much of the weighting effectively sits near the speciation time, which provides insights on $w(t)$ design choices. These results follow from an exact high-dimensional analysis of the training dynamics of unbalanced and hierarchical Gaussian mixtures. Numerical experiments on image and human genome haplotype generation recover the predicted hierarchy of learning timescales in more complex settings.

1 Introduction

Prior analyses locate when generation commits to a mode, but do not explain which target features a trained score learns or when. This work analyzes how noise levels and weighting schedules shape feature acquisition in high-dimensional multimodal data.

1 Introduction

  • The training objective quantitatively sets the rates at which mode directions, relative weights, and mode amplitudes are acquired.The paper studies unbalanced and hierarchical Gaussian mixture models and empirically tests whether the findings generalize to standard stochastic-interpolant pipelines.Source: §1 · p. 1
  • The analysis targets finite-time, finite-precision training, when noise and weighting schedules influence learning dynamics despite exact minimization recovering the score at every noise level.The schedule’s influence is the paper’s concern for high-dimensional data.Source: §1 · p. 2
  • At fixed noise levels, training well before speciation recovers mode directions but not relative weights, whereas training near speciation learns weights and directions jointly.Near speciation, direction-learning rates depend on the modes’ relative amplitudes; at earlier times, direction-learning timescales are insensitive to mode amplitudes.Source: §1 · p. 2
  • The effective weighting over signal-to-noise ratios determines which noise regions dominate learning, and Diffusion and Flow Matching induce different effective weights.These differences yield distinct learning dynamics.Source: §1 · p. 2
  • Experiments on unbalanced and hierarchical MNIST, human-genome haplotypes, and multiple stochastic-interpolant frameworks recover quantitative agreement with predicted feature-emergence timescales.The evaluated frameworks include Flow Matching and Masked Discrete Diffusion.Source: §1 · p. 2

Related works

Prior work analyzes when multimodal structure appears during generation and how training dynamics learn distributional structure, but the training-time question for multimodal targets remains distinct. Existing work also studies noise schedules and weighting schemes, while this paper focuses on their consequences for learning mode weights and hierarchy.

Where structure appears at generation time

  • Prior sampling analyses identify a narrow speciation window in which generation trajectories commit to target modes, assuming an exact score [32] [10] [11] [27] [37].The complementary question is when each feature becomes available during training.Source: Related works · p. 2

Where structure appears at training time

  • Training-dynamics studies address sequential learning in hierarchies, moments, or Gaussian covariance structure, whereas this work focuses on multimodal structure [7] [14] [30] [8] [41].Other cited work studies finite-sample optimal scores or mode-weight recovery without analyzing training dynamics [13] [5].Source: Related works · p. 2

Noising and weighting schedules

  • Prior work proposes SNR-based schedules and analyzes schedule optimization, mode-weight recovery, or weighting effects on performance; this paper studies mode weights and hierarchy in standard SI pipelines [26] [25] [21] [16].The related work also includes ad-hoc weighting analyses [5] and studies of weighting schemes [18].Source: Related works · p. 3

2 Setting

The section sets up stochastic interpolants, whose marginals bridge a target distribution and Gaussian noise, and the score-matching training framework with a weighting function. It motivates studying how weighting and schedule choices shape learning in high-dimensional multimodal targets, against the background of sampling-time speciation.

2 Setting

  • Stochastic interpolants define a bridge from target distribution P0 to Gaussian white noise, with marginals Pt and a deterministic flow whose backward integration generates samples from P0.The score s*(x, t) = ∇x log Pt(x) minimizes the Denoising Score Matching loss, and the score is approximated by a neural network trained with an arbitrary weighting function w(t).Source: §2 · p. 3
  • Schedule functions and weighting determine the relative emphasis on noise levels and thus which target features can be recovered, but their interplay with training and sampling dynamics remains poorly understood.Available results cited here assume an exact, optimally trained score.Source: §2 · p. 3
  • The analysis uses an unbalanced bimodal mixture and a quadrimodal hierarchical distribution, in the high-dimensional limit d ≫ 1, where the choice of w(t) matters most.For the unbalanced mixture, P0(x) = ω N(μ, σ²I_d) + (1 − ω) N(−μ, σ²I_d), with ||μ||² = d and ω, σ² = O_d(1); the hierarchical distribution has component means supported on separate coordinate blocks.Source: §2 · pp. 3–4
  • For a symmetric bimodal Gaussian target, the speciation time t_s = 1/2 log(d) marks a narrow transition from an effectively unimodal potential to two separated wells where trajectories commit to a mode.This sampling result identifies where sampling compute is best spent, but does not establish how training approaches the exact score.Source: §2 · p. 4
  • The paper studies whether speciation time matters for training and whether training-time learning extends to unequal mode weights and hierarchies of sub-modes at different scales.The stated aim is to study training dynamics under such structure for general stochastic interpolants.Source: §2 · p. 4

3 Analytical results

The analysis uses fixed-SNR training dynamics for unbalanced and hierarchical Gaussian mixtures to show that the SNR determines both which features are learned and how quickly. For integrated objectives, the effective weighting near speciation shapes the resulting learning signatures.

3 Analytical results

  • The analysis studies gradient-flow training of score models on unbalanced and hierarchical GMMs, reducing the dynamics to closed systems of feature-summary statistics.Single-noise-level training isolates the dependence on SNR before considering integrated objectives.Source: §3 · p. 4, §3.1 · p. 5

3.1 Single fixed noising time

  • At high SNR, the unbalanced model learns the mode direction on O(d) timescales but leaves the mode weight unchanged; at speciation, direction and weight are learned jointly.The skip connection adapts first on O(1) timescales in both regimes.Source: §3.1 · p. 5, Figure 1 · p. 6
  • In Regime II, both block overlaps evolve jointly, independently of block asymmetry; at speciation, κ sets each block’s SNR, so dominant direction µ2 is acquired before µ1.In the small-κ limit, τ1/τ2 ∼ 1/κ.Source: §3.1 · pp. 5–6
  • The closed-form hierarchical rate λ(γ) uses a fixed-norm heuristic that is not exact away from κ → 0, where it recovers the dynamics exactly.The rate supplies the reference timescale used to rescale the training curves.Source: §3.1 · p. 6

3.2 Integrated loss

  • Rewriting the integrated loss in SNR yields an effective weighting: diffusion spreads effort across SNR decades, while canonical Flow Matching tilts it toward low SNR near speciation.Because parameters are shared across noise levels, the fixed-SNR dynamics do not simply superpose; the integrated experiments show qualitative agreement.Source: §3.2 · p. 7
  • In ResNet experiments, uniform DSM first approaches an even mode split before converging to ω, while FM learns weight and direction together and more sharply separates the subdominant direction as κ varies.Under FM, rescaling training steps by λ(√κ) collapses the subdominant-direction curves.Source: §3.2 · p. 7, Figure 3 · p. 7

4 Structure learning in complex data

Experiments on MNIST and human haplotype data test whether the predicted feature-learning patterns extend beyond controlled mixtures. They find joint learning of class proportion and direction under Flow Matching, amplitude-dependent color learning, and sequential recovery of leading genomic features.

Weight learning in unbalanced MNIST

  • For MNIST restricted to 3s and 6s, Flow Matching learns class proportion jointly with direction, while the DSM comparison is inconclusive because the data offers too short an SNR range to separate the weightings.The out-of-the-box pipelines also differ in more than their weighting function.Source: §4 · p. 8, Figure 4 · p. 8

Learning hierarchy in Tinted MNIST/FM

  • On tinted MNIST, the digit direction is learned at a rate nearly insensitive to tint intensity α, while color learning varies with α and its onset curves collapse after rate rescaling.The GMM rescaling ansatz quantitatively captures color-mode onset when the small mode amplitude is identified as proportional to α.Source: §4 · p. 8, Figure 4 · p. 8

Sequentially learning structured features in genomic data (HGD/MDLM)

  • In HGD under Masked Discrete Diffusion, generated samples acquire the leading PCA directions sequentially; rate rescaling collapses the two leading curves but degrades for smaller directions without resolved modes.The leading two directions are separated by close to half a decade of training.Source: §4 · p. 8, Figure 5 · p. 9
  • The scale-free Masked Discrete Diffusion loss still permits sequential feature emergence in HGD.The mode-amplitude ansatz used for GMMs has no counterpart in this setting.Source: §4 · pp. 8–9

5 Conclusion and future work

Effective SNR weighting predicts which features a model learns and when: high-SNR training recovers mode directions but not relative weights, while training around speciation acquires both, with direction-learning rates set by relative amplitudes. The authors propose tailoring weighting schedules to target structure and compute budgets as future work.

5 Conclusion and future work

  • High-SNR training recovers all mode directions together but leaves relative weights untouched, whereas around speciation weights and directions are acquired jointly and direction rates reflect relative amplitudes.This phenomenology also appears in experiments on unbalanced MNIST, tinted MNIST, and genomic haplotype data.Source: §5 · p. 9
  • The authors leave target-tailored weighting schedules for acquiring chosen features within a fixed compute budget to future work.They frame schedule tailoring as a design principle rather than a completed result.Source: §5 · p. 9

Supplementary Material (SM)

The supplementary analysis identifies speciation through a transition in signal-to-noise ratio and derives unbalanced-mixture training dynamics from an exact loss and gradient flow. Its high-SNR and speciation-time regimes differ in whether mode imbalance is learned.

A.1 Dynamical regimes and signal-to-noise ratio

  • For Λt ≪ 1 trajectories are not driven toward either mode; for Λt ≫ 1, a linear cusp makes q = 0 unstable and commits trajectories to a branch; transition at Λt = O(1) defines speciation time ts.The adopted convention is Λts = 16.Source: §A.1 · p. 14

A.2.1 Unbalanced distribution: exact loss and gradient flow

  • The unbalanced-mixture analysis derives the DSM loss and gradient flow in finite summary statistics for any SNR, then specializes by scaling analysis.The exact reduction is followed by separate high-SNR and speciation-time regimes.Source: §A.2.1 · pp. 14, 17
  • For the unbalanced mixture at high SNR, the network learns mode direction on an O(d) timescale but leaves the imbalance bias frozen at initialization.The exact score is recovered up to that bias.Source: §A.2.1 · p. 15
  • At speciation, mode direction and imbalance are learned jointly on the same O(d) training timescale.The flow converges to a bias encoding the mixture imbalance.Source: §A.2.1 · pp. 15, 19

A.2.2 GMM with four modes: exact loss and gradient flow

The section analyzes how SNR and time weighting shape learning in Gaussian-mixture settings, including block-specific effects and fixed-norm speciation-time dynamics. It derives weighting behavior for flow matching and masked discrete diffusion, and describes fixed-noise DSM experiments on synthetic mixtures to assess learned score coefficients.

A.2.2 GMM with four modes: exact loss and gradient flow

  • At speciation, the blocks decouple and learn at rates proportional to κ_i; for κ ≪ 1, the subdominant direction takes O(d/κ) rather than O(d).The orthogonal component q_i decreases on the same timescale as overlap m_i, driving its growth rate upward as it decays.Source: §A.2.2 · pp. 20, 22
  • At high SNR, both mode directions are recovered together on the same O(d) timescale, regardless of block asymmetry.The exact loss decomposes into two single-mode problems at rescaled SNR κ_iΛ, coupled only through the shared skip connection.Source: §A.2.2 · pp. 21–22

A.2.3 Fixed-norm ansatz: a one-dimensional reduction at the speciation time

  • At speciation, constraining the weight vector to the target direction’s fixed norm removes q and reduces learning to a scalar overlap ODE with a positive, fixed growth rate λ(γ).For small γ, this rate agrees with the limiting rate from the unconstrained dynamics, but the descriptions do not coincide away from that limit.Source: §A.2.3 · pp. 23–24

A.3 Integrated loss: SNR reparametrization and effective weightings

  • Two pipelines with the same effective SNR weighting train identically regardless of time parametrization, while the same time weighting can produce different effective weightings for different interpolants.The integrated loss can be reparametrized by SNR, allowing single-time results to be applied directly.Source: §A.3 · p. 24
  • In log-SNR, Diffusion assigns equal weight per SNR decade, whereas Flow Matching increasingly up-weights low SNR, with its weighting per decade growing as SNR decreases.Flow Matching’s denoising-loss reweighting includes a 1/α²(t) factor, concentrating weight toward the noisy, low-SNR region.Source: §A.3 · p. 25

A.4 Matching Diffusion and Flow Matching at equal SNR

  • Uniform time sampling with the derived weight reproduces Flow Matching’s effective weighting, with large- and small-time limits as anticipated.The small-time divergence is regularized in practice by a training cutoff tmin > 0.Source: §A.4 · p. 26

A.5 Masked discrete diffusion

  • MDLM’s effective weighting is noise-schedule independent and has the same 1/Λ low-SNR behavior as diffusion, making it scale-free and more DSM-like than FM.Only the exposed Λ range depends on the noise-schedule endpoints.Source: §A.5 · p. 26

B.1 Fig. 1 and Fig. 2

  • The experiments use fixed-noise DSM to learn score coefficients for unbalanced two-mode and quadrimodal Gaussian mixtures, evaluating coefficients rather than generated samples.The unbalanced mixture sweeps dimension at ω = 0.8; the quadrimodal mixture fixes d = 1024 and sweeps κ, with comparisons at t = 0.5 and ts = 1/2 log d.Source: §B.1 · pp. 26–27

B.1.1 Data

The appendix specifies the Gaussian-mixture datasets, score models, and training procedures, then describes experiments using fixed-time training and continuous-time ResNet objectives.

B.1.1 Data

  • The unbalanced and quadrimodal GMMs vary dimension and block asymmetry, respectively, with specified sign probabilities and orthogonal block-indicator directions.The unbalanced sweep uses d ∈ {64, 128, 256} at ω = 0.8; the quadrimodal sweep uses d = 1024 and κ ∈ {0.1, 0.2, 0.3}.Source: §B.1.1 · p. 27

B.1.2 Model architecture

  • For the unbalanced mixture, the exact score is parameterized by a learned direction w, scalar c, and bias b, totaling d + 2 parameters.At convergence, b represents the mode log-weight through b → 1/2 log(ω/(1−ω)).Source: §B.1.2 · p. 27
  • For the quadrimodal mixture, the exact score uses one learned direction and bias per orthogonal block, plus c, totaling d + 3 parameters.The block directions are supported on their respective mode subspaces.Source: §B.1.2 · p. 27

B.1.3 Training procedure

  • Fixed-time DSM training uses x_t = α_t x_0 + σ_t ξ, with α_t = e^−t and σ_t = 1 − e^−2t.The expectation is over training samples and standard Gaussian noise.Source: §B.1.3 · p. 27

B.1.4 Unbalanced, dimension dependence, Fig. 1

  • Figure 1 compares overlap, precision-like c, and inferred mode weight at t = 0.5 and at speciation time, with SGD steps scaled by dimension.The plotted dimensions are 64, 128, and 256; the target mode weight is ω = 0.8.Source: §B.1.4 · p. 27

B.1.5 Kappa dependence, Fig. 2

  • Figure 2 plots normalized overlaps for two orthogonal directions across κ values, with the final panel rescaling steps by each block’s growth rate.The runs use sign probabilities ω1 = ω2 = 0.7, so the blocks are imbalanced; the rescaling uses the balanced closed-form rate as an ansatz.Source: §B.1.5 · p. 28

B.2 Fig. 3a and Fig. 3b

  • The continuous-time ResNet experiments compare standard DSM with FM-weighted DSM on unbalanced and balanced quadrimodal GMMs, evaluating generated samples against training-distribution statistics.The quadrimodal target uses equal sign probabilities, while the unbalanced target retains ω = 0.8.Source: §B.2 · p. 28

B.2.1 Model architecture

  • Both ResNet experiments use a fully connected epsilon-predictor with width 128, four residual blocks, sinusoidal time embeddings, and a skip-connection score parameterization.Only the input and output projections depend on dimension.Source: §B.2.1 · p. 28

B.2.2 Training procedure

The experiments use weighted denoising score matching and track how mode directions and mixture weights emerge in Gaussian-mixture and MNIST settings. Their probes distinguish directional structure from class or mode frequency.

B.2.2 Training procedure

  • A variance-preserving Ornstein–Uhlenbeck process adds noise; the network predicts it, the score is recovered with Tweedie’s formula, and sampling integrates the reverse SDE.The training procedure compares uniform DSM with FM-weighted DSM.Source: §B.2.2 · p. 29

B.2.3 Unbalanced GMM, Fig. 3a

  • For the unbalanced GMM, cosine similarity tracks directional learning, while the positive fraction tracks acquisition of the dominant mode’s target weight ω.The comparison uses Uniform DSM and FM-weighted DSM at ω = 0.8.Source: §B.2.3 · p. 29

B.2.4 Quadrimodal GMM, Fig. 3b

  • For the quadrimodal GMM, cosine similarity to the subdominant direction is tracked across κ values, with the FM-weighted curves also rescaled by λ(√κ).Smaller κ denotes a weaker subdominant mode.Source: §B.2.4 · pp. 29–30

B.3 Fig. 4a and Fig. 4b

  • The MNIST experiments separate class-weight learning in grayscale digits from structure learning in randomly tinted digits with orthogonal digit and color directions.Tint contrast α controls the strength of the weaker color signal.Source: §B.3 · p. 30

B.3.1 Data

The MNIST experiments use the training split of digits 3 and 6, either grayscale with controlled class imbalance or randomly tinted at selected contrasts.

B.3.1 Data

  • Both datasets use 28 × 28 images from the MNIST training split, restricted to digits 3 and 6, with 6,131 threes and 5,918 sixes available.Source: §B.3.1 · p. 30
  • For grayscale MNIST, all sixes are retained and threes are subsampled to set the sixes’ target fraction ω to 0.5 or 0.8.The resulting totals are 11,836 images at ω = 0.5 and 7,397 at ω = 0.8.Source: §B.3.1 · p. 30
  • For tinted MNIST, all 12,049 images receive a random red or green tint independent of digit label, with contrast α ranging from 0.4 to 1.0.α = 0 gives grayscale and α = 1 gives a fully saturated tint; Figure 6 shows a sample at α = 0.7.Source: §B.3.1 · p. 30, Figure 6 · p. 30

B.3.2 Model architecture

The experiments use a fully connected ResNet epsilon-predictor, while the MNIST experiments use a shared U-Net backbone in Flow Matching and DSM roles.

B.3.2 Model architecture

  • The MNIST U-Net has about 1.08 million parameters, with 3 input/output channels for tinted images and 1 for grayscale images.It uses 32 model channels, attention at the coarsest resolution, and no class conditioning.Source: §B.3.2 · p. 30
  • The U-Net represents the Flow Matching velocity field or, under DSM, the score through a trainable skip connection.Source: §B.3.2 · p. 30
  • At initialization, the DSM skip connection makes the score default to −x, avoiding division of an untrained noise prediction by a potentially tiny σ(t) during early reverse-SDE sampling.Source: §B.3.2 · p. 31

B.3.3 Flow-matching training procedure

The experiments apply Flow Matching and DSM to MNIST and evaluate feature learning with fixed probes, then extend the analysis to human haplotypes using masked discrete diffusion.

B.3.3 Flow-matching training procedure

  • Conditional Flow Matching regresses the constant velocity from Gaussian noise to a clean sample with an MSE loss, then samples by integrating the learned field from t = 0 to t = 1.Sampling uses an adaptive-step Dormand–Prince solver.Source: §B.3.3 · p. 31

B.3.4 Weight learning, Fig. 4a

  • The weight-learning experiment compares DSM with Flow Matching on the same imbalanced dataset and monitors generated class frequency and cosine similarity to the nearer class mean.A balanced-data auxiliary classifier estimates the generated fraction of sixes, while class-mean projections provide the cosine measure.Source: §B.3.4 · p. 31

B.3.5 Structure learning, Fig. 4b

  • The structure-learning experiment measures digit-direction overlap and normalized color spread, and rescales training time by the relative direction amplitude to compare curves across α.The color direction’s squared amplitude scales as α².Source: §B.3.5 · p. 32

B.4.1 Dataset

  • The haplotype dataset contains 5,008 phased binary sequences of length 805 from 2,504 individuals, with a fixed PCA basis fit on the training split for evaluation.Source: §B.4.1 · p. 32

B.4.2 Masked Discrete Diffusion process

  • The haplotype model uses masked discrete diffusion, independently masking tokens under a log-linear noise schedule and training with a continuous-time ELBO on masked positions.Sampling starts fully masked and uses 30 reverse steps; PCA overlap is evaluated on generated samples.Source: §B.4.2 · p. 32

B.4.3 Model architecture

The Human Genome haplotype experiment uses a small Transformer denoiser and tracks generated-sample PCA directions against a fixed reference. Its time-rescaled PCA curves collapse for the two leading, mode-resolving directions; the rescaling is less meaningful for later directions without visible mode separation.

B.4.3 Model architecture

  • A learned token embedding and absolute positional embedding feed a time-conditioned Transformer encoder with a three-token vocabulary and a SUBS-parameterized output head.For the analyzed run, the model has 64 dimensions, 4 attention heads, 2 layers, and approximately 1.85 × 10^5 trainable parameters.Source: §B.4.3 · p. 32
  • Generated-sample PCA directions are matched one-to-one to reference directions by maximizing total overlap, avoiding misleading rankings of raw overlap entries.The reported overlap curves are means and 95% confidence intervals across four training seeds for the top five modes.Source: §B.4.3 · p. 33
  • The figure’s left panel projects all 4,508 real training haplotypes onto the top two reference PCA directions, while its right panel plots PCA overlap over training steps and rescaled steps.The reference directions and left-panel projections use the same reference PCA fit.Source: §B.4.3 · p. 33
  • The two leading Human Genome PCA directions collapse when plotted against their rescaled training times, while later directions without visible mode separation do not.The rescaling is based on explained variance along each reference direction, identified with its squared amplitude.Source: §B.4.3 · p. 33

C LLM usage

  • The authors disclose using LLMs to assist manuscript drafting, proof derivation, and code writing, with the authors retaining responsibility for scientific integrity.Source: C LLM usage · p. 34

Related papers

Loading 2609.35322v1…