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

The gist The RefineMix framework introduces a method for training discrete diffusion models under severe data scarcity by using out-of-distribution data at selected diffusion times to improve generalization without biasing the sampling distribution

Problem Definition

Discrete diffusion modeling has recently gained significant research traction due to its potential for faster sampling compared to the autoregressive paradigm (Page 1). The problem of developing techniques for training discrete diffusion models in data-challenging settings remains underexplored (Page 1). The paper studies the setting where a few samples from a target domain and abundant samples from a distinct source domain are given, formally defined as having small set of NT samples from the desired target distribution p and NS >> NT samples from a distinct source distribution q (Page 2). The noisy distribution of the target domain at diffusion time t is denoted as pt(xt) and similarly qt(xt) for the source domain at time t (Page 2).

RefineMix Algorithm

The RefineMix algorithm improves over standard baselines by deciding, at each noise level, whether to mix the training data (Page 3). The core idea is to deviate from standard discrete diffusion training by changing the pool of sentences that can be sampled from based on the diffusion time t (Page 3). Specifically, below a cutoff tmax, the sequence comes from the source with probability λ and from the target otherwise, and above it, only from the target (Page 3). This strategy is justified by two observations: first, that in discrete diffusion for textual data, the noising mechanism is not an effective way to erase distribution shifts (Page 2). Second, at low noise levels, different domains do not have any shared data distributional support (Page 2).

Theoretical Justification

The theoretical justification separates the effect of mixing into what it can cost and gain (Page 3). The main risk of mixing is that the model samples from the wrong domain, which is analyzed through a binary classification problem (Page 3). Theorem 1 determines what mixing costs in sampling, showing that the sequence x0 that RefineMix generates from pure noise is a source sequence with probability P(ˆx0 ∈ A) = λ / (1 − λ) C(tmax) (Page 25). The overlap C(t) measures how often the classifier mistakes a noised source sequence for a target one, and the theorem shows that as long as the classifier can still tell the domains apart at the cutoff, RefineMix rarely generates samples from the wrong domain (Page 25).

What Mixing Can Gain

The gain of mixing is derived by assuming that domains share a latent content variable v (Page 3). The denoiser splits into two factors: a specific factor p(v xt) and a shared factor r(x0 xt, v) (Page 6). The specific factor is learned from the NT target samples, while the shared factor is learned from both corpora below the cutoff tmax (Page 6). This leads to Theorem 4, which states that the gain it brings below the cutoff is traded against a price that is paid once, at the cutoff (Page 27).

Empirical Results and Conclusion

Experimentally, RefineMix matches or outperforms in-domain finetuning and data mixing across five domain pairs (Page 1). 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 (Page 1). The method is shown to beat baselines on every pair and improves pretrained models from 130M to 7B parameters (Page 8). The conclusion is that discrete noise, unlike Gaussian noise on images, does not erase the domain of a text (Page 10). In summary, RefineMix admits source data only below a cutoff where a classifier can still tell the domains apart (Page 10).

How it works

The training step in RefineMix involves drawing a sequence from either the source or target based on t and then training the network with the loss equation 2 (Page 3). The denoiser of RefineMix is defined as pRefineMix(x0 xt) = (m(x0 xt), t ≤ tmax, p(x0 xt), t > tmax) (Page 3). This structure ensures that sampling stays on the target domain above the cutoff because only the target denoiser acts above it (Page 3). The best cutoff is found by maximizing a function Fκ(t) which balances the gain of learning shared realization against the bias introduced by mixing (Page 27).

What Mixing Can Cost

The cost of mixing is quantified by Theorem 1, which shows that the cost of mixing equals how often the classifier mistakes such a sequence for a source sequence (Page 3). This probability is expressed as Ex∼q(xt) [c(x)] (Page 25). The analysis shows that the optimal cutoff t⋆ is nonincreasing in NT and nondecreasing in κ, which decreases with NT (Page 27).

Best Cutoff Selection

