autoregressive-diffusion-pytorch: A Reference Implementation for Token-Free Autoregressive Image Generation

The vector quantization step this package removes

Most autoregressive image models predict discrete tokens, which means an image first has to be compressed through a learned codebook. That codebook is a training problem of its own: it can collapse, it caps the fidelity of whatever is decoded from it, and it adds a stage between the image and the model. The paper this repository implements, Autoregressive Image Generation without Vector Quantization, is an attempt to drop that stage and let the autoregressive model predict continuous values directly. The README states the package is an implementation of the architecture behind that paper, and it links the official repository at LTH14/mar.

The audience is narrow and specific. This is for researchers and engineers who want to read the mechanism in PyTorch rather than in a research codebase, or who want a starting point they can modify: swap the backbone, change the diffusion parameterization, or attach their own dataset. The README shows a results image captioned oxford flowers at 96k steps, which tells you the intended use is training runs on image datasets, not loading a checkpoint and calling generate. There is no pretrained weights download in the material, no model zoo, and no inference CLI. If you want to generate images today from a hosted model, this is not that.

Two entry points: a generic sequence model and an image wrapper

The package exposes a low-level module and a task-specific one. AutoregressiveDiffusion takes dim_input, dim, max_seq_len, depth, mlp_depth and mlp_width, accepts a tensor of shape (batch, sequence, dim_input), returns a scalar loss, and exposes a sample method that produces a tensor of the same shape as the input. That is the whole contract shown in the README: forward for loss, backward, then sample. The shape assertion in the example is the only statement about output correctness the documentation makes.

ImageAutoregressiveDiffusion wraps that with a patchification layer. You pass a nested model dict (dim, depth, heads), an image_size and a patch_size, feed images of shape (batch, channels, height, width), and get a loss and samples back in image shape. The README example uses image_size 64 with patch_size 8, and a later example uses 128 with patch_size 16. The patch size is what converts the image into the sequence the autoregressive model consumes, so it is the parameter that most directly trades sequence length against spatial resolution.

The third piece is ImageTrainer, which pairs a model with an ImageDataset pointed at a directory path and is invoked by calling the trainer object. The README does not describe the training loop, the optimizer, the schedule, or checkpointing. That is a real gap: you can see the call signature, but the behaviour behind it has to be read in the source.

Flow matching as an import swap, and what xm_candidates changes

The repository ships a second parameterization. The README says that for an improvised version using flow matching, you import ImageAutoregressiveFlow and AutoregressiveFlow instead, and that the rest is the same. The flow example is otherwise identical to the diffusion example except for the class name and one extra argument.

That argument is xm_candidates, set to 4 in the example, with the README describing it as explorative modeling where training goes against the best of k candidates per timestep. This is the one place in the material where a design decision is spelled out rather than implied. Training against the best of several candidates changes what the model is optimized toward, and it costs proportionally more compute per step, since you are evaluating k candidates rather than one. The README does not say how candidates are drawn or how the best one is selected, so anyone using this flag should read the implementation before assuming what it does.

The citation list gives the provenance of these choices: rectified flow and scaling rectified flow transformers for the flow-matching path, the Karras et al. paper on diffusion design space, and a 2025 paper on denoising generative models, plus a 2026 paper on explorative modeling that appears to be the source of xm_candidates. Note that the last two citations carry future years relative to the 2025 release dates in the repository metadata, which is worth checking against the actual papers if you plan to cite this work.

Install and the smallest runnable example

Installation is a single command from the README: pip install autoregressive-diffusion-pytorch. The current release listed is 0.3.0, dated 2025-12-04, following 0.2.8 in November 2024 and 0.2.7 in September 2024. That spacing is informative. There was a gap of roughly thirteen months between 0.2.8 and 0.3.0, and the 0.3.0 release lands after the last recorded push to main in September 2026, so the release line and the branch are not obviously in lockstep. Pin a version if you need reproducibility.

The minimal sequence example is short enough to quote in structure: construct AutoregressiveDiffusion with dim_input 512, dim 1024, max_seq_len 32, depth 8, mlp_depth 3 and mlp_width 1024; create a tensor with torch.randn(3, 32, 512); call the model on it to get a loss; call loss.backward(); then call model.sample(batch_size=3) and assert the sampled shape matches the input shape. The image example follows the same pattern with ImageAutoregressiveDiffusion, a model dict of dim 1024, depth 12 and heads 12, image_size 64 and patch_size 8, on a tensor of shape (3, 3, 64, 64).

