cs.LGOct 8, 2026
SaveAmbient Discrete Diffusion: Using the Wrong Data at the Right Time for Data Efficient Learning
Organizations: ETH Zürich · Harvard University · MIT
Abstract
We introduce RefineMix, a framework for training discrete diffusion models under severe data scarcity, a common constraint in scientific applications. RefineMix uses out-of-distribution data at selected diffusion times to improve generalization without biasing the sampling distribution. Although this strategy has been explored in continuous diffusion, discrete diffusion presents a distinct challenge: unlike Gaussian noise, masking preserves domain information in surviving tokens, limiting the use of related data at high noise levels. At low noise levels, however, the domains effectively disjoint supports become an advantage, allowing the model to learn from both in-domain and out-of-distribution data without biasing the sampler. We formalize these intuitions and provide a theoretical analysis for the proposed method. Experimentally, across five domain-shift settings, RefineMix matches or outperforms in-domain finetuning and data mixing. For protein sequence generation, finetuning with just 197 in-domain examples nearly doubles the fraction of generated proteins that are simultaneously novel, foldable, and in-family compared to standard finetuning.
Figures & tables
Figure 1: Noise-selective mixing. (a) Noising widens the target (green) and source (orange) distributions until both reach the prior at . They overlap (purple) only at high noise, so RefineMix admits source data only below the cutoff (blue). (b) How well a classifier tells the domains apart, against the noise level: for text (solid), it stays sharp until almost every token is masked and for images under Gaussian noise (dotted), it goes blind early. (c) Why text stays separable: masked sequences at decreasing noise (top to bottom), with the classifier’s probability of the target. A single target-specific token (green) settles the domain, whereas shared tokens (white) add little.
| Pair (src tgt) | Tgt. tokens | Fine-tune | Data Mixing | Conditional | RefineMix ( ) |
| Math Stat | 12.96 M | 19.92 0.01 | 19.61 0.01 | 19.10 0.00 | 18.29 0.01 (0.7) |
| Physics Hep-Ex | 4.71 M | 7.49 0.02 | 7.17 0.01 | 7.09 0.01 | 6.71 0.02 (0.8) |
| cc_news Hep-Ex | 4.71 M | 8.04 0.05 | 8.05 0.01 | 11.28 0.01 | 7.39 0.01 (0.8) |
| cc_news Textiles | 1.17 M | 20.40 0.04 | 19.73 0.12 | 30.12 1.19 | 19.70 0.06 (0.9) |
| Patents-G AG News | 5.67 M | 28.02 0.15 | 31.55 0.16 | 29.55 0.04 | 27.47 0.02 (0.2) |
Table 1: RefineMix beats the baselines when adapting a source-pretrained model. Target validation perplexity (lower is better) after a source checkpoint (masked kernel, three seeds). Fine-tuning uses the target data alone, data mixing adds the source at every noise level, the conditional model adds it with a domain label, and RefineMix adds it only below .
Figure 2: Testing the theory (masked kernel, three seeds). (a) CS Math. A fully mixed model completes a noised sequence with revealed tokens from CS, Math, or both. (b) cc_news Hep-Ex, from the source checkpoint. Target validation perplexity against the cutoff, whose endpoints are fine-tuning and data mixing. (c) CS Math, from scratch. Target validation perplexity against the fraction of the target corpus kept, for target-only training and RefineMix at its best cutoff.
| Model | Target | Fine-tuning | Data mixing | LoRA | RefineMix ( ) |
|---|---|---|---|---|---|
| MDLM 130M | Lean/mathlib | 3.289 0.010 | 3.439 | 3.775 0.007 | 3.234 0.003 (0.2) |
| ( Sahoo et al., 2024 ) | Rare disease | 15.703 0.033 | 15.770 | 18.859 0.117 | 14.933 0.016 (0.4) |
| ChemRxiv | 18.739 0.002 | 20.260 | 25.248 0.164 | 18.048 0.017 (0.4) | |
| SEDD 424M | Lean/mathlib | 3.027 0.004 | 3.147 | 3.314 0.013 | 2.987 0.005 (0.3) |
| ( Lou et al., 2024 ) | Rare disease | 14.072 0.002 | 13.720 | 15.473 0.009 | 13.275 0.055 (0.5) |
| ChemRxiv | 16.588 0.031 | 17.340 | 19.659 0.126 | 15.915 0.055 (0.4) |
Table 2: RefineMix adapts pretrained diffusion language models to scarce domains. Target validation perplexity (lower is better). LoRA (rank ) fine-tunes on the target alone. Mean std. over three seeds (data mixing: one). ∗ DiffuLLaMA uses LoRA (rank ) in every column.
| Family | Folds | Novel | All three | |
| Real peroxiredoxins (ceiling) | 94.0 0.8 | 98.2 0.9 | 78.1 6.8 | 75.1 5.7 |
| Fine-tuning | 51.8 5.0 | 37.5 4.6 | 99.0 0.3 | 35.8 4.6 |
| Data mixing | 36.2 2.2 | 35.7 2.2 | 98.7 0.2 | 32.9 2.2 |
| RefineMix ( ) | 79.8 2.4 | 74.3 2.9 | 89.8 1.8 | 63.3 2.0 |
| Generic UniRef50 (floor) | 0.0 | 65.1 | 100.0 | 0.0 |
Table 3: RefineMix writes new, folding members of a protein family from examples. Percentage of sequences generated by DPLM-650M that pass each test of Section 4 ; All three counts sequences that pass all tests at once. Real held-out peroxiredoxins and generic proteins mark the ceiling and the floor. Mean standard deviation over three train–test splits, each with three seeds.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
| Corpus | Tokenizer | Tokens |
|---|---|---|
| arXiv Math | bert-base-uncased | M |
| arXiv Stat | (our M DiT) | M |
| arXiv Physics | M | |
| arXiv Hep-Ex | M | |
| arXiv CS | M | |
| CC-News | M |
Table 4: Corpora , in tokens of the listed tokenizer. The upper block serves our M model (Table 1 and Figure 2 ), the lower two the pretrained models of Table 2 . Of the two web corpora we use only one parquet file each of OpenWebText ( documents) and of DKYoon/SlimPajama-6B . Figure 2(c) keeps down to of the Math training data.
| Pair (src tgt) | Tgt. tokens | Fine-tune | Data mixing | Conditional | RefineMix ( ) |
|---|---|---|---|---|---|
| cc_news Hep-Ex | 4.71 M | 8.49 0.03 | 8.76 0.02 | 12.29 † | 7.96 0.01 (0.7) |
Table 5: Comparison to baselines under the uniform kernel. Target validation perplexity, same setting as Table 1 . Mean standard deviation over three seeds; † one seed.
Figure 3: The three predictions of the theory under the uniform kernel (three seeds). Same experiments as Figure 2 ; in (b), the source checkpoint is a DUO model of the same size, trained on same datasets.
Figure 4: Ablation of the source rate (CS Math, of the target, trained from scratch, masked kernel, one seed). (a) Target validation perplexity against . (b) On-target rate, the fraction of samples that the domain classifier assigns to the target, at sampling temperature , against .
Figure 5: Samples generated by the models of Table 1 . For each of the five source–target pairs, all models start from the same source-pretrained checkpoint and are adapted to the target (masked kernel, mean over three seeds). We sample from pure noise over a sweep of sampling temperatures. (a) On-target rate: the fraction of samples that a domain classifier assigns to the target, against the sampling temperature, for fine-tuning, data mixing, the conditional model, and RefineMix. RefineMix stays on target at every temperature. (b) Generation frontier: generative perplexity of the samples, scored by OPT-2.7B, against their entropy, for fine-tuning and RefineMix. Perplexity is inverted, so up and to the right means more fluent and more diverse samples, and the dashed line marks real target text. The RefineMix frontier matches or lies above fine-tuning on every pair.
Figure 6: Domains stay separable almost to . Balanced accuracy of the domain classifier vs. noise level , six source–target pairs of our text experiments; chance is . Solid = far source, dashed = near. is a DDiT-style time-conditioned transformer ( hidden, blocks, heads, attention-mean pooling) under the log-linear schedule, whose corrupted fraction is linear in ; steps, batch , lr with cosine decay and warmup, no label smoothing. The two panels differ only in the corruption kernel.
Figure 7: Where the source helps the target. For each pair of Table 1 , data mixing and fine-tuning start from the same source-pretrained checkpoint (masked kernel, mean over three seeds). We plot the loss gap , the per-level cross-entropy of data mixing minus that of fine-tuning, on held-out target text. Below zero, the source improves the target at that noise level; above zero, the price of mixing has taken over. On the arXiv pairs, the gap turns positive at the chosen cutoff; elsewhere, it is a rough guide.
Figure 8: RefineMix keeps what the source checkpoint knows. For each pair of Table 1 , fine-tuning and RefineMix start from the same source-pretrained checkpoint (masked kernel, mean over three seeds). We plot the change in per-level cross-entropy on held-out source text between the checkpoint and the end of training; positive values mean that source knowledge was lost. The dashed line marks the cutoff of RefineMix, and the shaded region is where it admits the source. Fine-tuning loses source knowledge at every level, whereas RefineMix keeps it below the cutoff and moves away only above it.
Figure 9: Generation frontiers of the adapted models of Table 2 . Each of the four pretrained models is adapted to each of the three targets, and we sample from pure noise over a sweep of sampling temperatures (mean over three seeds). We plot the generative perplexity of the samples, scored by OPT-2.7B, against their entropy, for fine-tuning and RefineMix. Perplexity is inverted, so up and to the right means more fluent and more diverse samples, and the dashed line marks real target text.
Figure 10: The adapted models of Table 2 stay on target. Each of the four pretrained models is adapted to each of the three targets, and we sample from pure noise over a sweep of sampling temperatures (mean over three seeds). We plot the on-target rate, the fraction of samples that a domain classifier assigns to the target rather than the source, against the sampling temperature, for all methods.
Figure 11: Testing Theorem 1 directly (CS Math, three seeds). A fully mixed model ( , of the target) completes noised sequences in which target-specific tokens are revealed and the rest is masked. We plot two estimates of how often a completion is source text, against : the classifier’s source probability of the noised input, shifted to the prior equation 20 , and the classifier’s score of the completion. By Theorem 1 , the two agree for the optimal classifier; a trained classifier makes the input side an upper bound. Bands are bootstrap intervals over samples per 3 seeds; the dashed line marks .