[P] Consistency: Diffusion in a Single Forward Pass 🚀

Submitted by Beautiful-Gur-9456 t3_124jfoa in MachineLearning

Hey all!

Recently, researchers from OpenAI proposed consistency models, a new family of generative models. It allows us to generate high quality images in a single forward pass, just like good-old GANs and VAEs.

​

training progress on cifar10

​

I have been working on it and found it definetly works! You can try it with diffusers.

import diffusers

from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "consistency/cifar10-32-demo",
    custom_pipeline="consistency/pipeline",
)

pipeline().images[0]  # Super Fast Generation! 🤯
pipeline(steps=5).images[0]  # More steps for sample quality

It would be fascinating if we could train these models on different datasets and share our results and ideas! 🤗 So, I've made a simple library called consistency that makes it easy to train your own consistency models and publish them. You can check it out here:

https://github.com/junhsss/consistency-models

I would appreciate any feedback you could provide!

61

Comments

You must log in or register to comment.

noraizon t1_je10328 wrote

x0-parametrization has been used for some time now. imo, nothing new under the sun. maybe it's something else I don't see

5

geekfolk t1_je23p7c wrote

How is it better than GANs though? or in other words, what's so bad about adversarial training? modern GANs (with zero centered gradient penalties) are pretty easy to train.

1

Beautiful-Gur-9456 OP t1_je3hxbn wrote

Replying to geekfolk (#2,428,433)

The training pipeline, honestly, is significantly simpler without adversarial training, so the design space is much smaller.

It's actually reminiscent of GANs since it uses pre-trained networks as a loss function to improve the quality, though it's completely optional. Still, it's a lot easier than trying to solve any kind of minimax problem.

2

geekfolk t1_je3io3b wrote

Replying to Beautiful-Gur-9456 (#2,438,096)

using pretrained models is kind of cheating, some GANs use this trick too (projected GANs). But as a standalone model, it does not seem to work as well as SOTA GANs (judged by the numbers in the paper)

​

>Still, it's a lot easier than trying to solve any kind of minimax problem.

This is true for GANs in the early days; however, modern GANs are proved to not have mode collapse and the training is proved to converge.

>It's actually reminiscent of GANs since it uses pre-trained networks

I assume you mean distilling a diffusion model in the paper. There have been some attempts to combine diffusion and GANs to get the best of both worlds but afaik none involved distillation, I'm curious if anyone has tried distilling diffusion models into GANs.

0

hebweb t1_je3ofr1 wrote

Cool! this is amazing. You already created a pip package out of it. Have you measured the fid on your model? Does it match the numbers of the paper? I think their batch size and model size were pretty large even for the CIFAR10 training. Not sure if we can match that..

1

huehue9812 t1_je3q0pl wrote

Replying to geekfolk (#2,428,433)

Hey, can I ask something about 0-GP GANs? This is the first time I've ever heard of them. I was wondering what makes them superior over R1 regularization. Also, why is it that most papers mention R1 reg., but not 0-GP?

1

Beautiful-Gur-9456 OP t1_je3qsdu wrote

Replying to geekfolk (#2,438,235)

Nope. I mean the LPIPS loss, which kinda acts like a discriminator in GANs. We can replace it to MSE without much degradation.

Distilling SOTA diffusion model is obviously cheating 😂, so I didn't even think of it. In my view, they are just apples and oranges. We can augment diffusion models with GANs and vice versa to get the most out of them, but what's the point? That would make things way more complex. It's clear that diffusion models cannot beat SOTA GANs for one-step generation; they've been tailored for that particular task for years. But we're just exploring possibilities, right?

Aside from the complexity, I think it's worth a shot to replace LPIPS loss and adversarially train it as a discriminator. Using pre-trained VGG is cheating anyway. That would be an interesting direction to see!

2

geekfolk t1_je3qyfr wrote

Replying to Username912773 (#2,439,251)

I don’t know about this model, but GANs are typically smaller than diffusion models in terms of num of params. The image structure thing probably has something to do with the network architecture since GANs rarely use attention blocks and the network architecture of diffusion models is more hybrid (typically CNN + attention)

1

Beautiful-Gur-9456 OP t1_je3uesm wrote

Replying to hebweb (#2,439,099)

I haven't done it yet, but I'm working on it! Their suggested sampling procedure requires multiple FID calculation, so I'm thinking of how to incorparate it efficiently.

Their scale is indeed large, it would cost me a few hundread bucks to train CIFAR10. My checkpoint was trained with much smaller size 😆

1