For training, the README shows ImageDataset taking a directory path and an image_size of 128, ImageAutoregressiveDiffusion configured with dim 512, image_size 128 and patch_size 16, and then ImageTrainer(model=model, dataset=dataset) followed by trainer(). Nothing else is configured in the example. There is no learning rate, batch size, epoch count, or output directory in the snippet, which means either the trainer supplies defaults or the example is incomplete. Treat the snippet as a starting point and read the constructor signature before running a long job.

Where this implementation stops short

The most concrete limitation is that the README documents nothing about sampling quality, sampling cost, or the number of steps required. The sample method returns the right shape; whether the contents are any good is not addressed. For a diffusion-style model that is a significant omission, because the number of denoising or integration steps is usually the dominant cost at inference, and the README never mentions a step count, a sampler name, or a schedule.

The trainer is the second gap. ImageTrainer is presented as a one-line call, and the README does not describe checkpointing, resumption, distributed training, mixed precision, or logging. For a 96k-step run on flowers, as the results caption implies, you would want at least checkpointing. If the trainer does not provide it, you are writing that yourself, which changes the effort estimate considerably. The README also gives no guidance on dataset layout beyond a directory path, so the expected on-disk structure is unclear from the documentation alone.

Third, this is a single-author research implementation. The release cadence shows long quiet periods, and the material contains no statement about supported PyTorch versions, CUDA versions, or tested hardware. If your environment is pinned, verify compatibility yourself. None of this makes the code wrong; it makes it a reference implementation rather than a dependency you can lean on without reading it.

Against the official MAR repository and the Transfusion route

The README names two alternatives directly. The first is the official repository for the paper, linked as LTH14/mar. The difference in approach is one of intent rather than algorithm: the official repository is the authors' code, which typically means it is tied to the exact training recipe and evaluation setup reported in the paper, while this package is a reimplementation organized around a small, importable module surface and a reusable trainer. If you want to reproduce reported numbers, the official code is the closer match. If you want to lift the architecture into your own project and change parts of it, the module boundaries here (AutoregressiveDiffusion, ImageAutoregressiveDiffusion, ImageAutoregressiveFlow, ImageTrainer) are easier to work with.

The second alternative is linked as transfusion-pytorch, described in the README as an alternative route. Transfusion-style models combine autoregressive and diffusion objectives over the same transformer rather than replacing token prediction with a diffusion head over continuous tokens, which is the distinction the MAR line of work is built on. The practical difference is what you get out of the model: a Transfusion-style setup keeps a next-token objective in the mix, while this package is aimed at continuous-valued autoregressive prediction with no codebook. If your data is naturally discrete (text, code, audio tokens), the vector-quantization-free premise buys you less and the Transfusion route may fit better. If your data is continuous and you want to avoid a learned codebook, this is the more direct match.

Licence and the cost of keeping up

The repository is MIT licensed, which permits commercial use, modification, and redistribution provided the copyright notice and permission notice are retained. That is permissive and low-friction. Two things to keep in mind without treating this as legal advice: the package implements methods described in several cited papers, and the citations in the README are to the papers, not to any code licence, so if you are reimplementing or redistributing anything derived from the official MAR repository you should check that repository's terms separately. Also, the README image assets and the results figure are part of the repository; MIT covers the code, and asset provenance is a separate question.

Maintenance cost is the more practical concern. With releases at 0.2.7, 0.2.8 and 0.3.0 spread across roughly fifteen months, and a last push recorded in September 2026, the project moves in bursts. Upgrading between minor versions is therefore a real task: read the diff, not just the version number. Because the public surface is small (three model classes, one dataset class, one trainer), the blast radius of an upgrade is limited, and pinning to 0.3.0 in a requirements file is a reasonable default. The cost you should budget for is not dependency churn but the reading time: the README gives you the shape of the API and nothing about the training loop, the sampler, or the defaults, so the source is the documentation.

Editorial conclusion

Adopt it if you are reimplementing or extending the MAR architecture and want a readable PyTorch reference with an image trainer already wired up. Do not adopt it if you need a packaged inference server, a stable API surface, or a maintained benchmark suite; nothing in the repository promises any of those. Before you commit, check the release history against the main branch, confirm which of AutoregressiveDiffusion, ImageAutoregressiveDiffusion and ImageAutoregressiveFlow matches the paper variant you care about, and read the training loop in image_trainer.py to see what the trainer does and does not handle.

Official sources

1. Issues
2. License: MIT
3. lucidrains/autoregressive-diffusion-pytorch on GitHub
4. README
5. Releases