Research map / Image generation / Autoregressive and latent generation
Learning Stackable and Skippable LEGO Bricks for Efficient, Reconfigurable, and Variable-Resolution Diffusion Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes LEGO bricks, stackable and skippable diffusion backbone units that enrich local patches with MLP tokens and orchestrate them with DiT-style Transformer blocks, enabling faster training, cheaper sampling via brick skipping, and variable-resolution generation.
Problem
Diffusion models are computationally expensive to train and sample, and existing backbones (U-Net, ViT/DiT) are deep, rigid at sampling time (cannot be reconfigured or shrunk without retraining), and struggle to generate resolutions above the training resolution.
Approach
A LEGO brick enriches non-overlapping local patches into token embeddings via an MLP and mixes them with short-span Transformer (DiT) blocks while keeping the full-resolution image at input and output. Bricks of different patch sizes are stacked vertically (progressive growth, refinement, or U-shaped), trained on sampled patches, and selectively skipped at different timesteps during sampling.
Key results
- On ImageNet 256x256 with classifier-free guidance, LEGO-XL-PG reaches FID 2.05 and IS 289.12, versus DiT-XL/2 at 2.27 FID.
- On CelebA 64x64, LEGO-S-PR achieves FID 2.09 at 0.2G FLOPs, versus DiT-S/2 at 2.52 FID with 0.4G FLOPs.
- The paper reports about 60% sampling-time reduction compared with DiT while keeping the same NFE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use LEGO-PG instead of DiT | ImageNet 64×64 class-conditional generation · ImageNet 64×64 | FID: 2.91 -> 2.16 | Helped |
| Use LEGO-PR instead of DiT | CelebA unconditional generation · CelebA 64×64 | FID: 2.52 -> 2.09 | Helped |
| Use LEGO-XL-U instead of MaskDiT | ImageNet generation without classifier-free guidance · ImageNet 256×256 | FID: 5.69 -> 5.28 | Helped |
| Use LEGO-XL-PG instead of MDT-XL/2 | ImageNet generation with classifier-free guidance · ImageNet 256×256 | FID: 1.79 -> 2.05 | Hurt |
| Use LEGO-XL-PG instead of DiT-XL/2 | ImageNet generation without classifier-free guidance · ImageNet 512×512 | FID: 12.03 -> 10.06 | Helped |
| Use LEGO-XL-PG instead of DiT-XL/2 | ImageNet generation with classifier-free guidance · ImageNet 512×512 | FID: 3.04 -> 3.74 | Hurt |
| Combine progressive refinement and progressive growth in LEGO-U (ablation) | LEGO-XL · ImageNet 256×256 | IS: LEGO-XL-PG 289.12; LEGO-XL-PR 284.73; LEGO-XL-U 338.08 | Helped |
| Add a LEGO-PR refinement brick to the pretrained VP model while freezing VP | Pretrained VP diffusion model · CIFAR-10 | FID: 2.5 -> 1.92 | Helped |
| Add a LEGO-PR refinement brick and fine-tune the pretrained VP model | Pretrained VP diffusion model · CIFAR-10 | FID: 2.5 -> 1.88 | Helped |
| Skip the top-level image-brick at low-noise timesteps (ablation) | LEGO-PG sampling · ImageNet 256×256 | sampling time | Helped |
| Skip the top-level image-brick during high-noise timesteps (ablation) | LEGO-PG sampling · ImageNet 256×256 | FID | Hurt |
| Skip the top-level patch-brick at high-noise timesteps (ablation) | LEGO-PR sampling · ImageNet 256×256 | FID | Hurt |
Surprising findings
- LEGO-XL-U achieves the highest IS (338.08 with guidance) among compared models on ImageNet 256x256, while its FID is worse (2.59 with guidance vs. 2.05 for LEGO-XL-PG).
Limitations
- Only explored progressive-growth, progressive-refinement, and U-shaped stacking; brick-size combinations and skipping schedules are heuristic.
- Class-conditional generation at much higher resolution than training can yield low-photorealism results for low-diversity categories.
- Training the 512x512 model required 32 A100 GPUs for about 14 days, so resource requirements remain high.
Experimental setup
Datasets: ImageNet CelebA CIFAR-10
Benchmarks: CelebA 64x64 ImageNet 64x64 ImageNet 256x256 ImageNet 512x512 CIFAR-10
Metrics: FID sFID IS precision recall FLOPs sampling time
Baselines: DiT-XL/2 U-ViT-H/2 MDT-XL/2 MaskDiT ADM LDM-4 LDM-8 U-ViT-S/4 U-ViT-L/4 DiT-S/2 DiT-L/2
Builds on: DiT EDM iDDPM Stable Diffusion VAE
Closest papers by meaning
- Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models (CVPR 2025 · similarity 0.88)
- Plug-and-Play Diffusion Distillation (CVPR 2024 · similarity 0.88)
- Revisiting Diffusion Models: From Generative Pre-training to One-Step Generation (ICML 2025 · similarity 0.88)
- Scaling Diffusion Transformers Efficiently via $\mu$P (NeurIPS 2025 · similarity 0.87)
- ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation (NeurIPS 2025 · similarity 0.87)
- Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.