Generative modelling of compressed image file bits

The holy grail

Tales have long told of a great and mythical image-text architecture. Stories passed down through the generations called it the holy grail. Parents, after a long hard day of prompting, would come home to their children and speak of a day when their tedious prompting would be no more. It would be the holy grail and would generate images using the direct bits/bytes of a compressed file format.

Most called it a fools errand but many were desperate to search for the grail. Despite much research [note] some found that LLMs can use VQ-coded images and use bidirectional attention as well as a parallel sampling scheme for image tokens. (Chamel3on and Effective Long-Context Scaling of Foundation Models)
others discovered that llms could input and output CLIP embeddings (Grounding Language Models to Images for Multimodal Inputs and Outputs) to this day the holy grain remains out of reach. Whether it was to do with anti-research or proprietary datasets, none could tell.

Images and pixels

Images have two dimensions. They have a height and a width. At each position in a 2D grid is a pixel which has 3 bytes that are the pixel's colors. The pixels themselves are captured by a piece of photosensitive material which samples light rebounding off of objects in the real world. The real world is where physics happens and makes light bounce around from stuff going on.

So there are two worlds: pixel world and the real world. Pixel world, though being just a sample of the real world, can nonetheless still communicate complicated information about the real world. Sampling down to 2D doesn't rid all of the intracacies of the real world. Images still have to portray scenes that maintain coherent portrayal objects moving through time and 3D space. Moreover, images aren't restricted to physical 3D objects, they also can show written text and numbers. Imagine an entity that lives in 2D pixel world and can only ever see and interact in a 2D plane of existance. Would this entity have enough information to figure out all of the underlying facts of the real world that their 2D existence is sampled from? [note] This is an idea explored in the Three-Body-Problem sci-fi series. Also this forum has some interesting thoughts on this. It turns out that in 2d pixel world life isn't really any more simple than in real world.

Images being virtually unrestricted in their expressive power is bad news for our generative models. If we train a model on pixels and we don't specify any other restrictions, we are in effect expecting the model to not only figure out snapshots of the natural world, but also understand all bodies of written text and mathematical proofs and also physics and geometry and psychology and biology. That is why I think people are so eager for an all-in-one holy grail model that does everything. Because in the skillsets demanded of LLMs and image gen is the same.

It used to be the case that LMs couldn't generate text with proper grammar and a whole bunch of research was done on decoding within the bounds of correct grammatical sytax. Well, nowdays our image models can't generate within the bounds of physics. I think physics is more complicated than grammar. Why not force our image models to generate within some sort of constrained system?

We aren't designing handmade generation constraints and features these days, or at not as much as we used to. Models get trained on data that is as close to the pure input and output modalities as possible, end-to-end. For images this means means no more SIFT features or handmade edge kernels. [note] Lex Fridman's interview with Jitendra Malik is very interesting, here he talks about the transition from the 'old-style' computer vision to the more end-to-end approach This avoids fettering inputs and outputs with any prior assumptions about structure. In image generation this idea takes shape in the form of training on raw pixels. In video generation, this means using pixels from videos frames. For 3D data, this means something like implicit neural representations. I'd say the vast majority of the recent sucessful generative models use this end-to-end approach. [note] LLMs which use tokenizers which are arguably 'hand-picked' features. The less restrictive approach would be training on the direct characters/bytes.

There are still some modelling strategies being used that are less expressive but impose more structure on the inputs and outputs. The model is not given the 'raw' data, it is given a limited version. You're still expecting your model to perform well on the raw data but you only are giving it some gimped version of the data, which is not ideal. If given a choice, most researchers would prefer to use the end to end, but are forced to make concessions for stability or performance reasons. Some examples of this for 3D data are techniques that don't use nerfs Octree transformer where 3D data is generated by autoregressively decoding Octrees, and Mesh GPT for autoregressive generation of vector quantized codings of vertices. For images and sound, the recent paper Bytes Are All You Need: Transformers Operating Directly On File Bytes showed high quality image and sound classification done using the raw bytes of encoded files. These strategies tie in with the idea of constrained decoding. It's similar to generating structured data (like json) with an LLM. Octree-transformer and Mesh GPT during generation have deterministic code than runs in step with the model, allowing and disallowing actions.

