When Worse is Better: Navigating the compression-generation tradeoff in visual tokenization

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Ramanujan, Vivek, Tirumala, Kushal, Aghajanyan, Armen, Zettlemoyer, Luke, Farhadi, Ali
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866917137661034496
author Ramanujan, Vivek
Tirumala, Kushal
Aghajanyan, Armen
Zettlemoyer, Luke
Farhadi, Ali
author_facet Ramanujan, Vivek
Tirumala, Kushal
Aghajanyan, Armen
Zettlemoyer, Luke
Farhadi, Ali
contents Current image generation methods are based on a two-stage training approach. In stage 1, an auto-encoder is trained to compress an image into a latent space; in stage 2, a generative model is trained to learn a distribution over that latent space. This reveals a fundamental trade-off, do we compress more aggressively to make the latent distribution easier for the stage 2 model to learn even if it makes reconstruction worse? We study this problem in the context of discrete, auto-regressive image generation. Through the lens of scaling laws, we show that smaller stage 2 models can benefit from more compressed stage 1 latents even if reconstruction performance worsens, demonstrating that generation modeling capacity plays a role in this trade-off. Diving deeper, we rigorously study the connection between compute scaling and the stage 1 rate-distortion trade-off. Next, we introduce Causally Regularized Tokenization (CRT), which uses knowledge of the stage 2 generation modeling procedure to embed useful inductive biases in stage 1 latents. This regularization improves stage 2 generation performance better by making the tokens easier to model without affecting the stage 1 compression rate and marginally affecting distortion: we are able to improve compute efficiency 2-3$\times$ over baseline. Finally, we use CRT with further optimizations to the visual tokenizer setup to result in a generative pipeline that matches LlamaGen-3B generation performance (2.18 FID) with half the tokens per image (256 vs. 576) and a fourth the total model parameters (775M vs. 3.1B) while using the same architecture and inference procedure.
format Preprint
id arxiv_https___arxiv_org_abs_2412_16326
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle When Worse is Better: Navigating the compression-generation tradeoff in visual tokenization
Ramanujan, Vivek
Tirumala, Kushal
Aghajanyan, Armen
Zettlemoyer, Luke
Farhadi, Ali
Computer Vision and Pattern Recognition
Machine Learning
Current image generation methods are based on a two-stage training approach. In stage 1, an auto-encoder is trained to compress an image into a latent space; in stage 2, a generative model is trained to learn a distribution over that latent space. This reveals a fundamental trade-off, do we compress more aggressively to make the latent distribution easier for the stage 2 model to learn even if it makes reconstruction worse? We study this problem in the context of discrete, auto-regressive image generation. Through the lens of scaling laws, we show that smaller stage 2 models can benefit from more compressed stage 1 latents even if reconstruction performance worsens, demonstrating that generation modeling capacity plays a role in this trade-off. Diving deeper, we rigorously study the connection between compute scaling and the stage 1 rate-distortion trade-off. Next, we introduce Causally Regularized Tokenization (CRT), which uses knowledge of the stage 2 generation modeling procedure to embed useful inductive biases in stage 1 latents. This regularization improves stage 2 generation performance better by making the tokens easier to model without affecting the stage 1 compression rate and marginally affecting distortion: we are able to improve compute efficiency 2-3$\times$ over baseline. Finally, we use CRT with further optimizations to the visual tokenizer setup to result in a generative pipeline that matches LlamaGen-3B generation performance (2.18 FID) with half the tokens per image (256 vs. 576) and a fourth the total model parameters (775M vs. 3.1B) while using the same architecture and inference procedure.
title When Worse is Better: Navigating the compression-generation tradeoff in visual tokenization
topic Computer Vision and Pattern Recognition
Machine Learning
url https://arxiv.org/abs/2412.16326