Learning from Context: Leave-One-Out Prediction in Jump Diffusions

Continuous diffusion models generate discrete data through stochastic dynamics in an embedding space. We consider jump diffusions that additionally reset coordinates to embeddings of categorical values. A coordinate's own corrupted state can make prediction easy, potentially reducing the pressure to learn dependencies between coordinates. We study leave-one-out (LOO) prediction, in which each coordinate is predicted from the corrupted states of the other coordinates. We show that, under stated assumptions, jump diffusions with complete coordinate resets and state-independent clocks preserve the prescribed marginals when reset destinations use exact LOO probabilities and the drift uses the full endpoint posterior. We further show that this drift can be constructed from LOO probabilities by analytic likelihood completion without additional network evaluations. We evaluate unrestricted, LOO, and completed predictions for character-level text generation on Text8 and molecular graph generation on ZINC250k. Compared with unrestricted prediction, LOO improves lexical statistics on Text8 and the fraction of valid, connected molecules for absorbing jump diffusion on ZINC250k. For molecular pure diffusion, completion outperforms using LOO probabilities directly in the drift.