Marrying Autoregressive Transformer and Diffusion with Multi-Reference Autoregression

TL;DR

TransDiff is a class-conditional image generator that puts an AR Transformer in front of a DiT-style diffusion decoder: the transformer encodes the class label (and any previously generated images of the same class) into high-level semantic features over continuous VAE latents, and a flow-matching DiT decodes those features into the final image. On ImageNet 256×256 the 1.3B Huge model reaches FID 1.61 with one AR step and 1.42 under Multi-Reference Autoregression (MRAR), which re-runs the AR encoder over up to 4 previously generated images of the same class to enrich the conditioning. Same family at 512×512 hits FID 2.51. Inference is faster than DiT-XL/2 and competitive with VAR/RAR at matched parameter count. Open-source (code + 7 checkpoints on HuggingFace), so it’s also the cleanest hybrid-AR+diffusion training recipe currently on the table.

Key claims

TransDiff is presented as the first framework that jointly trains an AR Transformer (as a semantic encoder over continuous VAE latents) with a diffusion decoder (DiT with flow matching) in a unified objective, rather than coupling them via tokenization (VQ + AR) or via a frozen text encoder (DiT + T5) [§3.1].
On ImageNet 256×256 with CFG, TransDiff-H MRAR (1.3B) reaches FID 1.42 / IS 301.2, beating DiT-XL/2 (2.27 / 278.2 at 675M), MDTv2-XL/2 (1.58 / 314.7 at 676M), VAR-d30 (1.92 / 323.1 at 2B), RAR-XXL (1.48 / 326.0 at 1.5B), and MAR-H (1.55 / 303.7 at 943M) [Table 3].
At matched parameter count (~683M), TransDiff-L MRAR reaches FID 1.49 versus DiT-XL/2 2.27 and MAR-L 1.78, while the 1-Step AR variant reaches FID 1.69 at 0.2s/image — faster than MAR-L’s 1.1s/image at 256-step diffusion-loss inference [Table 2 / Table LABEL:tab:DiffdecvsDiffloss].
Multi-Reference Autoregression (MRAR) is the paper’s generation paradigm contribution: instead of predicting the next token (token-AR) or the next scale (VAR-style scale-AR), the model predicts a complete next image whose features are appended to the AR Transformer’s context, then re-encoded with the class label to produce richer semantic features for the next diffusion decode [§3.2].
MRAR lowers FID from 1.69 (1-Step AR) to 1.49 at TransDiff-L and from 1.61 to 1.42 at TransDiff-H; the paper’s cosine-similarity diversity measure $M$ drops from 0.44 (1-Step) to 0.39 (MRAR) and rises to 0.64 under Scale-AR (which the paper reports as failing — FID 2.78) [Table 4.2.3].
Replacing MAR’s diffusion loss with a diffusion decoder (DiT) on the same AR encoder lowers FID from 1.78 to 1.49 and supports one-step inference (MAR-L at 1 step collapses to FID 336.58) [Table LABEL:tab:DiffdecvsDiffloss].
The number of references in MRAR has an optimal value: experiments at 0 / 4 / 16 / 64 references show 4 is the best operating point on ImageNet 256 [Fig. 6].
“Tiger-striped macaw” feature-fusion experiment: linearly interpolating 128 AR features from a Macaw image with 128 from a Tiger image produces an image with characteristics of both classes, evidence that the AR encoder’s output is a high-level semantic representation rather than a low-level pixel reconstruction [§4.2.2, Fig. 5 bottom].
512×512 results: TransDiff-L at 512² reaches FID 2.51 / IS 286.6 with 683M parameters under fine-tuning from the 256² checkpoint [Table 3 / Appendix D].
The VAE used is borrowed from MAR (KL-16, embed dim 16), so the architectural delta from MAR is entirely in the decoder (diffusion loss → DiT decoder) and in the generation paradigm (token-AR → MRAR) [§3.1, GitHub README].

Method