To summarize, there's the idea that we divvy our creations up into two distict worlds. The real world, where all the complicated signals are, and the world that programs live in. There is a massive chasm between realness and virtualness. The trend seems to be that machine learning needs to learn from the real world. But I'd argue that the digital world holds a lot, probably enough. I think that it shouldn't matter to train image models on full .RAW images and decode uncompressed pixels and videos. The algorithms that provided the media that fueled the digital boom: html,jpg,mp3, are all handmade and intricate and beautiful. Why should we suddenly desert them? By the 90s we had given up on all of the raw data formats. These encoding algorithms are suboptimal, but they do a fantastic job at linking the real world and the digital world with minimal overhead, and they do it more than well enough for human perception. It shouldn't matter that our models aren't given the real deal.

An image as a left-to-right sequence of tokens

Autoregressive decoder-only token-based transformers work pretty well. They have an objective that matches well with the idea of max entropy, and they scale easily, and can usually generate samples matching their data distribution without many tweaks. For images, the landscape of autoregressive generation is a bit desolate. Granted, there have been advancements in discrete modelling of images as a 2D array of tokens. This is not the same as autoregressive generation though, tokens are sampled from a 2D grid of probabilities. [note] masked tokenized image generation. see: magvit, vqvae, taming transformers, Meta's Chamel3on, Fewer attempts are made at left-to-right autoregressive conditional modelling. [note] Using residual VQ you can sort of do part masked image sampling and part autoregressive modelling, see Autoregressive Image Generation using Residual Quantization

Previous works have flattened images in a variety of creative ways. The simplest way to flatten an image into a 1D sequence is to use raster ordering of all pixels. This is what is experimented with in the classic pixel rnn. MEGABYTE used a slightly fancier patch scan order.

Conceptually, raster pixel ordering and even patch scan ordering is a very bad way to treat an image as a sequence of tokens. Image asking a painter to create an art piece by starting in the top left corner, and from there only allowing them to make one brush stroke at every position while moving left-right, top-down in the painting. It would make everything much more difficult for the painter. A more natural way to add detail to a painting over time is to first start with a rough background, then create the overall composition of colors, then add particular strong edges and shapes, and then move in and add fine details to certain areas.

Bob Ross and other artists do not make paintings by placing brush strokes in a raster scan ordering.

We want to turn an image into a sequence of actions that create the image, in a way that is easy to explain and follow. In a way, the image itself is a place for planning. The first actions are noncommittal but set the general stage for what the image will finally look like. Human artists have a sense of what to start with when making a painting. Unless you have a dataset of millions of brush stroke histories we can't really mimic this sequential behavior. [note] Looking at you adobe But I think that image compression algorithms are expert demonstrators of step-by-step image synthesis.

Tokenizing an image can be interpreted as translating an image into a sequence of expert instructions. The instructions have to be 'intuitive', and when followed, actually result in the creation of something close to the original image. When we train our LLM to produce the image, we want it to generally mimic these expert instructions.

Image compression algorithm

There are lots of image compression algorithms, but we want one with certain properties. Firstly, the bit-stream should have a left-to-right bias. Secondly, it would be nice if bit stream could be able to be interrupted at any point and decoded into an image. This would let us have previews during the slow autoregressive generation. Also, we want the coding to be very resiliant to noisy tokens. Perhaps most importantly, there should be a way for the decoder/encoder of the compression algorithm to provide extra conditioning information to the autoregressive model about what is going on. We want this property so it removes the burden of the autoregressive model needing to reverse engineer in its weights the compression algorithm.

We want to obtain a stream of tokens that describes a sequence of actions to construct the image. We also want to get conditioning tokens. Conditioning tokens aren't used directly by the encoder/decoder, but are given as additional information for the downstream autoregressive generative models. After the autoregressive model generates each token, it feeds the new sequence to the decoder and gets an additional conditioning vector. It is difficult to train both the image compression model and the autoregressive model at the same time and pass gradients between them. How do we know what information will be useful for the generative model? In this project I use handmade features.

SPIHT (pdf here) is an old image compression algorithm from the 90s that satisfies many of our requirements. First, it uses the DWT transform to decorrelate RGB pixel values. The DWT decomposition of an image makes a 2D grid of sparse coefficients. They are placed in a way that the coefficients that aren't close to zero are arranged in predictable clusters. Importantly, MSE distance is preserved in the DWT.

The goal of the SPIHT encoder is to output a sequence of bits from most importance to least importance (where importance is measured by MSE between the output image and the current encoded representation). DWT coefficients can be arranged in a tree structure. SPIHT does a kind of BFS of this tree in an ordering of most-to-least significance in the bit plane.

