Ambient-Consistent Distillation (AC-DMD)

• The paper introduces AC-DMD as a cold-start method that matches student and teacher latent distributions using subinterval-wise reverse-KL objectives.
• The methodology employs a fake score model with ambient denoising score-matching and a consistency regularizer to ensure stability during generator updates.
• Empirical results show that AC-DMD reduces variance and improves performance metrics like PickScore, establishing its value as an effective initialization stage.

Searching arXiv for the cited paper and closely related distribution-matching distillation work. Ambient-Consistent Distribution Matching Distillation (AC-DMD) is the first-stage cold-start procedure introduced within Reward-Tilted Distribution Matching Distillation (RTDMD) for few-step flow generators. It matches a $K$-step student generator to a pretrained teacher on a sequence of subintervals $[t_k,1]$, while jointly training an auxiliary fake score model to approximate the student’s marginal scores on those same subintervals. Its distinguishing feature is a consistency regularizer applied to the fake score model, intended to help that model remain stable as the generator distribution shifts under limited updates. In the reported formulation, AC-DMD supplies a low-variance, scheduler-agnostic initialization for downstream reward-tilted optimization rather than serving as the complete alignment procedure by itself (Huang et al., 25 May 2026).

1. Role within few-step reward-tilted distillation

AC-DMD appears as the first of two stages in RTDMD. The full framework is described as unifying distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. Within that decomposition, AC-DMD is the distribution-matching stage: it trains the student generator $G_\theta$ against a teacher distribution $p_\psi$ and simultaneously trains a fake score model $s_\phi$ that estimates the student’s marginal scores over re-noised student latents. The second RTDMD stage then jointly optimizes the distribution-matching and reward-maximization terms; AC-DMD is the precursor that supplies the cold-start needed for that later optimization (Huang et al., 25 May 2026).

The setup is defined on latent variables. Let $p_\psi$ denote the pretrained teacher distribution on latents $x_0$, and let $G_\theta$ denote the $K$-step student. After step $k$, the student latent $[t_k,1]$0 has distribution $[t_k,1]$1. Rather than matching only a single terminal distribution, AC-DMD re-noises $[t_k,1]$2 forward to any $[t_k,1]$3 and matches the resulting student marginal to the teacher over that subinterval. This subinterval-wise construction is the central structural difference emphasized in the method description.

A plausible implication is that AC-DMD addresses a specific instability of few-step distillation: the score estimator used inside the generator update is trained on samples from a student distribution that changes during optimization. The consistency term is introduced precisely to regularize this estimator against such temporal drift.

2. Objective formulation and subinterval-wise matching

The generator objective at step $[t_k,1]$4 is a reverse-KL over the subinterval $[t_k,1]$5:

$[t_k,1]$6

The corresponding pathwise gradient is approximated as

$[t_k,1]$7

where

$[t_k,1]$8

and

$[t_k,1]$9

The induced student marginal on the subinterval is

$G_\theta$0

with Gaussian forward kernel

$G_\theta$1

The fake score model is trained with an ambient denoising score-matching objective:

$G_\theta$2

where

$G_\theta$3

In words, the generator update matches the student’s noised-latent distributions to the teacher on each subinterval, while the fake score model learns the student’s marginal scores on the same region. The summary identifies the DSM distance metric as $G_\theta$4 on velocity/score difference, equivalent to KL (Huang et al., 25 May 2026).

3. Forward re-noising construction and derivation assumptions

The subinterval KL term is derived by starting from the student latent after $G_\theta$5 steps, $G_\theta$6, and forwarding it to a later noise level $G_\theta$7 by

$G_\theta$8

This yields the student marginal $G_\theta$9 used inside the reverse-KL objective. The generator gradient then takes the usual reverse-KL form involving the difference between the student marginal score $p_\psi$0 and the teacher score $p_\psi$1, after which $p_\psi$2 is approximated by the learned fake score $p_\psi$3 (Huang et al., 25 May 2026).

Three assumptions are stated explicitly for this derivation. First, the noise schedule is rectified-flow, with $p_\psi$4 and $p_\psi$5. Second, the sampling kernel from $p_\psi$6 to $p_\psi$7 is Gaussian with known $p_\psi$8 and $p_\psi$9. Third, the generator’s current mapping $s_\phi$0 induces $s_\phi$1 exactly via coefficient-preserving sampling (CPS).

This framing distinguishes AC-DMD from full-interval formulations that match only over $s_\phi$2. The paper summary states that base A-DMD generalizes DMD2 when $s_\phi$3 and the full interval $s_\phi$4 is used, while remaining unbiased under any CPS scheduler. It also reports no quality loss when specializing to $s_\phi$5, but extension to $s_\phi$6. A common source of confusion is therefore to treat AC-DMD as merely a restatement of full-interval DMD2; the stated relation is instead one of generalization under CPS, with subinterval-wise matching as the operative construction.

4. Consistency regularization and fake-score stabilization

The consistency term is introduced because the ambient DSM loss $s_\phi$7 is described as unbiased but high-variance when $s_\phi$8 is trained on fresh student latents that shift each iteration. The formulation uses the denoiser or “$s_\phi$9-prediction” form $p_\psi$0 of the fake score model and imposes a self-consistency relation across adjacent noise levels:

