ml-mdm: Efficient High-Resolution Diffusion Training with Matryoshka Models

The Multi-Resolution Training Problem ml-mdm Addresses

Training a high-resolution diffusion model directly at 1024x1024 pixels is computationally expensive and optimisation-intensive. Most approaches cascade separate models at increasing resolutions, which adds complexity and inference cost because each stage must be trained and run independently.

The Matryoshka Diffusion Models approach published in the accompanying arXiv paper (arXiv:2310.15111) addresses this by training a single pixel-space model that denoises images at multiple resolutions simultaneously. The name is a reference to the Russian nesting dolls: the model operates at nested resolutions within a single unified architecture.

ml-mdm is the Python implementation of this approach, released as an open-source package by Luke Carlson, Jiatao Gu, Shuangfei Zhai, and Navdeep Jaitly. The README states that a single model can be trained at resolutions up to 1024x1024 and demonstrates strong zero-shot generalisation using CC12M, a dataset of 12 million text-image pairs. This is the research claim to verify if you plan to reproduce the results; the README does not document the hardware or wall-clock time required for a full CC12M run.

Nested U-Nets and the Matryoshka Denoising Mechanism

The core model implementations are in the ml_mdm.models submodule, which contains U-Net and Nested U-Net architectures. The U-Net is the standard diffusion model building block. The Nested U-Net adds an inner network that processes a lower-resolution version of the noisy image alongside the full-resolution computation.

The ml_mdm.diffusion submodule holds the diffusion pipelines, including DDPM (Denoising Diffusion Probabilistic Models). Training and sampling share the same pipeline class, so a configuration used for training can be reused directly at inference time.

Configuration is handled through ml_mdm.config, which connects configuration dataclasses with the models, pipelines, and command-line interfaces using SimpleParsing. The README describes this as converting all CLI arguments and YAML files into clean configuration classes like ml_mdm.reader.ReaderConfig and ml_mdm.diffusion.DiffusionConfig. Configuration files for different resolution models live in the /configs directory, with separate entries for dataset creation, datasets, and models.

This architecture has a concrete implication for usage: the model selection and the dataset configuration are tightly coupled through the config system. Switching resolutions is a matter of pointing the CLI at a different YAML config file rather than changing code.

Installing ml-mdm and Downloading Pretrained Checkpoints

The default installation installs only the CPU-compatible dependencies, so the package can be installed even without a GPU. Install it in editable mode from the repository root:

pip install -e .

Developers should also set up pre-commit hooks:

pre-commit install

Three pretrained checkpoints trained on 50 million text-image pairs from Flickr are available for download at resolutions of 64x64, 256x256, and 1024x1024. Download them using curl:

export ASSET_PATH=https://docs-assets.developer.apple.com/ml-research/models/mdm

curl $ASSET_PATH/flickr64/vis_model.pth --output vis_model_64x64.pth
curl $ASSET_PATH/flickr256/vis_model.pth --output vis_model_256x256.pth
curl $ASSET_PATH/flickr1024/vis_model.pth --output vis_model_1024x1024.pth

The README notes that the repository also provides scripts for downloading CC12M and configs for training equivalent models on that dataset, so you are not restricted to the Flickr checkpoints if you want to train from scratch.

Running the Web Demo and the Test Suite

Once you have downloaded a pretrained checkpoint, you can run a local web demo that generates images from text prompts. The demo uses torchrun:

torchrun --standalone --nproc_per_node=1  ml_mdm/clis/generate_sample.py --port $YOUR_PORT

Replace $YOUR_PORT with a port number of your choice. The web UI will be available at that port in your browser.

The test suite separates CPU and GPU tests. The default run covers only CPU-compatible tests, which can be used to verify installation without GPU access:

pytest

To run all tests including GPU-specific ones:

pytest  -m ''

To run only the GPU tests:

pytest -m gpu

The README notes that users have run the codebase with Python 3.9 and 3.10 combined with CUDA 12 and CUDA 11.8. Python 3.11 and later versions are not explicitly listed, which means compatibility may require testing if you plan to use a newer Python runtime.

Repository Structure and the Configuration System

The top-level directory contains ml-mdm/ (the main package), ml-mdm-matryoshka/ (likely a variant), configs/, data/, docs/, and tests/. The /data directory contains multiple vocabulary files for different tokeniser types (BERT, C4, CIFAR10, T5, and a custom SPM vocabulary) along with prompt files for several datasets at different resolutions.

The /configs directory is split into dataset_creation, datasets, and models subdirectories. Each resolution model has its own configuration file. This is the primary lever for customising training: rather than modifying code, you copy and edit the appropriate YAML file and pass it to the CLI with --config_path.

The ml_mdm.clis submodule contains all CLI tools, with train_parallel.py described as the most relevant for training. The SimpleParsing integration means every dataclass field in the configuration classes is also available as a CLI argument, so short-running experiments can override individual settings without writing a new config file.

The /tests/test_files directory holds sample files for the test suite. These are used to verify that the dataset loading and model configuration code behaves correctly, independent of the full training pipeline.

Limitations and Cases Where ml-mdm Is the Wrong Choice

ml-mdm is a research codebase, not a production image generation service. It has no API server, no user management, and no content safety filtering. Deploying it behind a public endpoint without adding those layers would be inappropriate.

Training at 1024x1024 resolution requires substantial GPU resources. The README does not document specific GPU memory requirements or estimated training time on standard hardware. Before starting a training run, expect to experiment with batch size and gradient accumulation settings to fit your available memory.

The pretrained checkpoints were trained on Flickr data (50 million pairs). The README provides this as a starting point; it does not claim the Flickr-trained checkpoints generalise to all image domains. The CC12M configs allow retraining on a different dataset, but CC12M still has a creative commons licence that limits commercial use of images derived from it.

ml-mdm is a separate tool from production-oriented diffusion model frameworks like Hugging Face Diffusers. Diffusers has a broader model zoo, better documentation for deployment, and active community support. Researchers extending the Matryoshka technique should use ml-mdm; developers who just need a text-to-image generation endpoint should look at Diffusers instead.

Maintenance, Paper Reference, and Licence

The last push to the repository was on 2026-09-11. The project carries no GitHub releases. It accompanies the paper "Matryoshka Diffusion Models" by Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Josh Susskind, and Navdeep Jaitly (arXiv:2310.15111).

The pre-commit configuration and the .gitattributes file in the repository suggest the codebase is maintained with code quality tools, though these are primarily for contributors rather than users.

The project is released under the MIT licence. The pretrained checkpoints are hosted separately on Apple's developer assets server; the README does not specify a separate licence for the checkpoints, so their usage terms may differ from the package code. Checking the terms at the hosting URL before incorporating the checkpoints into a commercial workflow is advisable.

Editorial conclusion

ml-mdm is a good fit for researchers who want to replicate or extend the Matryoshka Diffusion Models paper, and for engineers who need a working end-to-end diffusion training codebase in Python with support for multiple resolutions. It is not a production image generation service and requires GPU access for meaningful training. Before starting, confirm that your Python version is 3.9, 3.10, or compatible with the cuda version in your environment, and verify that the pretrained checkpoint URLs at docs-assets.developer.apple.com are still accessible before building a workflow around them.

Frequently asked questions

Official sources

1. apple/ml-mdm on GitHub
2. Issues
3. License: MIT
4. Project website
5. README