I think SPIHT is better for autoregressive modelling than JPG or JPGXL or HEIC or png. At any point, the bitstream can be interrupted and the full image decoded. Each bit has a direct relationship to a single coefficient (or it's decendents). Each bit is explainable, and clearly attributable to pixels in the image. It does not use entropy encoding; one bit is one action [note] although entropy coding isn't something that would stop us from training an autoregressive model. Entropy coding reminds me a lot of tokenization; how to find an efficient sub-word vocab. . Each bit has a simple and explainable role in the resulting final image. So coming up with a way to produce SPIHT specific token conditioning information is not too difficult.

Assisted SPIHT decoding - Conditioning tokens

It's possible to take an AR model, and train it directly on SPIHT bits where each bit is a token. This would be wasteful because it completely discards all advantage of how explainable each bit of the SPIHT encoding is. By default, by just training on the bits themselves, we're basically expecting the NN model to fully memorize how the decoding algorithm works and the current variables at each step of the SPIHT decoder/encoder. SPIHT is close to being a tree decoding algorithm. You can 'tell' the NN where it is in the tree at each next token. We can infuse our input tokens with extra information about the current state of the SPIHT decoder right as it is consuming the next bit. This is information that is available during inference time.

What can be used from the spiht algorithm as conditioning at each bit? There's the height and width and channel position of the next coefficient to be queried. [note] When I say 'queried' I mean that it is the singular coefficient that the next bit is answering some information about There's also the next coefficient's value in the coefficient array. Also, there's the local variable 'n'. Also, there's the coefficient's depth, and it's filter array. There's also the type of query being asked of the coefficient; there are 7 lines in the SPIHT encoder/decoder code where bits are emitted/imitted. So overall, 8 concrete values.

Spihtter

Put altogether I call this model architecture 'Spihtter'. Spihtter inputs and outputs a left to right sequence of discrete tokens. It recieves additional embeddings per-token, representing the state of the spiht decoder. Technically, any model could be used to produce next token logits.

The only tricky part is keeping the model fed with spiht metadata during inference. Per generated token, this is a linear time operation. But you have to keep track of where you are in the SPIHT decoding algorithm. I implemented a spiht encoder/decoder in rust which is used to build a dataset of training examples. During inference I use a different python implementation of the spiht algorithm. I do this so I can use python's yield feature and access spiht bits in a stream.

I place image bits inside of html tags, like this: <spiht h=28 w=28 n=11101>010101011110111101</spiht>
Spiht needs some extra information to start decoding. Namely, the height and width and starting value of the variable 'n'. I encode these values as html attributes.

Experiments

I train different Spihtter models on some small toy datasets. I am currently gpu poor and did these small experiments using an M1 macbook air with 8GB of memory.

For generation I use greedy search. I did not play around with sampling settings. I left hyperparameters at their recommended values

MNIST

28x28 greyscale MNIST digits can be encoded in as few as 256 tokens using SPIHT. The small llama model can be trained to produce coherent digits in a very short wall clock time, about 4 minutes. Mamba fails to produce correct mnist digits.

Random dataset samples, compressed to 256 bits. These are the images decoded from the ground truth bits.

digits 0-9 after 1 hour of training, decoded from .3M parameter mamba model, batch size 16, 1000 steps

digits 0-9 after 10 minutes of training, decoded from 2M parameter llama model, batch size 16, 1000 steps

mamba: purple, vs llama: grey, NLL loss. To be fair, LLaMa had almost 10x the number of parameters, but also trained 6x faster. The blue loss curve is without any conditioning information from the spiht decoder. The grey loss curve is from giving the LLM embeddings based on the internal programmatic state of the spiht decoder. The conditioning information from the spiht algorithm vastly improves the ability to model the sequence. model used: llama-small.json

The mamba model didn't match the performance of the llama model. I don't want to be overly critical. I'm relying on a hacked together implementation. I couldn't train using the 'MPS' device because 4D cumsum is not supported as of pytorch 2.2.3. So the training time was much longer. Worryingly, throughout training, most of the images generated from the 0-9 digits are exactly the same. I checked for bugs and couldn't find any; the mamba model is receiving the proper labels, but it seems that the information from the class label at the start of the sequence goes completely forgotten.

Future work

I figured out that you can use spiht to encode latents from the stable diffusion VAE. It make perfect sense why it works. It looks kinda trippy, and the reconstruction quality is quite high at around .25 bpp for most images. I'm doing more research on neural image compression so I'm going to come back to this topic later once I understand the literature better.

Using this trick you can get good images at around ~25k bits. It's also extremely useful if you are trying to save a huge dataset of VAE embeddings on disk. You can just save them as torch tensors, but using spiht as lossy compression you can save 60% of the disk space.