d3pm

Discrete denoising diffusion probabilistic models in MLX

d3pm is a small library for discrete denoising diffusion probabilistic models (D3PM) in MLX.

Important

🚀 DiffusionGemma

🔥 d3pm now features a 💎 DiffusionGemma 💎 example. DiffusionGemma trains on tiny_shakespeare and generates text in a block-autoregressive fashion with a Gemma-4 MoE backbone, uniform (mask-free) diffusion, self-conditioning, and D-CFG guidance. It supports prompt-conditioned generation via causal prompt prefill with bidirectional denoising.

Examples

Self-contained examples can be found in examples.

Installation

To install the latest GitHub , just call the following on the command line:

pip install git+https://github.com/dirmeier/d3pm-mlx@<RELEASE>

Please do not directly install from the main branch, since these commits are not tested, but prefer using PyPI or a stable release.

Contributing

Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled good first issue.

In order to contribute:

1. Clone d3pm-mlx and install uv from here.

2. Install all dependencies using uv sync --all-groups.

3. Install the Git hooks:

   uv run pre-commit install -t pre-commit -t commit-msg

4. Create a new branch locally, e.g. git checkout -b feature/my-new-feature.

5. Implement your contribution and ideally a test case.

6. Check your work (see below).

7. Submit a PR 🙂.

Development commands

The project uses uv for everything:

uv sync --all-groups
uv run ruff format src examples
uv run ruff check --fix src examples
uv run mypy src examples
uv run pre-commit run --all-files

References

• Arriola et al., Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models (BD3LM), 2025.
• Schiff et al., Simple Guidance Mechanisms for Discrete Diffusion Models (UDLM), 2024.
• Sahoo et al., The Diffusion Duality (Duo), 2025.
• Sahoo et al., Simple and Effective Masked Diffusion Language Models (MDLM), 2024.
• Lou et al., Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution (SEDD), 2024.
• Austin et al., Structured Denoising Diffusion Models in Discrete State-Spaces (D3PM), 2021.