driftax

JAX implementation of Generative Modeling via Drifting: https://arxiv.org/abs/2602.04770

This repo includes:

• 1D inverse-problem / conditional inference demos (with comparisons against other conditional density baselines).
• MNIST DiT-style drifting model following the authors' approach.

Results

1D Inverse Problems: Conditional Inference

We evaluate conditional inference on a toy inverse problem with a multi-modal conditional structure, and compare drifting against CFM and Mixture Density Network.

MNIST: DiT Implementation

We implement the authors' DiT-style generator setup for MNIST and visualize conditional samples.

---

If you use this library please cite the original authors here:

@article{deng2026drifting, title={Generative Modeling via Drifting}, author={Deng, Mingyang and Li, He and Li, Tianhong and Du, Yilun and He, Kaiming}, journal={arXiv preprint arXiv:2602.04770}, year={2026} }

---

To see more JAX projects check out Awesome JAX