Ambient Discrete Diffusion: Using the Wrong Data at the Right Time for Data Efficient Learning

1 INTRODUCTION

• Data-challenging training for discrete diffusion remains underexplored, despite data scarcity limiting effective model training in some application domains.The introduction notes that final-model performance and generalization still depend on available training data.Source: §1 · p. 1
• Uniformly mixing related-domain data can bias the learned distribution, while pretraining followed by finetuning is prone to overfitting and requires careful tuning.Continuous-diffusion approaches also use out-of-distribution data at selected diffusion times, but that strategy had not been explored for discrete diffusion.Source: §1 · pp. 1–2
• Discrete masking can preserve domain separability even at extreme noise, whereas disjoint low-noise supports let source data help without biasing target sampling.RefineMix therefore uses target data at high noise and both domains at low noise; its theory analyzes the safe cutoff and the trade-off between shared structure and bias.Source: §1 · p. 2
• Across five domain pairs, RefineMix matches or outperforms finetuning, data mixing, and conditional training; with 197 protein-family examples, it nearly doubles generated sequences that are new, fold, and belong to the family.The paper also reports improvements when adapting pretrained diffusion language models up to 7B parameters.Source: §1 · p. 2

2 BACKGROUND

• Discrete diffusion corrupts clean categorical tokens with a forward process and trains a network to predict clean tokens for the reverse process.The forward process mixes each clean token with a fixed prior; the masked kernel uses a [MASK] token as that prior, while the uniform kernel uses a uniform prior.Source: §2 · pp. 2–3
• For sequences, tokens are noised independently, while the network sees the whole noisy sequence and predicts each clean token; the loss sums over positions.The text uses x for sequences and z for tokens.Source: §2 · p. 3
• The problem setting has few samples from target distribution p and many more samples from a distinct source distribution q.The background denotes their noisy distributions at time t by p_t and q_t.Source: §2 · p. 3

3.1 THE REFINEMIX ALGORITHM

• Below cutoff tmax, RefineMix samples source data with probability λ and target data otherwise; above it, training uses only target data.The schedule, loss, network, and sampler otherwise remain as in regular discrete diffusion training.Source: §3.1 · p. 3
• At high noise, domain-specific tokens can steer generation toward the wrong domain, while at low noise visible target tokens identify the domain and source data can teach shared structure.Unlike Gaussian noise in continuous diffusion, masking can preserve domain information in surviving text tokens.Source: §3.1 · pp. 3–4

3.3 WHAT MIXING CAN COST

• Under distinct-domain assumptions, the probability of generating a source sequence is λ/(1 − λ) times the overlap C(tmax), and any later noise level bounds that cost from above.C(t) measures how often a classifier mistakes a noised source sequence for a target one; the overlap grows with noise.Source: §3.3 · p. 5

3.4 WHAT MIXING CAN GAIN

• With finite target data, source samples can improve the shared realization learned across domains, while introducing bias into the domain-specific content factor.The analysis assumes distinct domains, a shared realization, and equally informative source and target samples for learning that shared factor.Source: §3.4 · pp. 5–7
• For source-pretrained model adaptation, Table 1 reports that RefineMix beats the baselines on target validation perplexity (lower is better), using a masked-kernel source checkpoint and three seeds.Source: Table 1 · p. 7
• The best cutoff rises as the target corpus shrinks and is expected to fall when domains share fewer writing rules; a loss-gap sign change or a cutoff sweep can guide its selection.The loss gap compares fully pooled and target-only models on held-out target text; the pooled model must be trained to convergence for this proxy.Source: §3.4 · p. 7

4 EXPERIMENTS

• The experiments compare target validation perplexity across source–target pairs using the same 139M-parameter DiT, with three seeds per setting.Generation frontier and target-domain sample fraction are also reported in Appendix B.3.Source: §4 · p. 8
• RefineMix reaches the lowest or equal target perplexity on every domain pair and keeps samples on target with a better generation frontier than fine-tuning.The best cutoff ranges from 0.7–0.9 for close pairs to 0.2 for Patents-G → AG News.Source: §4 · p. 9
• On CS → Math with 1% target data, about four revealed target tokens shift fully mixed sampling to the target domain, while shared tokens do not.The best cutoff is 0.8, and every tested cutoff beats fine-tuning and data mixing; as target data shrinks to 1%, RefineMix more than halves target-only perplexity.Source: §4 · p. 9
• In twelve pretrained-model settings spanning 130M to 7B parameters and three scarce targets, RefineMix beats fine-tuning.LoRA does worse than full fine-tuning in these experiments.Source: §4 · p. 9
• With 197 in-family protein examples, RefineMix produces roughly twice as many sequences that are novel, fold, and belong to the family as fine-tuning.The evaluation requires each generated sequence to pass all three tests.Source: §4 · p. 9
• The cutoff remains a tuning requirement, and the theory describes an ideal model rather than the trained network.The paper identifies a cheap pre-training cutoff diagnostic and per-sample cutoffs as directions for future work.Source: §4 · p. 10