$p_\psi$1

The optimal denoiser is stated to satisfy

$p_\psi$2

which provides the justification for the regularizer. In practice, the expectation over the learned reverse kernel is replaced by two independent samples $p_\psi$3, yielding the unbiased estimator

$p_\psi$4

The total fake-score loss is

$p_\psi$5

The reported rationale is variance reduction and stability as $p_\psi$6 shifts. This suggests that the consistency term should be read less as a change in the target score and more as a temporal regularizer on the estimator that approximates that target. The theoretical summary supports this interpretation: Proposition 1 states that the ambient DSM objective is unbiased and has the true marginal score $p_\psi$7 as its unique minimizer, while Proposition 2 supplies the self-consistency property used to motivate $p_\psi$8 (Huang et al., 25 May 2026).

5. Optimization procedure, hyperparameters, and architecture

The reported AC-DMD pseudocode takes as inputs the pretrained teacher score $p_\psi$9, student generator $x_0$0, fake score $x_0$1, subinterval endpoints $x_0$2, weights $x_0$3, $x_0$4, consistency weight $x_0$5, and the number of cold-start steps $x_0$6. For each update $x_0$7, the sequence is:

1. Sample a mini-batch of prompts, or pure Gaussian noise for the unconditional case.
2. Roll out the $x_0$8-step generator $x_0$9, recording each intermediate $G_\theta$0.
3. Uniformly sample a step index $G_\theta$1 and a noise level $G_\theta$2.
4. Form $G_\theta$3, with $G_\theta$4.
5. Update $G_\theta$5 by descending $G_\theta$6.
6. Update $G_\theta$7 by descending $G_\theta$8.

The notes accompanying the pseudocode state that subintervals $G_\theta$9 are inherent in the $K$0-step schedule, that $K$1 is applied at every fake-score update, and that $K$2 and $K$3 may be chosen constant or may follow the teacher’s weighting $K$4.

The hyperparameter and architecture choices are specified as follows. Typical $K$5, with

$K$6

The weights $K$7 and $K$8 are typically uniform in $K$9 or inherited from the teacher’s $k$0. The consistency weight $k$1 is ablated over $k$2, with best performance at $k$3. The CPS stochasticity parameter $k$4 controls injection noise; $k$5 works well, with default $k$6. Learning rates are reported as $k$7 for SD3* and $k$8 for FLUX.2 4B during the cold-start stage, shared by $k$9 and $[t_k,1]$00; during the RL stage, $[t_k,1]$01 retains the cold-start learning rate and $[t_k,1]$02 is reduced to $[t_k,1]$03. The generator $[t_k,1]$04 uses the same UNet/flow backbone as the teacher and is finetuned via LoRA with rank $[t_k,1]$05 and $[t_k,1]$06, while the fake score model $[t_k,1]$07 copies the teacher’s score-network architecture (Huang et al., 25 May 2026).

6. Guarantees, ablations, and empirical role

The theoretical guarantees summarized for AC-DMD are limited but concrete. Proposition 1, located in Appendix C, states that the ambient DSM objective $[t_k,1]$08 is unbiased and that its unique minimizer is the true marginal score $[t_k,1]$09. Proposition 2, in Appendix D, states that the optimal denoiser satisfies the self-consistency property used to justify $[t_k,1]$10. No formal global convergence proof is provided. The paper summary nevertheless reports that AC-DMD cold-starts reliably in practice and produces low-variance score estimates that accelerate downstream RL (Huang et al., 25 May 2026).

The ablations attribute a measurable share of performance to the consistency regularizer alone, before any RL is applied. For the consistency study in Table 6, base A-DMD with $[t_k,1]$11 yields CLIPScore $[t_k,1]$12, PickScore $[t_k,1]$13, and HPSv2 $[t_k,1]$14. AC-DMD with $[t_k,1]$15 yields CLIPScore $[t_k,1]$16, PickScore $[t_k,1]$17, and HPSv2 $[t_k,1]$18, corresponding to gains of $[t_k,1]$19, $[t_k,1]$20, and $[t_k,1]$21, respectively. Improvement is reported as robust across $[t_k,1]$22.

A separate ablation on the second RTDMD stage reports that naïve GRPO without consistency reaches PickScore $[t_k,1]$23, whereas SubGRPO with $[t_k,1]$24 plus deterministic final-step gradient reaches PickScore $[t_k,1]$25 and CLIPScore $[t_k,1]$26, with gains of $[t_k,1]$27 and $[t_k,1]$28. Each component, including shared-noise and final-step backpropagation, is said to yield consistent gains. Although those results belong to the later reinforcement stage, they reinforce the stated role of AC-DMD: it is the cold-start mechanism that makes the downstream reward-tilted optimization tractable under limited updates.

Taken together, the empirical summary characterizes AC-DMD as a low-variance, scheduler-agnostic cold-start whose most visible standalone contribution is the consistency regularizer. The most direct evidence for that characterization is the reported improvement of over $[t_k,1]$29 PickScore points before any RL is applied.