8 ms·
The VAE Used for Stable Diffusion Is Flawed
- october8140 3y ago[flagged]
- prashp 3y agoThe image compression model used in SD is not evenly (in a spatial sense) compressing the images
- nness 3y agoCould someone ELI5? What is the impact of this issue?
- gchadwick 3y agoStable diffusion (along with other text to image models like Dall-E) use a process called 'latent diffusion'. At the core of a latent diffusion model is a de-noising process. It takes a noisy image and predicts what is noise vs what is the real image without noise. You use this to remove a bit of noise from the image and repeat to iteratively denoise an image. You can use this to generate entirely new images by just starting with complete random noise and denoising til you get a 'proper' image. Obviously this would not give you any control over what you generated. So you incorporate 'guidance' which controls how the denoise works. For stable diffusion this guidance comes from a different neural network called CLIP (https://openai.com/research/clip https://openai.com/research/clip) which can take some text and produce a numerical representation of it that can be correlated to an image of what the text describes (I won't go into more detail here as it's not really relevant to the VAE). The problem you have with the denoising process is the larger the image you want to denoise the bigger the model you need, and even at a modest 512x512 (the native resolution of stable diffusion) training the model is far too expensive. This is where the latent bit comes in. Rather than train your model on a 512x512x3 representation (3 channels R,G,B per pixel) use a compressed representation that is 64x64x4, significantly smaller than the uncompressed image and thus requiring a significantly smaller denoising model. This 64x64x4 representation is known as the 'latent' and it is said to be in a 'latent space'. How do we produce the latent representation? A VAE, a variational autoencoder, yet another neural network. You train an encoder and decoder together to encode an image to the 64x64x4 space and decode it back to 512x512x3 with as little loss as possible. The issue pointed out here is the VAE for stable diffusion has a flaw, it seems to put global information in a particular point of the image (to a crude approximation it might store information like 'green is the dominant colour of this image' in that point). So if you touch that point in the latent you effect the entire image. This is bad because the denoising network is constructed in such a way that it expects that points close in the latent only effect other points close in the latent. When that's not the case it ends up 'wasting' a bunch of the network on extracting that global data from that point and fanning it out to the rest of the image (as the entire image needs to know it to denoise correctly). So without this flaw it may be the stable diffusion denoising model could be more effective as it doesn't need to work hard to work around the flaw. Edit: Pressed enter too early, post is now complete.
- ants_everywhere 3y ago> 64x64x4 I'm curious why 4? Is this just what works in practice, or do the 4 channels have known interpretations?
- Sharlin 3y agoI'm not sure why 4 was chosen (maybe just because it's a power of two?) but a while ago a SD user found that RGB can be approximated fairly well with a simple linear combination of the latents: https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/2 https://discuss.huggingface.co/t/decoding-latents-to-rgb-wit...
- matthewsinclair 3y agoThat’s a great explanation. Thanks!
- phkahler 3y agoWhy does denoising require global information? Shouldn't it be mostly local? I would expect individual pixels to have very little global information and a lot of information from neighboring pixels, and nothing of a local scale from other areas of the image. Isn't this the case?
- Robin_Message 3y agoIf the latent space is meant to be highly spatially correlated, could you simply apply random rotations of rows and columns† to the latent space throughout the process? That way, there wouldn't be specific areas data could be smuggled through. † As in, move a random vertical stripe of the image from the right to the left, and a random horizontal portion from the top to the bottom. Or, if that introduces unacceptable edge effects, simply slice the space into 4 randomly sized spaces (although that might encourage smuggling in all of the corners at once.)
- wokwokwok 3y ago> It's a spot where the VAE is trying to smuggle global information about the image through latent space. This is exactly the problem that KL-divergence loss is supposed to prevent. Is that what KL divergence does? I thought it was supposed to (when combined with reconstruction loss) “smooth” the latent space out so that you could interpolate over it. Doesn’t increasing the weight of the KL term just result in random output in the latent; eg. What you get if you opt purely for KL divergence? I honestly have no idea at all what the OP has found or what it means, but it doesnt seem that surprising that modifying the latent results in global changes in the output. Is manually editing latents a thing? Surely you would interpolate from another latent…? And if the result is chaos, you dont have well clustered latents? (Which is what happens from too much KL, not too little right?) I'd feel a lot more 'across' this if the OP had demonstrated it on a trivial MNIST vae with both the issue, the result and quantitatively what fixing it does. > What are the implications? > Somewhat subtle, but significant. Mm. I have to say I don't really get it.
- a-dub 3y agoedit: nevermind, i need to read up on these.
- ants_everywhere 3y agoI can't comment on what changing the weights of the KL divergence does in this context, but generally > Is that what KL divergence does? KL divergence is basically a distance "metric" in the space of probability distributions. If you have two probability distributions A and B, you can ask how similar they are. "Metric" is in scare quotes because you can't actually get a distance function in the usual sense. For example, dist(A,B) != dist(B,A). If you think about the distribution as giving information about things, then the distance function should say two things are close if they provide similar information and are distant if one provides more information about something than the other. The comment claims (and I assume they know what they're talking about) that after training we want the KL divergence to be close to a standard Gaussian. So that would mean that our statistical distribution gives roughly the same information as a standard Gaussian. It sounds like this distribution has a whole lot of information in one heavily localized area though (or maybe too little information in that area, I'm not sure which way it goes).
- ants_everywhere 3y agoI'm curious, was this well-known by experts already? How surprising is this? I enjoyed the write up.
- GaggiX 3y agoI have never heard of this problem before, and I have seen a lot of discussion about VAE from researchers.
- dwringer 3y agoIf one ever tried to make edits to the latents prior to decoding them with a VAE in SD1.5 and then in SDXL, it could be seen that that local changes had somewhat unpredictable and global effects on the image in SD1.5, while in SDXL the changes have more predictable impacts to the output image and some of the different latent channels end up corresponding more directly to the resulting image channels. Definitely a fascinating write-up. I have been curious about these differences for a while, though I had never considered this a "problem" per se.
- numpad0 3y agoI've once seen someone on Twitter wondering about to-them-obviously-bug with VAE leading to oddly saturated images in anime space, just my dumb brain keyword search though
- joefourier 3y agoI’ve done a lot of experiments with latent diffusion and also discovered a few flaws in the SD VAE’s training and architecture, which have hardly no attention brought to them. This is concerning as the VAE is a crucial competent when it comes to image quality and is responsible for many of the artefacts associated with AI generated imagery, and no amount of training the diffusion model will fix them. A few I’ve seen are: - The goal should be to have latent outputs as closely resemble gaussian distributed terms between -1 and 1 with a variance of 1, but the outputs are unbounded (you could easily clamp or apply tanh to force them to be between -1 and 1), and the KL loss weight is too low, hence why the latents are weighed by a magic number to more closely fit the -1 to 1 range before being invested by the diffusion model. - To decrease the computational load of the diffusion model, you should reduce the spatial dimensions of the input - having a low number of channels is irrelevant. The SD VAE turns each 8x8x3 block into a 1x1x4 block when it could be turning it into a 1x1x8 (or even higher) block and preserve much more detail at basically 0 computational cost, since the first operation the diffusion model does is apply a convolution to greatly increase the number of channels. - The discriminator is based on a tiny PatchGAN, which is an ancient model by modern standards. You can have much better results by applying some of the GAN research of the last few years, or of course using a diffusion decoder which is then distilled either with consistency or adversarial distillation. - KL divergence in general is not even the most optimal way to achieve the goals of a latent diffusion model’s VAE, which is to decrease the spatial dimensions of the input images and have a latent space that’s robust to noise and local perturbations. I’ve had better results with a vanilla AE, clamping the outputs, having a variance loss term and applying various perturbations to the latents before they are ingested by the decoder.
- parlancex 3y agoEverything you've said is _intuitively_ correct, but empirically wrong. I've experimented with training VAEs for audio diffusion for the last few months and here's what I found: - Although the best results for a stand-alone VAE might require increasing the KL loss weight as high as you can to reach an isotropic gaussian latent space without compromising reconstruction quality, beyond a certain point this actually substantially decreases the ability of the diffusion model to properly interpret the latent space and degrades generation quality. The motivation behind constraining the KL loss weight is to ensure the VAE only provides _perceptual_ compression, which VAEs are quite good at, not _semantic_ compression, for which VAEs are a poor generative model compared to diffusion. This is explained in the original latent diffusion paper on which Stable Diffusion was based: https://arxiv.org/pdf/2112.10752.pdf https://arxiv.org/pdf/2112.10752.pdf - You're correct that trading dimensions for channels is a very easy way to increase reconstruction quality of a stand-alone VAE, but it is a very poor choice when the latents are going into a diffusion model. This again makes the latent space harder for the diffusion model to interpret, and again isn't needed if the VAE is strictly operating in the perceptual compression regime as opposed to the semantic compression regime. The underlying reason is channel-wise degrees of freedom have no inherent structure imposed by the underlying convolutional network; in the limit where you hypothetically compress dimensions to a single point with a large number of channels the latent space is completely unstructured and the entropy of the latents is fully maximized; there are no patterns left whatsoever for the diffusion model to work with. TLDR: Designing VAEs for latent diffusion has a different set of design constraints than designing a VAE as a stand-alone generative model.
- prashp 3y agoEmad (StabilityAI founder) posted on the reddit thread: "Nice post, you'd be surprised at the number of errors like this that pop up and persist. This is one reason we have multiple teams working on stuff.. But you still get them"
- GaggiX 3y agoAnother example is when people realized that SD v1.5 wasn't able to generate images that were too dark or too bright. The problem in the end was that during training even the noisiest step still has enough signal for the model to be able to detect the mean of the actual image (signal), this is done because you cannot have pure Gaussian noise during training of an epsilon objective model or it will cause a division by zero. Of course during inference there is no signal in the first step, so the model would read the mean of the input (so zero as the input is Gaussian noise) and it will output an image of mean 0. It's not uncommon to find major problems with these systems, I remember inspecting the VQGAN used by Dalle Mega (the largest version of Dalle Mini) and discovering that the vast majority of entries in the codebook had a magnitude very close to zero, making them completely unusable by the model.
- msp26 3y ago> and I would also like to thank the Glaze Team, because I accidentally discovered this while analyzing latent images perturbed by Nightshade and wouldn't have found it without them, because I guess nobody else ever had a reason to inspect the log variance of the latent distributions created by the VAE That's just hilarious
- dawnofdusk 3y agoThis is one of the cool things about various neural network architectures that I've found in my own work: you can make a lot of dumb mistakes in coding certain aspects but because the model has so many degrees of freedom it can actually "learn away" your mistakes.
- CamperBob2 3y agoIt's also one of the scariest things about NNs. Traditionally, if you had a bug that was causing serious performance or quality issues, it was a safe bet that you'd eventually discover it and fix it. It would fail one test or another, crash the program, or otherwise come up short against the expectations you'd have for a working implementation. Now it's almost impossible to know if what you've implemented is really performing at its best. IMO the ability for a NN to compensate for bugs and unfounded assumptions in the model isn't a Good Thing in the slightest. Building latent-space diagnostics that can determine whether a network is wasting time working around bugs sounds like a worthwhile research topic in itself (and probably already is.)
- dawnofdusk 3y agoIt is a good thing for deep NNs to be expressive enough to do this, because it is this level of expressivity that let's it find answers to otherwise ill-posed problems. If it were not able to do this there would be no point in using them. The only thing that is scary is the hype, because this will make people sloppily use deep learning architectures for problems that do not need that level of expressive power, and because deep learning is challenging and not theoretically well understood, there will be little to no attempts made to ensure safe operation/quality assurance of the implemented solution.
- nyanpasu64 3y agoTBH I think this is the case for even classical heuristic algorithms.
- nullc 3y agoIt's a common problem for network protocols, IO subsystems, etc. and really even any software with error handling. It's been a few years since I worked on any program using boost asio, but at least back then if you straced it you'd find it constantly attempting to malloc hundreds of TB of ram, failing harmlessly, then continuing on with its life. (bet that will be fun when someone tries to run it on a system that supports that much virtual address space) Similarly anything with any kind of feedback correction. PID controllers, codecs that code residuals-- you can get things horribly wrong and the later steps will paper it over. Taking a step back you can even say that common software development practices-- a kind of meta program-- have the issue: A drunk squirrel sends you a patch full of errors, your test suite flags some which you fix. Then you ship all the bugs you didn't catch, because the test suite caused you to fix some issues but didn't change the fact that you were accepting code from a dubious source. So I would say that the ML world is only special in that they exist almost entirely of self-correcting mechanisms and that inconsistent performance is broadly expected to a vastly greater degree, so when errors leak through you still may not react. If a calculator app told you that 2+2=5 you'd immediately be sure that something is actually broken, while if some LLM does it, it could just be an expected limitation (or even just sampling bad luck).
- treprinum 3y agoThe author might be right though what I've noticed with DL models is that the theory is often leading to underwhelming results after training and "bugs" in models sometimes lead to much better real-world performance, pointing out some disconnect between theory and what gradient-based optimization can achieve. One could see it also in the deep reinforcement learning where in theory the model should converge due to being Markovian via Banach fixed point but in practice the monstrous neural networks that estimate rewards can override this and change the character of the convergence.
- Sharlin 3y agoOne interesting example is how the OG algorithm for solving diff equations, the venerable and almost trivial Euler's method, in fact works very well with SD compared to many much newer, slower and fancier solvers. This likely has to do with the fact that we're dealing with an optimization problem rather than actually trying to find solutions to the diffusion DE.
- stealthcat 3y agoImagine all those millions of $$$ GPU of cloud credits for training, only to overlook this bug
- Traubenfuchs 3y agoFrom a year ago: https://twitter.com/MosaicML/status/1617944401744957443?lang=en# https://twitter.com/MosaicML/status/1617944401744957443?lang... This probably got even better and cheaper.
- urbandw311er 3y agoRelated (coincidentally) — Google also posted research on a much more efficient approach to image generation: https://news.ycombinator.com/item?id=39210458 https://news.ycombinator.com/item?id=39210458
- mmastrac 3y agoThere seems to be a convincing debunking thread on Twitter, but I definitely don't have the chops to evaluate either claim: https://twitter.com/Ethan_smith_20/status/1753062604292198740 https://twitter.com/Ethan_smith_20/status/175306260429219874...
- bjornsing 3y agoOnly took a quick glance. But it looks like a good debunking to me, especially the part where he refers to a section of the original paper clearly stating that the VAE has been trained with a 10e-6 weight factor on the KL divergence term. I think what would happen if this problem was fixed is that the VAE would produce less appealing more blurry images. This is a classic problem with VAEs. So, more mathematically correct, but less visually appealing.
- joefourier 3y agoNot necessarily if the model is trained with an appropriate adversarial loss. The reason that VAEs are blurry isn’t directly because of the KL divergence loss term but because of the L1/L2 loss. Since VAEs sample from a Gaussian distribution, a high KL weight will make the latent of different images overlap (towards a -1 to 1.0 with a variance of 1 and a mean of 0), and the output of the decoder will tend towards the mean of possible values to try and minimise the pixel-wise loss. With the appropriate GAN loss, you will instead get a plausible sharp image that differs more and more from the original the more you weigh the KL loss term. A classic GAN that samples from the normal distribution in fact has the best possible KL divergence loss and none of the blurriness from a VAE’s pixel based loss.
- probably_wrong 3y agoLink for those without Twitter: https://threadreaderapp.com/thread/1753062604292198740.html https://threadreaderapp.com/thread/1753062604292198740.html
- smusamashah 3y agoReddit thread by the same person https://www.reddit.com/r/StableDiffusion/comments/1agd5pz/a_recent_post_went_viral_claiming_that_the_vae_is/?sort=new https://www.reddit.com/r/StableDiffusion/comments/1agd5pz/a_...
- mzs 3y agodetailed counter argument: https://twitter.com/Ethan_smith_20/status/1753062604292198740 https://twitter.com/Ethan_smith_20/status/175306260429219874...
- deleted 3y ago[deleted]
- bornfreddy 3y agoOff-topic: is anyone aware of good tutorials / howtos / books / videos on the NN developments of last few years? (Attention, SD, LLM,...)
- nerdponx 3y agoDive Into Deep Learning, free online at https://d2l.ai/ https://d2l.ai/ (and soon to be published in dead-tree format).
- bornfreddy 3y agoThank you!
- firechickenbird 3y agoBackprop doesn’t give a shit
- itsTyrion 3y agoso many kWh wasted
- brcmthrowaway 3y agoThese people talk like high school kids with too much time on their hands to be honest. Have we created a whole new class of practitioners hacking away with no theory background?