A RELATED WORK

• RefineMix changes which data the objective uses at each noise level, unlike prior approaches that keep training data fixed or use one corpus proportion throughout training.For discrete diffusion, the related source is used only at low noise, with a criterion opposite to Ambient Diffusion Omni’s image criterion.Source: A RELATED WORK · p. 16

• Text uses bert-base-uncased tokens in non-overlapping 512-token blocks, with documents split 90/10 into training and validation sets.The [MASK] token is the absorbing state, and blocks contain [CLS], 510 text tokens, and [CLS].Source: B EXPERIMENTAL DETAILS · p. 16
• All text models use a 139.3M-parameter DiT, a masked SUBS kernel, continuous-time NELBO training, and EMA weights for validation and sampling.Training uses AdamW, a constant 3 × 10^-4 learning rate after warm-up, and 256 blocks per step.Source: B EXPERIMENTAL DETAILS · p. 16
• Data mixing and RefineMix share a loader that selects source blocks with probability 0.9 below the cutoff and target blocks otherwise; data mixing sets the cutoff to 1.All settings use seeds 1, 2, and 3.Source: B EXPERIMENTAL DETAILS · p. 17
• Target validation perplexity is the exponential of per-token NELBO over the full target validation split, and validation selects checkpoints only.Source: B EXPERIMENTAL DETAILS · p. 17
• Baseline comparisons start from a common source-pretrained checkpoint and compare fine-tuning, data mixing, conditional training, and RefineMix.The cutoff search starts where the loss gap turns positive and advances in 0.1 steps until both neighboring values are worse.Source: B EXPERIMENTAL DETAILS · p. 17
• The revealed-token study samples completions after revealing 0–16 selected tokens from Math validation blocks, using 1,000 ancestral steps and a domain classifier.Token groups distinguish Math, CS, and shared words; each point averages 256 completions per model over three models.Source: B EXPERIMENTAL DETAILS · p. 17
• The cutoff sweep tests cc_news → Hep-Ex at values from 0 to 1 in 0.1 increments, while the target-size study reduces Math training data from 100% to 1%.The uniform-kernel cutoff experiment uses a DUO model.Source: B EXPERIMENTAL DETAILS · p. 17
• Protein experiments use 197–200 peroxiredoxin training sequences per cluster-held-out split and generate 256 sequences for evaluation.A generated sequence counts as “All three” only if it is family-matched, folds above the stated pLDDT threshold, and is novel under the identity criterion.Source: B EXPERIMENTAL DETAILS · p. 18

B.1 RESULTS UNDER THE UNIFORM KERNEL

• RefineMix again reaches the lowest target perplexity under the uniform kernel.Table 5 reports target validation perplexity; results are mean ± standard deviation over three seeds, except †one seed.Source: §B.1 · p. 18, Table 5 · p. 18
• A few revealed target tokens keep sampling on target, and the best cutoff lies between the two baselines and rises as the target shrinks.The three predictions of the theory are tested under the uniform kernel in Figure 3.Source: §B.1 · p. 18, Figure 3 · p. 18

B.2 ABLATION OF THE SOURCE RATE

• On CS →Math with 1% of the target, λ = 0.9 gives the lowest target validation perplexity, while samples stay on target at higher source rates.The sweep is trained from scratch with cutoff tmax = 0.9; Figure 4 plots validation perplexity and on-target rate against λ.Source: §B.2 · p. 19, Figure 4 · p. 19
• For RefineMix, λ is the probability that a training item below the cutoff is a source sample.The experiments fix λ = 0.9 apart from the source-rate sweep.Source: §B.2 · p. 19
• Data mixing is the tmax = 1 endpoint, but λ is not tuned because its generated samples come from the source at rate λ.It remains a comparison because it admits source data at every noise level and indicates how much the source can improve target likelihood; perplexity is reported with on-target rate.Source: §B.2 · p. 19

B.3 ADDITIONAL RESULTS FOR §4: REFINEMIX AGAINST FINE-TUNING

• The supporting analyses also measure what each model retains from its source checkpoint.Figure 8 compares per-level cross-entropy on held-out source text.Source: §B.3 · p. 19
• RefineMix stays on target at every sampling temperature, and its generation frontier matches or lies above fine-tuning on every pair.Figure 5 plots on-target rate against temperature and generative perplexity against entropy; higher and farther right indicates more fluent and diverse samples.Source: Figure 5 · p. 20
• Figure 6 measures domain-classifier balanced accuracy against noise level for six source–target pairs, with chance marked at 0.5.Solid lines denote far sources, dashed lines near sources; the panels differ by corruption kernel.Source: Figure 6 · p. 20

