US 2024/0404238 A1 · Vector-Quantized Image Modeling

Abstract

Systems and methods are provided for vector-quantized image modeling using vision transformers and improved codebook handling. In particular, the present disclosure provides a Vector-quantized Image Modeling (VIM) approach that involves pre-training a machine learning model (e.g., Transformer model) to predict rasterized image tokens autoregressively. The discrete image tokens can be encoded from a learned Vision-Transformer-based VQGAN (example implementations of which can be referred to as ViT-VQGAN). The present disclosure proposes multiple improvements over vanilla VQGAN from architecture to codebook learning, yielding better efficiency and reconstruction fidelity. The improved ViT-VQGAN further improves vector-quantized image modeling tasks, including unconditional image generation, conditioned image generation (e.g., class-conditioned image generation), and unsupervised representation learning.

Claims30

Dependent claims

A computer-implemented method to perform vector quantization of imagery, the method comprising:

obtaining, by a computing system comprising one or more computing devices, a plurality of input image patches of an image;
processing, by the computing system, the plurality of input image patches with a machine-learned image encoder to generate a plurality of image tokens in a latent space, wherein the plurality of image tokens correspond to the plurality of input image patches, and wherein the machine-learned image encoder performs one or more self-attention operations to process the plurality of input image patches to generate the plurality of image tokens in the latent space;
mapping, by the computing system, the plurality of image tokens to a plurality of quantized codes contained in a quantization codebook that contains a plurality of candidate codes;
and
providing, by the computing system, the plurality of quantized codes as an encoded version of the image.

2 more elements

A computer-implemented method to train a code prediction model, the method comprising:

obtaining, by a computing system comprising one or more computing devices, a training example comprising a training image and a natural language input descriptive of content of the training image;
processing, by the computing system, a plurality of image patches from the training image with a machine-learned image encoder to generate a plurality of image tokens in a latent space, wherein the plurality of image tokens correspond to the plurality of image patches;
mapping, by the computing system, the plurality of image tokens to a plurality of quantized codes contained in a quantization codebook that contains a plurality of candidate codes;
processing, by the computing system, the natural language input with a text encoder portion of the code prediction model to generate a text embedding;
processing, by the computing system, the text embedding with an autoregressive code selection portion of the code prediction model to autoregressively predict a sequence of predicted codes from the quantization codebook;
evaluating, by the computing system, a code prediction loss function that compares the sequence of predicted codes to the plurality of quantized codes;
and
modifying, by the computing system, one or more parameters of the code prediction model based at least in part on the code prediction loss function.

5 more elements

Dependencies are inferred from the claim text.

Vector-Quantized Image Modeling

Life of this patent

Vector-Quantized Image Modeling

Life of this patent

Abstract

Claims30