TransDiff has two trainable modules and one frozen one. The frozen module is an off-the-shelf KL-16 VAE that maps 256×256 images to a 16×16 continuous latent grid (the same VAE as MAR). The first trainable module is an AR Transformer encoder that takes a sequence of inputs — a class embedding, optionally a sequence of latents from previously generated images of the same class, and a learned mask embedding — and emits a sequence of conditioning features $z \in \mathbb{R}^{N \times C}$ . The second is a Diffusion Decoder: a DiT with rectified-flow / Flow Matching objective that, conditioned on $z$ , denoises Gaussian noise into the target VAE latent. Training is end-to-end with one loss — Flow Matching MSE on the velocity field — and gradients flow through both modules. Inference has two modes. 1-Step AR runs the AR transformer once over a class embedding plus mask, then runs the diffusion decoder; attention in the AR transformer is bidirectional. MRAR runs the AR transformer iteratively: it predicts an image, appends that image’s VAE latent into the AR context, re-runs the AR transformer (now under a causal attention mask), and predicts again; the final emitted image is the answer. The paper sweeps reference counts and finds 4 is best. CFG is applied at the diffusion-decoder level; per-model cfg, scale_0, scale_1 settings are tabulated in the README. Three sizes (Base 290M, Large 683M, Huge 1.3B) plus a 512×512 fine-tuned variant are released on HuggingFace.

Results

Headline ImageNet 256×256 numbers with classifier-free guidance [Table 3]:

Model	Params	FID	IS	Inference time
DiT-XL/2	675M	2.27	278.2	45s
MDTv2-XL/2	676M	1.58	314.7	—
MAR-L	479M	1.78	296.0	1.1s
MAR-H	943M	1.55	303.7	—
RAR-XXL	1.5B	1.48	326.0	6.4s
VAR-d30	2.0B	1.92	323.1	1s
TransDiff-L 1Step	683M	1.69	282.0	0.2s
TransDiff-L MRAR	683M	1.49	282.2	0.8s
TransDiff-H MRAR	1.3B	1.42	301.2	1.6s

Diffusion-decoder vs diffusion-loss head-to-head on the same AR encoder: TransDiff-L 1Step beats MAR-L at every step count and is the only one of the two that supports a single inference step (MAR-L at 1 step has FID 336.58) [Table LABEL:tab:DiffdecvsDiffloss]. MRAR vs alternative AR paradigms on TransDiff-L: 1-Step AR FID 1.69, Scale-AR 2.78, MRAR 1.49 — and on the diversity measure $M$ (cosine similarity over normalized AR features), Scale-AR is worst (0.64) and MRAR is best (0.39), supporting the paper’s “diversity drives FID” thesis [Table 4.2.3]. At 512×512, TransDiff-L MRAR reaches FID 2.51 / IS 286.6. Training was 800 epochs at 256² on 64 A100s (~115h) for the 1-Step phase; MRAR is an additional 40 fine-tuning epochs [GitHub README]. IS is consistently lower than RAR/VAR at matched FID — the paper does not discuss this trade-off in depth.

Why it’s interesting

TransDiff is a clean datapoint in an architectural direction that’s drawn several wiki entries recently — replace VQ-tokenization-and-cross-entropy with a continuous-token-and-diffusion-head — but is one of the few that does it for image generation rather than for LM. The closest sibling on the wiki is MAR (which TransDiff explicitly ablates against): MAR keeps the AR transformer as a per-token diffusion loss head, while TransDiff replaces the diffusion loss with a full DiT decoder conditioned on the entire AR output sequence at once, which is what enables one-step inference and the 1.78→1.49 FID drop. On the unified-AR+diffusion-LM axis, Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding does a structurally similar thing for text: train a single transformer under joint AR + diffusion objectives, then switch decoding mode by switching the attention mask — though NLD uses the diffusion mode as a drafter for AR speculation, whereas TransDiff uses the AR mode as a conditioner for diffusion decoding. The MRAR-as-paradigm framing is also worth flagging against The Diffusion Duality and Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking — both ask “what does autoregression mean in latent space?” — and against scale-AR systems like VAR. The fact that TransDiff fine-tunes a 1-Step AR checkpoint into MRAR for only 40 epochs is a useful operational signal: the paradigm switch is cheap if the underlying AR-encoder + DiT-decoder is already converged.