The best cutoff t⋆ is determined by finding the largest maximizer of Fκ(t) where κ = 1/NT − 1/Neff (Page 27). The argument shows that t⋆ is nonincreasing in NT, meaning it rises as the target corpus shrinks (Page 27).

Testing Theorem 1 Directly

To test Theorem 1 directly, the authors reveal n target-specific tokens and compare how often the model completes a sequence as a source sequence against the classifier's probability of assigning it to be a source (Page 29). The input side of this comparison is shown to be an upper bound for the output side when using a trained classifier cˆ, which is calibrated (Page 29).

A Related Work

Existing methods often act only after training data is fixed, converting weights of a pretrained autoregressive model or post-training a diffusion model (Page 16). RefineMix differs because it uses samples of another domain unaltered, with the standard loss, and only below a cutoff (Page 16).

Experimental Details

All text models use the same DiT denoiser with 139.3M parameters (Page 17). Training involves AdamW with a constant learning rate of 3 × 10−4 after 2,500 linear warm-up steps on one NVIDIA RTX PRO 6000 or B200 (Page 17). Data mixing and RefineMix share one data loader where each item pairs a target block with a source block, drawing t per item (Page 17).

Conclusion

RefineMix beats data mixing, conditional training, LoRA, and full fine-tuning across domain pairs (Page 8). It also improves pretrained models and protein families (Page 8). The main limitations are that the cutoff must still be chosen and the theory describes an ideal model rather than the trained network (Page 10).

References

Asad Aali, Marius Arvinte, Sidharth Kumar, and Jonathan I. Tamir. Solving inverse problems with score-based generative priors learned from noisy data. In 2023 57th Asilomar Conference on Signals, Systems, and Computers, pp. 837–843. IEEE, October 2023 (Page 11).

Asad Aali, Giannis Daras, Brett Levac, Sidharth Kumar, Alex Dimakis, and Jon Tamir. Ambient diffusion posterior sampling: Solving inverse problems with diffusion models trained on corrupted data. In The Thirteenth International Conference on Learning Representations, 2025 (Page 11).

Alon Albalak, Liangming Pan, Colin Raffel, and William Yang Wang. Efficient online data mixing for language model pre-training, 2023 (Page 16).

arXiv.org submitters. arxiv dataset, 2024 (Page 17).

Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2023 (Page 17).

Weimin Bai, Yifei Wang, Wenzheng Chen, and He Sun. An expectation-maximization algorithm for training clean diffusion models from corrupted observations. In Proceedings of the 38th International Conference on Neural Information Processing Systems, NIPS ’24, Red Hook, NY, USA, 2025 (Page 17).

Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine Learning, 79(1–2):151–175, 2010 (Page 16).

1. A data-efficient discrete diffusion model that leverages out-of-distribution data selectively at low noise levels to improve generalization without biasing sampling, as formalized by RefineMix. This system can adapt pretrained models to scarce domains, such as protein sequence families (e.g., generating new peroxiredoxins), by using data only from the target domain in the high-noise regime and data from both domains in the low-noise regime.

2. A method to rigorously determine a safe cutoff time, denoted as t⋆, for mixing source and target data based on empirical loss gap analysis. This allows researchers to find the best cutoff t⋆ which is nonincreasing in NT (target size) and nonincreasing in NT, balancing the gain of shared realization against the bias introduced by the source data.

3. A mechanism to quantify the trade-off between learning from a source and introducing bias, specifically through the price term B(tmax). This allows for understanding what it costs when mixing, as the price is governed by the same classifier as Theorem 1, and analyzing how this cost relates to the gain derived from shared realization.

4. A diagnostic tool to predict optimal data mixing cutoffs before training by monitoring the loss gap g(t). This allows researchers to select t⋆ where g turns positive almost exactly at the chosen cutoff, which is a practical way to find the best balance between source knowledge and target performance.

5. An adaptation framework for large pretrained diffusion language models (up to 7B parameters) to highly specific, scarce domains by using RefineMix instead of fine-tuning or LoRA. This results in outperforming fine-tuning approaches and improving protein family generation by nearly doubling the fraction of generated sequences that are simultaneously novel, foldable, and in-family.