Jet: A Modern Transformer-Based Normalizing Flow

TL;DR

This paper introduces Jet, a transformer-based normalizing flow model that achieves state-of-the-art performance in image generation, demonstrating the potential of transformer architectures in normalizing flows.

Contribution

The paper presents a novel transformer-based normalizing flow architecture that outperforms previous models in image generation tasks.

Findings

Achieved state-of-the-art quantitative results.

Produced high-quality visual samples.

Simplified architecture compared to prior models.

Abstract

In the past, normalizing generative flows have emerged as a promising class of generative models for natural images. This type of model has many modeling advantages: the ability to efficiently compute log-likelihood of the input data, fast generation and simple overall structure. Normalizing flows remained a topic of active research but later fell out of favor, as visual quality of the samples was not competitive with other model classes, such as GANs, VQ-VAE-based approaches or diffusion models. In this paper we revisit the design of the coupling-based normalizing flow models by carefully ablating prior design choices and using computational blocks based on the Vision Transformer architecture, not convolutional neural networks. As a result, we achieve state-of-the-art quantitative and qualitative performance with a much simpler architecture. While the overall visual quality is still behind the current state-of-the-art models, we argue that strong normalizing flow models can help advancing research frontier by serving as building components of more powerful generative models.