B.4 WHERE THE SOURCE HELPS, AND WHAT FINE-TUNING FORGETS

• On the two arXiv pairs, the loss gap turns positive near the chosen cutoff; on cc_news →Textiles it stays near or below zero, consistent with its high cutoff.The gap is data-mixing loss minus fine-tuning loss on held-out target text; below zero, source data improves target loss.Source: §B.4 · pp. 20–21, Figure 7 · p. 21
• The loss gap is only a rough cutoff guide on the remaining pairs: it turns positive too early on cc_news →Hep-Ex and stays positive on Patents-G →AG News, where RefineMix at 0.2 still beats fine-tuning.The plotted gap compares data mixing with fine-tuning at each noise level on held-out target text.Source: §B.4 · pp. 20–21, Figure 7 · p. 21
• Fine-tuning ends above the source checkpoint’s held-out-source cross-entropy at every noise level, while RefineMix stays at or below it below the cutoff and moves away only above it.The curves show how much source knowledge is lost, not when, because only final checkpoints were kept.Source: §B.4 · p. 21, Figure 8 · p. 21

C SAMPLES OF THE ADAPTED PRETRAINED MODELS

• Figure 9 compares generative perplexity, scored by OPT-2.7B, with sample entropy across temperatures for fine-tuning and RefineMix.Perplexity is inverted, so up and right indicate more fluent and diverse samples; the dashed line marks real target text.Source: C SAMPLES OF THE ADAPTED PRETRAINED… · p. 21, Figure 9 · p. 22
• Figure 10 plots the fraction of samples classified as target-domain against sampling temperature across adapted models and methods.The caption reports that the adapted models stay on target, averaging over three seeds.Source: C SAMPLES OF THE ADAPTED PRETRAINED… · p. 21, Figure 10 · p. 23

D POSTERIORS AND LOSSES FOR THE TWO KERNELS

• For masked diffusion, the posterior keeps unmasked tokens and reveals masked ones with probability (αs − αt)/(1 − αt), yielding a weighted cross-entropy only at masked positions.This is the continuous-time objective used in the experiments.Source: D POSTERIORS AND LOSSES FOR THE TWO… · p. 24
• For uniform-state diffusion, no state is absorbing, and every position carries signal, so the KL is evaluated in full across all L positions.The posterior places mass on the current token, the clean token, and uniformly on all others.Source: D POSTERIORS AND LOSSES FOR THE TWO… · p. 24

• Both proofs concern exact whole-sequence denoisers, with tmax on the sampler’s time grid; each reverse step is the exact reverse of the forward chain started at the denoiser.Source: E PROOFS · p. 24

E.1 PROOF OF THEOREM 1

• Under distinct, disjoint domains, RefineMix’s source-sequence probability is bounded by the mixture-scaled overlap at tmax, which is no greater than the corresponding overlap at any later t′.The result holds for both masked and uniform kernels.Source: §E.1 · pp. 24–25
• Below the cutoff, the sampler uses the mixed denoiser, while above it target-only denoising preserves the target marginal through tmax.Source: §E.1 · p. 25
• The overlap C(t) is nondecreasing in t, supplying the monotonicity used in the source-probability bound.Source: §E.1 · p. 25

E.2 PROOF OF THEOREM 4

• The price of RefineMix is governed by the same classifier as Theorem 1 and is bounded by B, which increases from 0 to −log(1 − λ).The argument uses the shared Markov forward process and applies to both kernels.Source: §E.2 · pp. 26–27
• The denoiser error separates into specific and shared factors; the specific factor uses the same NT target samples in both models, while the shared factor’s effective sample size changes at the cutoff.Source: §E.2 · p. 27
• The best cutoff is nonincreasing in the number of target samples NT, for any price that does not depend on NT.Source: §E.2 · p. 27

• The test compares completion-based source scores with a classifier’s source probability on controlled inputs, using a balanced classifier whose output is shifted to the mixture prior.The input-side estimate is an upper bound for the generated source rate when the trained classifier is calibrated.Source: F TESTING THEOREM 1 DIRECTLY · pp. 27–28
• Four revealed target tokens bring the generated source curve near zero under masking; from eight tokens onward, the model no longer writes source text.Source: F TESTING THEOREM 1 DIRECTLY · p. 28, Figure 11 · p. 29
• Under uniform noise, the generated curve needs twelve to sixteen revealed tokens to reach zero, as scrambled tokens obscure which positions carry evidence.Source: F TESTING THEOREM 1 DIRECTLY · p. 28