I Patched The Pixels Myself

sampel tokenizations
2 layers of hierarchy compared
less tokens. more meaning
we tried to remove as many inductive biases as possible
more sample tokenizations
training results! we are more flops efficient (even using the same JiT backbone) bc we use less tokens

Inspiration

We are inspired by the recent Cartesia post on improved and scalable H-nets, which does dynamic chunking to tokenize language adaptively.

What it does

Because there was no released code nor architecture for adaptive image tokenization (at least in the visage of H-net), we explored the search space of model architectures that were both performant and faithful to the idea of 'operating on the byte-level' and 'removing as many unneeded inductive biases as possible.' We present a core model architecture that operates on the pixel level, dynamically chunking pixels (not necessarily connected ones) together into chunks, with up to 3 levels of hierarchy. We compare against important baselines like JiT, as well as other dynamic tokenizers with baked-in biases. Most importantly, we observe two characteristics of 'learned chunking': 1) Compression. the model compresses boring parts of the image into just a few tokens. 2) Semantics. the model tokenizes semantically similar parts of the image together. We observe this behavior despite the objective simply being very basic (typical image diffusion model objective).

How we built it

We do x-prediction with v-loss following the JiT paper. We stack up to 3 'chunkers' together to go to higher levels of abstraction, with each chunker using attention (pixel-wise, later on chunk-wise) to do sequence mixing. Each chunker, after the sequence-mixing blocks, picks some 'centroids' to merge into using 'importance scores' inspired by TCFormer (from the last attn block, pixel i's strength is the row-wise sum of attn scores S_{ij} over all j. The importance for pixel i then 1 - (max S_{ij} over all j with higher strengths. Intuitively, high-strength pixels are similar to a lot of other ones, and the importance makes sure that high-strength pixels that are similar to each other sort of cancel each other out, so each patch has only one representative, the one with the highest strength). Centers are pixels with importance above a threshold, and they get keys k_i, which are attended to by queries. Each query gets a distribution over the centers, and the chunk embedding is a weighted sum over all of these pixels that attend to it. We present this method as an ~assumption free method that works on arbitrary embedding vectors in any geometry, in contrast to CNN and pixel-specific methods, which work when you want to stay in the pixel-space or grid-like-spaces, but fail when you want to chunk arbitrarily-think-cursed-tetris-shaped regions of the image (what does convolving these chunks even look like?) Models are around 30M parameters, we train on ImageNet in 64x64 resolution.

Challenges we ran into

Pretraining science is hard! It is impossible to do a "thorough enough" job sweeping through the huge space of architectures (emphasis on huuuuge), tuning hyperparams, trying to tune baselines to do these papers justice, in ~24h (we spent a lot of time talking to people, and a bit of time playing Dominion...) It felt like we were the models themselves: faced with an extremely high-dimensional loss landscape with nothing but some naturalistic intuitions to guide us (inductive biases begone!) As with all architecture work, there are a lot of infra issues to optimize down the line in the future as well... (something something the impact you have as a researcher is proportional to how much pain you can cause the infra team?)

Accomplishments that we're proud of

The tokenizations look meaningful. Even if pretraining-efficiency improvements require more careful tuning to confirm, qualitatively, the tokens are meaningful, and they compress in the places we expect them to compress (low information content).

We also know for a fact that we are much more FLOP efficient... like up to 3x more FLOP efficient, using the same JiT backbone and dimensions, all because the tokenizer is better...