VaeZip: store VAE-encoded latent datasets in 5-10x fewer space with almost no quality loss

Hi! I sometimes train models and one thing always seemed very suboptimal for me. This is VAE latent files that take much more space than original data, while VAE should compress data, not making it more heavy. This is mostly because jpeg and video codecs do there job and compress data very well, while latents are almost uncompressed. So I tried to make something like jpeg but for latents. I want to give a heads-up right away that this thing have very narrow use case when your latent dataset is large enough that it doesn't fit your local disk, so you have to use some s3 storage which may cost a lot when the dataset is so large. Anyway, VaeZip is a PyTorch-based library that implements several latent quantization codecs and stores the compressed tensors in common Safetensors files that are much smaller then if you store these tensors in fp32. The codecs are designed specifically for VAE latents: some of them exploit both spatial redundancy and that the VAE latent distribution is usually normal. That allows to store latent tensor with about 3 bits/value precision (10x smaller files) while keeping pixel MSE at 0.0003 after decoding. For example, one of the best codecs uses a combination of global range scaling, local range scaling with square spatial blocks and so called “Lloyd-Max table” which is essentially codebook that is tuned for the sample distribution. After that it quantizes values to 2-8 bits and packs them. The codecs were selected from the benchmark that compared ~50 algorithms on 8 image and video VAEs by the reconstruction quality, fidelity, space usage and I/O latency. I have also measured the impact on pretraining a small T2I flow matching DiT with latents compressed with different codecs. In one run the loss curve is indistinguishable from the fp32 baseline with just 6 bit storage. Maybe some of you here do some large scale pretraining or tuning, hope you find this useful! submitted by /u/TensorForger [link] [comments]

Source: r/StableDiffusion | 2026-10-16