Direct pixel-space megapixel image generation with diffusion models(crowsonkb.github.io) |
Direct pixel-space megapixel image generation with diffusion models(crowsonkb.github.io) |
activations are small! so you can enjoy bigger batch sizes. this is due to the 4x patching we do on the ingress to the model, and the effectiveness of neighbourhood attention in joining patches at the seams.
the model's inductive biases are pretty different than (for example) a convolutional UNet's. the innermost levels seem to train easily, so images can have good global coherence early in training.
there's no convolutions! so you don't need to worry about artifacts stemming from convolution padding, or having canvas edge padding artifacts leak an implicit position bias.
we can finally see what high-resolution diffusion outputs look like _without_ latents! personally I think current latent VAEs don't _really_ achieve the high resolutions they claim (otherwise fine details like text would survive a VAE roundtrip faithfully); it's common to see latent diffusion outputs with smudgy skin or blurry fur. what I'd like to see in the future of latent diffusion is to listen to the Emu paper and use more channels, or a less ambitious upsample.
it's a transformer! so we can try applying to it everything we know about transformers, like sigma reparameterisation or multimodality. some tricks like masked training will require extra support in [NATTEN](https://github.com/SHI-Labs/NATTEN), but we're very happy with its featureset and performance so far.
but honestly I'm most excited about the efficiency. there's too little work on making pretraining possible at GPU-poor scale. so I was very happy to see HDiT could succeed at small-scale tasks within the resources I had at home (you can get nice oxford flowers samples at 256x256px with half an hour on a 4090). I think with models that are better fits for the problem, perhaps we can get good results with smaller models. and I'd like to see big tech go that direction too!
-Alex Birch
regarding IPAdapter: I'm not familiar with it, but from the code it looks like they just run cross-attention again and sum the two attention outputs. feels a bit weird to me, because the attention probabilities add up to 2 instead of 1. and they scale the bonus attention output only instead of lerping. it'd make more sense to me to formulate it as a cross-cross attention (Q against cat([key0, key1]) and cat([val0, val1])), but maybe they wanted it to begin as a no-op at the start of training or something. anyway.. yes, all of that should work fine with HDiT. the paper doesn't implement cross-attention, but it can be added in the standard way (e.g. like stable-diffusion) or as self-cross attention (e.g. DeepFloyd IF or Imagen).
I'd recommend though to make use of HDiT's mapping network. in our attention blocks, the input gets AdaNormed against the condition from the mapping network. this is currently used to convey stuff like class conditions, Karras augmentation conditions and timestep embeddings. but it supports conditioning on custom (single-token) conditions of your choosing. so you could use this to condition on an image embed (this would give you the same image-conditioning control as IPAdapter but via a simpler mechanism).
I see your headline speed comparison is to "Pixel-space DiT-B/4" - but how does your model compare to the likes of SDXL? I gather they spent $$$$$$ on training etc, so I'd understand if direct comparisons don't make sense.
And do you have any results on things that are traditionally challenging for generative AI, like clocks and mirrors?
Thanks for your work - it’s been impactful since the early days of the project.
Excited to see where we get to this year.
this year I'm hoping for efficiency and small models! even if it's proprietary. if our work can reduce some energy usage behind closed doors that'd still be a good outcome.
The popular large scale diffusion models like StableDiffusion are CNN based at their heart, with attention layers sprinkled throughout. This paper builds on recent research exploring whether competitive image diffusion models can be built out of purely transformers, no CNN layers.
In this paper they build a similar U-Net like structure, but out of transformer layers, to improve efficiency compared to a straight Transformer. They also use local attention when the resolution is high to save on computational cost, but regular global attention in the middle to maintain global coherence.
Based on ablation studies this allows them to maintain or slightly improve FID score compared to Transformer-only diffusion models that don't do U-net like structures, but at 1/10th the computation cost. An incredible feat for sure.
There is a variety of details: RoPE positional encoding, GEGELU activations, RMSNorm, learnable skip connections, learnable cosine-sim attention, neighborhood attention for the local attention, etc.
The biggest gains in FID occur when the authors use "soft-min-snr" as the loss function; FID drops from 41 to 28!
Lots of ablation study was done across all their changes (see Table 1).
Training is otherwise completely standard AdamW, 5e-4, 0.01, 256 batch, constant LR, 400k steps for most experiments at 128x128 resolution.
So yeah, overall seems like solid work that combines a great mixture of techniques and pushes Transformer based diffusion forward.
If scaled up I'm not sure it would be "revolutionary" in terms of FID compared to SDXL or DALLE3, mostly because SD and DALLE already use attention obviating the scaling issue, and lots of tricks like diffusion based VAEs. But it's likely to provide a nice incremental improvement in FID, since in general Transformers perform better than CNNs unless the CNNs are _heavily_ tuned.
And being pixel based rather than latent based has many advantages.
it also doesn't punish poor detail either! this advantages latent diffusion, which can claim to achieve a high resolution but without actually needing to have correct textures to get good metrics.
It's not entirely clear from the comparison numbers at the end, but I think the big argument here is efficiency for the amount of performance achieved. One can get lower FID numbers, but also with a ton of compute.
I can't really speak technically to it as I've not given it a super in depth look, but this seems like a nice set of motifs for going halfway between a standard attention network and a convnet in terms of compute cost (and maybe performance)?
The large-resolution scaling seems to be a strong suit as a result. :)
Yeah, the main motivation was trying to find a way to enable transformers to do high-resolution image synthesis: transformers are known to scale well to extreme, multi-billion parameter scales and typically offer superior coherency & composition in image generation, but current architectures are too expensive to train at scale for high-resolution inputs.
By using a hierarchical architecture and local attention at high-resolution scales (but retaining global attention at low-resolution scales), it becomes viable to apply transformers at these scales. Additionally, this architecture can now directly be trained on megapixel-scale inputs and generate high-quality results without having to progressively grow the resolution over the training or applying other "tricks" typically needed to make models at these resolutions work well.
A few people have finetuned Stable Diffusion models on v-objective and solved the problem from the root.
I tried adding "stable-diffusion-style" cross-attn to HDiT, text-conditioning on small class-conditional datasets (oxford flowers), embedding the class labels as text prompts with Phi-1.5. trained it for a few minutes, and the images were relevant to the prompts, so it seemed to be working fine.
but if instead of a text condition you have a single-token condition (class label) then yeah the adanorm would be a simpler way.
>Does class-conditional mean that it takes a class like "car" or "face" and generate a new random image of that class?
Yup
nope, we don't do Imagen-style super-resolution. we go direct to high resolution with a single-stage model.
the ImageNet-256 examples also aren't typically blurry (but they are 256x256 so your viewer may be bicubic scaling them or something). the ImageNet dataset _can_ have blurry, compressed or low resolution training samples, which can afflict some classes more than others, and we learn to produce samples like the training set.
[0] https://crowsonkb.github.io/hourglass-diffusion-transformers...
Energy efficiency is key - Especially with some of these extremely inefficient (wasteful, even) features like real-time canvas.
Good luck - Let us know if/how we can help.
But generally, most other UIs support it. It has serious limitations though, for example it center-crops the input to 224x224px. (which is enough for a surprisingly large amount of uses, but not enough for many others)
A latent space CLIPVision model would be cool too. Presumably you could leverage the semantic richness of the latent space to efficiently train a more powerful CLIPVision. I don’t know whether anyone has tried this. Maybe there is a good reason for that.
The overall architecture diagram does not explicitly show the conditioning mechanism, which is a small separate network. For this paper, we only trained on class-conditional ImageNet and completely unconditional megapixel-scale FFHQ.
Training large-scale text-to-image models with this architecture is something we have not yet attempted, although there's no indication that this shouldn't work with a few tweaks.
Can this architecture be used to distill models that need fewer timesteps like LCMs or SDXL turbo?