Vector Quantized Image Modeling

Abstract

A system and method are provided for vector quantized image modeling using a vision transformer and improved codebook handling. In particular, the present disclosure provides a vector quantized image modeling (VIM) approach that includes pre-training a machine learning model (e.g., a transformer model) to autoregressively predict rasterized image tokens. Individual image tokens can be encoded from a learned vision transformer-based VQGAN (an exemplary implementation of which can be referred to as ViT-VQGAN). The present disclosure proposes multiple improvements over vanilla VQGAN, from architecture to codebook learning, resulting in better performance and reconstruction fidelity. The improved ViT-VQGAN further improves vector quantized image modeling tasks, including unconditional image generation, conditional image generation (e.g., class-conditional image generation), and unsupervised representation learning.

Description

Technical Field

[0001] Cross - Reference to Related Applications This application claims priority and benefit to U.S. Provisional Patent Application No. 63 / 351,131, filed Jun. 10, 2022, and U.S. Provisional Patent Application No. 63 / 252,452, filed Oct. 5, 2021. Each of the above - mentioned applications is hereby incorporated by reference in its entirety.

[0002] This disclosure generally relates to image modeling such as image synthesis and / or image generation. More specifically, this disclosure relates to vector - quantized image modeling using vision transformers and improved codebook handling.

Background Art

[0003] In computer vision, most recent unsupervised or self - supervised learning techniques have focused on applying different random augmentations to images during pre - training for the purpose of distinguishing image instances, and the quality of the learned representations depends on manually picked augmentations such as random brightness, cropping, blurring, etc.

[0004] Some other approaches perform GPT - style generative pre - training on images to predict pixels autoregressively without incorporating knowledge of 2D structure. Here, the pixels are 9 - bit values generated by clustering (R, G, B) pixel values using k - means with k = 512. However, due to a much longer sequence length (e.g., 50176 tokens per image for a 224x224 resolution), much larger memory consumption, and more training computations than those of language models, color encoding does not scale to typical image resolutions. As a result, these artifacts are only applicable at relatively small maximum resolutions (e.g., 64x64) for image recognition, severely limiting the representational power.

[0005] On the other hand, remarkable image generation results, including VQVAE (Oord et al., 2017), DALL-E (Ramesh et al., 2021), and VQGAN (Esser et al., 2021), have been achieved by pre-quantizing images into individual latent variables and autoregressively modeling them. In these methods, a convolutional neural network (CNN) learns to autoencode images, and a second-stage CNN or transformer learns to model the density of the input data. These have been proven effective for image generation, but have not evaluated the learned representations in characteristic tasks.

[0006] One subfield of image generation is the generation of images from text. Here, the image is synthesized based on the input text, and the synthesized image depicts the content described by the input text. The ability to generate images through language is appealing because language is the most natural form of communication, and such generative capabilities can potentially unlock creative applications in many fields such as art, design, and multimedia content creation. One line of research that has recently gained momentum in the space of generating images from text is techniques that leverage diffusion-based text-to-image models that have the ability to generate higher-fidelity images. These models, as a core modeling technique for image stills, avoid the use of individual image tokens and instead use diffusion models to achieve better zero-shot Frechet inception distance (FID) scores and aesthetically pleasing visual outputs for MS-COCO. Despite these advancements, assuming there is a rich body of prior work on large language models and general-purpose interfaces for individual tokens that are more easily applicable to more modalities, applying autoregressive modeling to the task of generating images from text remains practically appealing.

Summary of the Invention

Means for Solving the Problems

[0007] Aspects and advantages of embodiments of the present disclosure are described in part in the following description, or can be learned from the description, or can be learned through the implementation of the embodiments.

[0008] One exemplary aspect of the present disclosure is directed to a computer-implemented method for performing vector quantization of an image. The method includes obtaining, by a computing system comprising one or more computing devices, a plurality of input image patches of the image. The method includes 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, where the plurality of image tokens correspond to the plurality of input image patches and the machine-learned image encoder performs one or more self-attention operations to process the plurality of input image patches and generate the plurality of image tokens in the latent space. The method includes mapping, by the computing system, the plurality of image tokens to a plurality of quantization codes included in a quantization codebook including a plurality of candidate codes. The method includes providing, by the computing system, the plurality of quantization codes as an encoded version of the image.

[0009] In some implementations, the machine-learned image encoder comprises a vision transformer model.

[0010] In some implementations, the machine-learned image encoder performs one of one or more self-attention operations on the plurality of input image patches.

[0011] In some implementations, the method further includes: a step of processing, by a computing system, a plurality of quantization codes with a machine-learned image decoder to generate a plurality of synthetic image patches that form a synthetic image; a step of evaluating, by the computing system, a loss function that provides a loss value based at least in part on the synthetic image; and a step of modifying, by the computing system, at least in part based on the loss function, a machine-learned image encoder, a machine-learned image decoder, and one or more of a plurality of candidate codes.

[0012] In some implementations, the machine-learned image decoder comprises a vision transformer model.

[0013] In some implementations, the loss function includes a logit Laplace loss term, an L2 loss term, a perceptual loss term, and / or an adversarial generation network loss term.

[0014] In some implementations, the step of mapping, by a computing system, a plurality of image tokens to a plurality of quantization codes included in a quantization codebook including a plurality of candidate codes includes: a step of projecting, by the computing system, the plurality of image tokens into a lower-dimensional space; and after the step of projecting the image tokens into the lower-dimensional space, a step of mapping, by the computing system, the plurality of image tokens to the plurality of quantization codes included in the quantization codebook.

[0015] In some implementations, the step of mapping, by a computing system, a plurality of image tokens to a plurality of quantization codes included in a quantization codebook including a plurality of candidate codes includes: a step of applying L2 regularization to one or both of the plurality of image tokens and the plurality of candidate codes by the computing system; and after the step of applying L2 regularization, a step of mapping, by the computing system, the plurality of image tokens to the plurality of quantization codes included in the quantization codebook.

[0016] In some implementations, the method further includes: a step of autoregressively predicting, by a computing system using a machine-learned code prediction model, a plurality of predicted codes based at least in part on one or more of a plurality of quantized codes from a quantized codebook; and a step of generating, by the computing system, a plurality of synthetic image patches that form a synthetic image by processing the plurality of predicted codes with a machine-learned image decoder.

[0017] In some implementations, the method further includes: a step of evaluating, by a computing system, a code prediction loss function that evaluates a negative log-likelihood based on the plurality of predicted codes; and a step of modifying, by the computing system, one or more parameters of the machine-learned code prediction model based on the code prediction loss function.

[0018] In some implementations, the step of autoregressively predicting, by a computing system using a machine-learned code prediction model, a plurality of predicted codes includes conditioning, by the computing system, the machine-learned code prediction model with auxiliary conditioning data that describes one or more desired characteristics of a synthetic image.

[0019] In some implementations, the auxiliary conditioning data includes a class label that describes a desired class of the synthetic image.

[0020] In some implementations, the auxiliary conditioning data includes natural language text tokens.

[0021] In some implementations, the method further includes: a step of extracting, by a computing system, one or more intermediate features from the machine-learned code prediction model; and a step of predicting, by the computing system, a class label for an image based at least in part on the intermediate features.

[0022] Another exemplary aspect is directed to a computer-implemented method for performing vector quantization of an image. The method includes obtaining, by a computing system comprising one or more computing devices, a plurality of quantization codes that form an encoded version of the image, where the plurality of quantization codes are selected by mapping a plurality of image tokens generated by a machine-learned image encoder model to a plurality of quantization codes included in a quantization codebook that includes a plurality of candidate codes. The method includes generating, by the computing system, a plurality of decoded image patches that form a decoded version of the image by processing the plurality of quantization codes with a machine-learned image decoder. In some implementations, one or both of the machine-learned image encoder model and the machine-learned image decoder are configured to perform one or more self-attention operations.

[0023] Another exemplary aspect is directed to a computer system comprising one or more processors and one or more non-transitory computer-readable media collectively storing instructions that, when executed by the one or more processors, cause the one or more processors to perform any of the methods described herein.

[0024] Another exemplary aspect is directed to one or more non-transitory computer-readable media collectively storing instructions that, when executed by one or more processors, cause the one or more processors to perform any of the methods described herein.

[0025] Another exemplary aspect is directed to a computer system comprising one or more processors and one or more non-transitory computer-readable media collectively storing a machine-learned image processing model. The machine-learned image processing model comprises an encoder configured to encode one or more input image patches into one or more image tokens in a latent space, a quantization unit configured to quantize one or more image tokens into one or more quantization codes selected from a codebook, a code prediction unit configured to predict one or more predicted quantization codes from the codebook, at least in part based on one or more quantization codes, and a feature prediction unit configured to generate one or more characteristic predictions for the input image patch, at least in part based on data extracted from the code prediction unit.

[0026] In some implementations, the machine-learned image processing model further comprises a decoder configured to generate a reconstructed image patch based on one or more quantization codes or to generate a synthetic image patch, at least in part based on one or more predicted quantization codes.

[0027] In some implementations, one or more characteristic predictions include an image classification prediction.

[0028] Another exemplary aspect is directed to a computer-implemented method for performing image generation from text. The method includes obtaining, by a computing system including one or more computing devices, a natural language input that describes desired image content. The method includes processing, by the computing system, the natural language input in a text encoder portion of a machine-learned code prediction model to generate a text embedding. The method includes autoregressively predicting, by the computing system, in an autoregressive code selection portion of the machine-learned code prediction model, a sequence of codes predicted from a quantized codebook including a plurality of candidate codes, by processing the text embedding. The method includes generating, by the computing system, a plurality of synthetic image patches that form a synthetic image, by processing the sequence of quantized codes with a machine-learned image decoder. The synthetic image depicts the desired image content.

[0029] In some implementations, one or more of the text encoder portion of the machine-learned code prediction model, the autoregressive code selection portion of the machine-learned code prediction model, and the machine-learned image decoder are configured to perform one or more self-attention operations.

[0030] In some implementations, one or more of the text encoder portion of the machine-learned code prediction model, the autoregressive code selection portion of the machine-learned code prediction model, and the machine-learned image decoder include a transformer neural network.

[0031] In some implementations, one or both of the machine-learned image decoder and the codebook are learned together with an image encoder model.

[0032] In some implementations, the text encoder portion of the machine-learned code prediction model is pre-trained on a pre-training task.

[0033] Another exemplary aspect is directed to a computer-implemented method for training a code prediction model. The method includes obtaining, by a computing system comprising one or more computing devices, training examples including training images and natural language inputs that describe the content of the training images. The method includes processing, by the computing system, a plurality of image patches from the training images with a machine-learned image encoder to generate a plurality of image tokens in a latent space, where the plurality of image tokens correspond to the plurality of image patches. The method includes mapping, by the computing system, the plurality of image tokens to a plurality of quantization codes included in a quantization codebook including a plurality of candidate codes. The method includes processing, by the computing system, the natural language inputs with a text encoder portion of the code prediction model to generate a text embedding. The method includes 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 codes predicted from the quantization codebook. The method includes evaluating, by the computing system, a code prediction loss function that compares the predicted sequence of codes to the plurality of quantization codes. The method includes 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.

[0034] In some implementations, the machine-learned image encoder performs one or more self-attention operations to process the plurality of input image patches and generate the plurality of image tokens in the latent space.

[0035] In some implementations, the code prediction loss function evaluates the negative log-likelihood of the predicted codes with respect to the quantization codes.

[0036] In some implementations, the machine-learned image encoder and the codebook are pre-trained with an image decoder in an autoencoder architecture.

[0037] In some implementations, the step of a computing system changing one or more parameters of a code prediction model, at least partially based on a code prediction loss function, includes the step of a computing system changing one or more parameters of both a text encoder part and an autoregressive code selection part of the code prediction model, at least partially based on the code prediction loss function.

[0038] Other aspects of the present disclosure are directed to various systems, devices, non-transitory computer-readable media, user interfaces, and electronic devices.

[0039] These and other features, aspects, and advantages of the various embodiments of the present disclosure will be better understood with reference to the following description and the appended claims. The accompanying drawings, which are incorporated herein and form a part of this specification, illustrate exemplary embodiments of the present disclosure along with the description that serves to explain the relevant principles.

[0040] A detailed discussion of embodiments directed to those skilled in the art is described herein, which refers to the accompanying drawings.

Brief Description of the Drawings

[0041]

Figure 1

Figure 2

Figure 3

Figure 4

Figure 5

Figure 6

Figure 7

Figure 8A

Figure 8B

Figure 8C

Mode for Carrying Out the Invention

[0042] Reference numerals repeated throughout a plurality of figures are intended to identify the same features in various implementations.

[0043] Summary Generally, the present disclosure is directed to 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 includes pre-training a machine learning model (e.g., a transformer model) to autoregressively predict rasterized image tokens. Individual image tokens can be encoded from a learned vision transformer-based VQGAN (an exemplary implementation of which can be referred to as ViT-VQGAN). The present disclosure proposes a number of improvements over vanilla VQGAN, from architecture to codebook learning, resulting in better efficiency and reconstruction fidelity. The improved ViT-VQGAN further improves vector quantized image modeling tasks, including unconditional image generation, conditional image generation (e.g., class-conditional image generation), and teacherless representation learning.

[0044] Exemplary experiments demonstrate the technical effects of the proposed techniques. In particular, when training on ImageNet at 256x256 resolution, an exemplary implementation of the present disclosure achieves a starting score (IS) of 175.1 and a Frechet inception distance (FID) of 4.17, a dramatic improvement over vanilla VQGAN, which obtains 70.6 and 17.04 for the starting score (IS) and Frechet inception distance (FID), respectively. Based on ViT-VQGAN and teacherless pre-training, the exemplary experiments further evaluate the pre-trained transformer by averaging intermediate features, similar to ImageGPT (iGPT). This ImageNet pre-trained VIM-L significantly outperforms iGPT-L in terms of linear probe accuracy, from 60.3% to 72.1%, for a similar model size. ViM-L also outperforms iGPT-XL, which is trained with additional web image data and a larger model size.

[0045] Another exemplary aspect of the present disclosure is directed to the application of the proposed vector quantization image model to the image generation task from text. Specifically, the present disclosure provides an autoregressive model that can generate realistic images such as photos from text descriptions, which has advantages different from diffusion-based methods. Exemplary implementations of these models can be referred to as "Babeldraw". In some implementations, Babeldraw uses an improved image tokenizer and the transformer-based ViT-VQGAN described herein to encode an image as a sequence of individual tokens. This naturally reduces the current task to the well-known problem of machine conversion from text to image tokens. By scaling the second, encoder-decoder transformer parameters from 350M and 750M to 3B and 20B, consistent quality improvements are obtained. Compared with recent diffusion-based models, when fine-tuned on the MS-COCO training split, the 20B Babeldraw model achieves an equivalent zero-shot MS-COCO Frechet inception distance (FID) score of 7.31 and a previously unheard-of FID score of 4.03. Using a post-super-resolution upsampler trained with a frozen 256x256 image tokenizer, Babeldraw reliably generates realistic images like high-resolution photos.

[0046] Accordingly, exemplary implementations of the present disclosure relate to vector quantization image modeling (VIM) methods for both image generation and image understanding tasks. Specifically, some exemplary implementations conform to the following two-step approach.

[0047] Stage 1: Vector quantization. Assuming an image (e.g., of resolution 256x256), the Vision Transformer-based VQGAN encodes it into several discretized latent codes (e.g., 32x32) contained within a codebook (e.g., the codebook size may be 8192). In addition to VQGAN (Esser et al., 2021), this disclosure proposes multiple improvements from architecture to codebook learning. The resulting ViT-VQGAN has better efficiency and reconstruction fidelity in terms of per-pixel reconstruction measurement, inception score (IS), and Fréchet inception distance (FID). ViT-VQGAN can be thoroughly trained on image-only data using a combined objective function of logit-laplace loss, L2 loss, adversarial loss, and / or perceptual loss.

[0048] Stage 2: Vector quantization image modeling. Some exemplary implementations can include the step of training a Transformer model to autoregressively predict rasterized (e.g., 32x32 = 1024) image tokens, where the image tokens are encoded by a frozen Stage 1 ViT-VQGAN. In unconditional image synthesis or unsupervised learning, a Transformer model with only a decoder can be trained to predict the next token. In conditional image synthesis (e.g., class-conditional image synthesis), conditional data (e.g., class id tokens) can be appended to the beginning of all image tokens.

[0049] To evaluate the quality of unsupervised learning, intermediate Transformer features can be extracted (and averaged), and a linear head can be trained to predict class logits (also known as linear probe accuracy).

[0050] One key component for improving both image generation and image understanding in VIM is to have a better image quantizer in terms of computer efficiency and reconstructed quality. An efficient quantizer can speed up the training in stage 2, where random augmentations are first applied to the image and then to the encoder of the image quantizer as input tokens. Further, an image quantizer with better reconstructed quality can potentially reduce the information loss compared to the original image in the pixel space, which is important for the image understanding task.

[0051] A further exemplary aspect of the present disclosure is directed to an autoregressive image generation model that can generate realistic images such as photos from text descriptions. The autoregressive model provided herein can benefit from a better image tokenizer and from the scaling of the model. More specifically, the present disclosure provides a transformer-based sequence-to-sequence model for image generation (an exemplary implementation thereof can be referred to as "Babeldraw"). The image generation model can take tokenized text tokens as input to the encoder and predict individual image tokens using the decoder in an autoregressive manner. The image tokens, in some examples, can generate a model better than VQGAN in terms of efficiency and image reconstruction fidelity by the transformer-based ViT-VQGAN image tokenizer described herein.

[0052] In some implementations, all components of the image generation model, namely the encoder, decoder, and image tokenizer, may be based on the standard Transformer model. This makes it easier to scale the model. To test the limits of the above two-stage text-to-image framework, the parameter sizes of the exemplary Babeldraw models were scaled from 350M to 750M, 3B, and 20B, and quality improvements were observed in both text-image alignment and image photorealism. The 20B Babeldraw model achieves a strong zero-shot FID score of 7.31, comparable to state-of-the-art diffusion-based models. Notably, when fine-tuned on MS-COCO, the exemplary Babeldraw model achieves an unprecedentedly low FID score of 4.07.

[0053] While most works stick to the MS-COCO benchmark, exemplary experiments show that strong zero-shot and fine-tuned results can still be achieved on datasets with localized annotations, which on average have descriptions four times longer than MS-COCO. These results clearly demonstrate the strong generalization ability of the proposed model for longer descriptions.

[0054] Thus, the present disclosure realizes several contributions related to image generation from text. First, it is understood that for a two-stage text-to-image generation model, the image tokenizer is one of the key factors and is directly related to the image quality of the generated images. Second, a simple scaling encoder-decoder using exactly the same image tokenizer and training data is effective for image generation from text. Third, a text encoder pre-trained only with BERT gently aids natural language understanding for a text-to-image system. The text encoder fine-tuned for image generation from text has worse results in language understanding tasks, suggesting a potential gap between general language understanding and vision-based language understanding.

[0055] The systems and methods of the present disclosure provide several technical effects and advantages. As one example, the proposed technique achieves equivalent results to current state-of-the-art models (e.g., iGPT) in certain tasks (e.g., image recognition using generative pre-training), but can be achieved with a smaller model and less data. By enabling the use of a smaller model with less data (while obtaining equivalent results), the consumption of computing resources such as processor usage, memory usage, network bandwidth, etc. is reduced. Thus, the proposed technique saves computing resources and improves the functionality of the computer.

[0056] As another example, the exemplary implementations of the present disclosure have excellent efficiency and reconstruction fidelity in terms of pixel-wise reconstruction measurement, inception score (IS), and Frechet inception distance (FID) with respect to existing models. Thus, the present disclosure represents an improvement in the ability to encode image data into an encoded representation. The encoded data has a number of advantages including savings in memory space and network bandwidth. The present disclosure enables the system to encode an image to achieve such advantages while reducing the data loss experienced when decoding the image into a decoded image.

[0057] As another exemplary technical effect and advantage, the present disclosure realizes improved image generation and image understanding using VIM by providing an improved image quantizer in terms of computer efficiency and reconstructed quality. An efficient quantizer can speed up the training in stage 2, where random augmentations are first applied to the image and then to the encoder of the image quantizer as input tokens. By increasing the training speed, savings in computing resources such as processor usage, memory usage, and / or network bandwidth usage are achieved. Further, an image quantizer with better reconstructed quality can potentially reduce information loss compared to the original image in the pixel space, which is important for image understanding tasks. Thus, the computer performance for image understanding tasks can be improved.

[0058] Referring now to the figures, exemplary embodiments of the present disclosure are discussed in further detail.

[0059] Exemplary Stage 1 Techniques FIG. 1 depicts a diagram of an exemplary machine learning model for performing vector quantization of an image, in accordance with an exemplary embodiment of the present disclosure.

[0060] Specifically, as shown in FIG. 1, a computing system can obtain a plurality of input image patches 12 of an image 14.

[0061] The computing system can process the plurality of input image patches 12 with a machine-learned image encoder 16 to generate a plurality of image tokens 18 in a latent space. The plurality of image tokens 18 may correspond to the plurality of input image patches 12. The machine-learned image encoder 16 can perform one or more self-attention operations to process the plurality of input image patches 12 and generate the plurality of image tokens 18 in the latent space.

[0062] The computing system can map a plurality of image tokens 18 to a plurality of quantization codes 20 included in a quantization codebook 22 that includes a plurality of candidate codes. The computing system can provide the plurality of quantization codes 20 as an encoded version of the image.

[0063] Further, in some implementations, the computing system (i.e., its computing devices of the same or different components) can process the plurality of quantization codes 20 with a machine-learned image encoder 24 to generate a plurality of synthetic image patches 26 that form a synthetic image 28.

[0064] The computing system can evaluate a loss function 30 that provides a loss value based at least in part on the synthetic image. The loss function can include a logit Laplace loss term, an L2 loss term, a perceptual loss term, and / or an adversarial generative network loss term (e.g., GAN evaluation is specifically illustrated at 30 while others are not illustrated).

[0065] The computing system can modify one or more of the machine-learned image encoder 16, the machine-learned image decoder 24, and the plurality of candidate codes 22 based at least in part on the loss function.

[0066] In some implementations, the step of mapping the plurality of image tokens 18 to the plurality of quantization codes 20 can include projecting the plurality of image tokens 18 into a lower-dimensional space and, after the step of projecting the image tokens 18 into the lower-dimensional space, mapping the plurality of image tokens to the plurality of quantization codes included in the quantization codebook. The quantization codes 20 can then be reprojected back into a higher-dimensional space before being processed by the decoder 24.

[0067] In some implementations, the step of mapping a plurality of image tokens 18 to a plurality of quantization codes 20 can include applying L2 regularization to one or both of the plurality of image tokens 18 and one of the plurality of candidate codes in the codebook 22, and after applying the L2 regularization step, mapping the plurality of image tokens 18 to the plurality of quantization codes 20 included in the quantization codebook 22.

[0068] Exemplary Stage 2 Techniques FIG. 2 depicts a diagram of an exemplary machine learning model for performing vector quantization of an image, in accordance with an exemplary embodiment of the present disclosure.

[0069] Specifically, as shown in FIG. 2, a computing system can obtain one or more input image patches 212 of an image 214.

[0070] The computing system can process one or more input image patches 212 with a machine-learned image encoder 216 to generate one or more image tokens 218 in a latent space. The one or more image tokens 218 can correspond to the one or more input image patches 212. The machine-learned image encoder 216 can perform one or more self-attention operations to process the one or more input image patches 212 and generate one or more image tokens 218 in the latent space. For example, the machine-learned image encoder 216 can be trained or learned according to the process shown in FIG. 1.

[0071] The computing system can map one or more image tokens 218 to one or more quantization codes 220 included in a quantization codebook 222 that includes a plurality of candidate codes. For example, the quantization codebook 222 can be learned according to the process shown in FIG. 1.

[0072] Furthermore, in some implementations, a computing system (i.e., computing devices of the same or different components thereof) can process one or more quantized codes 220 with a code prediction model 224 that has learned the one or more quantized codes to autoregressively predict a plurality of predicted codes 226 from a quantized codebook 222, at least partially based on the one or more quantized codes 220.

[0073] The computing system can process the plurality of predicted codes 226 with a machine-learned image decoder 228 to generate a plurality of synthetic image patches 230 that form a synthetic image 232.

[0074] In some implementations or instances, the computing system can further evaluate a code prediction loss function that evaluates a negative log-likelihood based on the plurality of predicted codes 226. The computing system can change one or more parameters of the machine-learned code prediction model 224 based on the code prediction loss function. For example, the code prediction model 224 can be trained to learn the distribution of tokens over a corpus of images, for example, using the code prediction loss function.

[0075] In some implementations, during training and / or inference, the machine-learned code prediction model 224 can be conditioned with auxiliary conditioning data that describes one or more desired characteristics of the synthetic image. As one example, the auxiliary conditioning data can include a class label that describes the desired class of the synthetic image 232. For example, in FIG. 2, the model 224 is conditioned with the label [Cat], and as a result, the synthetic image 232 depicts a cat, and the synthetic image 232 shares visual characteristics with the image 214. As another example, the auxiliary conditioning data can include natural language text tokens and / or any other conditioning data that provides a preference for image synthesis.

[0076] FIG. 2 illustrates an exemplary approach in which an image generation process is conditioned on both a patch 212 from an input image 214 and also a text input (e.g., a class label). However, in some implementations, image generation can be conditioned on text only. For example, FIG. 3 illustrates an example in which the generation of a synthetic image 332 is conditioned on a natural language input 334 of "angry cat". For example, the natural language input 334 can be converted into tokens and then provided as input to the code prediction model 224.

[0077] According to aspects of the present disclosure, in some implementations that perform image generation tasks from text, including text conditioning, the code prediction model 224 can include a text encoder that generates text embeddings from a text input.

[0078] As an example, FIG. 4 illustrates an exemplary image generation approach in which the code prediction model 224 includes a text encoder 402 and an autoregressive code selector 404. The text encoder 402 can receive an input text 334 and generate a text embedding 406. The text embedding 406 can be provided as input to the autoregressive code selector 404. For example, the text embedding 406 can be provided as an initial prompt token to the autoregressive code selector 404.

[0079] In some implementations, one or both of the text encoder 402 and the autoregressive code selector 404 can be a transformer model or another model that implements self-attention, or can include a transformer model or another model that implements self-attention. In some implementations, the text encoder can be pre-trained (e.g., using masked language modeling pre-training techniques, contrastive learning techniques, and / or other pre-training techniques). In some implementations, the parameters of the text encoder 402 and the autoregressive code selector 404 can be trained together (e.g., using a code prediction loss function) to learn the token distribution over a set of text and image training examples, each containing a pair of text and image.

[0080] Exemplary vector quantized image in ViT-VQGAN A vector quantized variational autoencoder (VQVAE) is a CNN-based autoencoder whose latent space is a matrix of individually learnable variables and is thoroughly trained by straight-through estimation. VQGAN is a model that improves VQVAE by introducing an adversarial loss generated by a discriminator. In this specification, further improvements to VQGAN are realized to improve efficiency and enhance reconstruction quality.

[0081] Exemplary VQGAN in Vision Transformer The core network architecture used by both VQVAE and VQGAN to encode and reconstruct images is a CNN. VQGAN introduces a transformer-like element in the form of a non-localized attention block, enabling it to capture distance interactions with fewer layers. An exemplary implementation of this disclosure takes a further step by replacing the CNN encoder and decoder with a vision transformer (ViT). Given sufficient data (where there is a lot of unlabeled image data), ViT-VQGAN has been found to be less suppressed by the inductive prior imposed by convolution. Furthermore, ViT-VQGAN provides better computer efficiency with an accelerator and generates higher-quality reconstructions, as shown in Table 1 below.

[0082]

Table 1

[0083] Throughput is benchmarked on the same 128 CloudTPUv4 devices.

[0084] An exemplary encoder of ViT-VQGAN first maps non-overlapping 8x8 image patches to image tokens, and then a transformer block encodes a 256x256 resolution image into a sequence of 32x32 = 1024 tokens. An exemplary decoder performs the reverse operation, mapping each image token back to an 8x8 image patch from the latent variable and regrouping them into a 256x256 image. In some implementations, a two-layer feed-forward network with a tanh activation layer in the middle is applied at the output of the transformer block. In some implementations, no activation is applied to the output of the ViT-VQGAN encoder or decoder (except for the average prediction of the logit-laplace loss). In some implementations, a sigmoid activation is applied to the average prediction of the decoder due to the logit-laplace loss. This exemplary approach results in high-quality reconstructions without visible grid-like artifacts.

[0085] Exemplary codebook learning Vanilla VQVAE typically suffers from a low codebook usage due to an insufficient initialization of the codebook. Thus, during training, a significant portion of the codes are hardly used or are invalid. As the effective codebook size decreases, worse reconstructions in stage 1 quantizer training and insufficient versatility in stage 2 image synthesis are brought about. As a result, VQGAN relies on top-k and top-p (nucleus) sampling discovery methods with a default codebook size of 1024 to obtain the best results in image synthesis. Exemplary implementations of the present disclosure include two improvements that can significantly boost the codebook usage even with a codebook size larger than 8192. During image synthesis, the exemplary implementations perform simple sampling at a temperature of 1.0 without top-k and top-p discovery methods.

[0086] One exemplary training objective for vector quantization is defined as follows.

[0087] [Number]

[0088] Here,

[0089] [Number]

[0090] is the stopping gradient operator, β is the commitment loss hyperparameter set to 0.25 in all current experiments, and e is the codebook vector. In some implementations, the quantization codebook index is determined by searching for the codebook vector closest to the input feature z e (x) in terms of Euclidean distance,

[0091] [Number]

[0092] is brought about.

[0093] Factorized code. Some exemplary implementations of the present disclosure include a linear projection from the output of the encoder to a low-dimensional latent variable space (reduced from 768d vectors per code to 32d or 8d vectors) for code index search. This has a direct increasing effect on the codebook usage. Factorization can be seen as separating the code index and the code embedding. Some exemplary implementations search for the closest variable encoded from the input on a lower-dimensional search space and then project the matched latent code into the high-dimensional embedding space. Exemplary experiments show that by reducing the dimensionality of the search space from 256d to 32d, the reconstruction quality is consistently improved.

[0094] l2-regularized code. Some exemplary implementations have the encoded latent variable z e(x) and the codebook latent variable e are also applied with l2 regularization. In some implementations, the codebook variables are initialized from a normal distribution. By mapping all latent variables onto a sphere, the l2-regularized latent variables

[0095]

Number

[0096] 's Euclidean distance evolves into the cosine similarity of two vectors between z e (x) and e, and the training stability and reconstruction quality shown in our experiments are further improved.

[0097] Exemplary ViT-VQGAN training losses Some exemplary implementations use a combination of logit Laplace loss, l2 loss, perceptual loss based on the VGG network, and GAN loss in the architecture of the StyleGAN discriminator. To optimize the image reconstruction quality, codebook usage, FID, and inception score, the loss balancing weights can be configured with a hyperparameter sweep. After the sweep, some exemplary implementations apply the same set of hyperparameters of the training loss to all datasets including CelebA-HQ, FFHQ, and ImageNet. The logit Laplace loss is a regularized l1 loss assuming that the pixel-level noise is from a Laplace distribution, while the l2 loss can be seen as assuming that the noise is from a Gaussian distribution. It has been found that the logit Laplace loss contributes to the codebook usage, while the l2 loss and perceptual loss significantly contribute to the FID. The final combination of losses used by default is L = L VQ + 0.1L Adv + 0.1L Perceptual + 0.1L Logit-laplace + 1.0L2.

[0098] One provisional application regarding VGG-based perceptual loss states that the VGG network is pre-trained with supervised classification loss, so there may be leakage of monitoring to stage 2 for linear probe accuracy measurement. Therefore, for all reported unsupervised learning results, some exemplary implementations exclude the perceptual loss during ViT-VQGAN training. For all unconditional and class-conditional image synthesis, some exemplary implementations use a ViT-VQGAN quantizer trained with perceptual loss, resulting in higher-fidelity reconstructions.

[0099] Exemplary Vector Quantization Image Modeling Using the learned ViT-VQGAN, an image can be encoded into individual latent codes id flattened in raster order, similar to Image GPT. Using only the decoder's transformer model, the density P(x) of the image data can be autoregressively modeled as follows.

[0100]

Equation

[0101] Here, θ are learnable weights. The training objective is to minimize the negative log-likelihood of the data

[0102]

Equation

[0103] by minimizing it.

[0104] Table 2 summarizes an exemplary architecture configuration for the transformer. Some exemplary implementations first embed at each position an individual image token id into a learnable embedding space together with an additional learnable 2D position embedding. Both embedding dimensions are the same as the model dimension. Some exemplary implementations apply a stack of transformer blocks to the input with causal attention over the entire sequence. A dropout rate of 0.1 can be used for all remaining outputs, activation outputs, and attention outputs. Some exemplary implementations apply additional layer regularization at the final layer of all transformer blocks.

[0105]

Table 2

[0106] Exemplary Image Synthesis Unconditional image generation can be achieved by simply sampling each token from the output softmax distribution in a pre-trained generative transformer model. All samples used in both qualitative and quantitative results can be obtained without temperature annealing. The sampled tokens can then be fed into the decoder of ViT-VQGAN to decode the output image. An exemplary default stage 1 ViT-VQGAN encodes an input image of resolution 256x256 into a latent code of 32x32 with a codebook size of 8192, while the stage 2 transformer takes flattened image tokens with a total length of 1024.

[0107] Class-conditional ImageNet generation is also a widely used benchmark for measuring the capabilities of models for image synthesis. Some exemplary implementations extend unconditional generation to class-conditional generation by adding class id tokens before the image tokens. A separate embedding layer can learn from the start for both the class id tokens and the image tokens, and the embedding dimension is the same as the Transformer model dimension. During sampling, the class id tokens can be placed at the first position and the remaining image tokens can be decoded autoregressively.

[0108] Exemplary unsupervised learning In image understanding tasks, some exemplary implementations feed all the image tokens of the input into a pre-trained Transformer and obtain a sequence of features of 1024 tokens. Some exemplary implementations take the layer output at a particular block l over a total of block L, average over the (frozen) sequence of token features, and insert a (learnable) softmax layer to project the averaged features into class logits. Some exemplary implementations simply take the output of one particular Transformer block instead of concatenated different block outputs as in iGPT. Often, the most discriminative features for linear probes are typically near the middle of all the Transformer blocks.

[0109] Exemplary Babeldraw model Some exemplary implementations of Babeldraw are two-stage models consisting of an image tokenizer and an autoregressive model over both language and image tokens.

[0110] Exemplary image tokenizer Auto-regressive text-image models rely on some form of linearization of 2D images into 1D sequences of patch representations. At the extreme, these are just pixels, but these require modeling very long sequences even for relatively small images (e.g., a 256x256 image results in 65536 rasterized pixels). Unfortunately, it is based on a very low-level representation of the input rather than something richer known by the position of the pixels in the context of the image. Many works have solved the above problem by using individual variational autoencoders to learn a quantized representation of image patches over a set of raw images. Instead of learning a representation that can take on any value in the latent space, a visual codebook is learned that fits patches into the entry of the nearest codebook that is learnable and indexable in the entire latent space. These entries are considered visual word types, and the representation of any of these words in a patch of a given image is thus an image token.

[0111] To be most useful for the second-stage model, the image tokenizer should (a) learn an effective visual codebook that supports a balanced use of its entries over a wide range of images, and (b) support the reconstruction of a sequence of visual tokens as a high-quality output image. Some exemplary implementations use ViT-VQGAN, an example described herein, which addresses both requirements, namely, training stability, reconstruction quality, and regularization and factored codes that contribute to codebook usage, by using techniques to improve codebook learning.

[0112] In the Babeldraw image tokenizer, some exemplary implementations train as described above for ViT-VQGAN, but use images in larger image-text training data (e.g., instead of ImageNet, CelebA-HQ, or FFHQ). Some exemplary implementations first train a ViT-VQGAN-Small configuration (8 blocks, 8 heads, model dimension 512, and hidden dimension 2048) and learn 8192 image token classes for the codebook. Note that only the second stage encoder-decoder training relies on the learned quantizer's encoder and codebook. To further improve the visual sensitivity of the reconstructed images, some exemplary implementations freeze the encoder and codebook after training and fine-tune with a decoder of a larger size (32 blocks, 16 heads, model dimension 1280, and hidden dimension 5120). Both the input and output of the image tokenizer may be of resolution 256x256.

[0113] Some exemplary implementations exhibit pixelation patterns (saturated pixel values) in the output of ViT-VQGAN when magnifying a part of the image. To address this issue, some exemplary implementations remove the final sigmoid layer and expose the raw values as RGB pixel values (in the range [0,1]). Conveniently, this modification can be hot-swapped into an already trained model by fine-tuning the decoder.

[0114] In addition, an image with a resolution of 256x256 captures most of the content, structure, and texture, while higher resolution images such as 512x512 or 1024x1024 are more visually satisfying. For this purpose, some exemplary implementations utilize a simple super-resolution module on top of the image tokenizer. The stacked convolutional layers with the remaining connections can be used as a super-resolution network module. It can be trained with the same losses (perceptual loss, StyleGAN loss, and L2 loss) in ViT-VQGAN to map from a reconstructed image to a higher resolution reconstructed image. Note that the diffusion model can also be used here as an iterative refinement super-resolution module, with or without conditioning, on the text input.

[0115] Image from text in an exemplary encoder-decoder In some exemplary implementations, a standard encoder-decoder transformer model is trained in a second stage by treating the image from text as a sequence-to-sequence modeling problem. The model takes the text as input and is trained using the next token prediction of the image latent code generated from the image tokenizer of the first stage. At inference time, the model can sample the image tokens autoregressively, which are then decoded to pixels using the decoder of the previously learned image tokenizer.

[0116] The text prompt can be truncated to a maximum of 128, and an image with a resolution of 256x256 can be encoded into tokens of length 1024 (32x32). Some exemplary implementations train with cross-entropy loss and use the int8-quantized Adafactor optimizer with a constant decay second moment factor (beta1 = 0.9, beta2 = 0.96). The data type can be bfloat16 or dropped to the attention projection and feed-forward transformer layers, while all layer norms and model outputs can be kept as float32. All models can use conv-shaped masked sparse attention, and some exemplary implementations train four size variants in the range of 350 million to 20 billion parameters.

[0117] Most existing two-stage text-to-image generation models are decoder-only models. An early investigation at the scale of 350 million to 750 million parameters found that the encoder-decoder variant of Babeldraw outperformed the decoder-only ones in both training loss and image generation quality from text.

[0118] Exemplary Text Encoder Pretraining The encoder-decoder architecture separates image token generation from text encoding. Thus, some exemplary implementations also include a straightforward way to warm start the model with a pre-trained text encoder. Intuitively, a good text encoder for visual synthesis should enable general language understanding in addition to visually strong prompts. Some exemplary implementations pre-train the text encoder on two datasets, namely, the Colossal Clean Crawled Corpus (C4) for the BERT pre-training objective and an image-text dataset for a contrastive learning objective. After pre-training, some exemplary implementations continue to train both the encoder and the decoder for image generation from text with a softmax cross-entropy loss over a vocabulary of 8192 individual image tokens.

[0119] Guidance and reranking without an exemplary classifier Guidance without a classifier is typically used in situations where it improves the sample quality of a diffusion model without a pre-trained classifier. In this setup, the generative model G is trained to be able to perform unconditional generation G(z) (where z represents random noise) and conditional generation G(z, c) (where c represents some conditions such as the language description in this case). It is simply implemented as a conditional vector that is randomly omitted (masked or switched to a learned embedding) with a certain probability. Then, during the inference process, sampling of the output I is performed by using a linear combination of unconditional and conditional predictions. I = G(z) + λ(G(z, c) - G(z)) (3) Here, λ is a hyperparameter representing the weight of the guidance without a classifier. Intuitively, it decreases the unconditional likelihood of the sample while increasing the conditional likelihood, which can be seen as improving the consistency of the samples generated with respect to the conditions.

[0120] Guidance without a classifier has a large effect and is similarly applied in the context of autoregressive models for image generation from text. One exemplary approach is to fine-tune the model while randomly replacing text prompts with padding tokens. During the inference period, a linear combination of logits sampled from the unconditional model (conditioned on padding tokens) and the model conditioned on the original text prompt is taken.

[0121] Some exemplary implementations also apply guidance without a classifier to Babeldraw and it is found that it has a positive impact on the output quality, especially for difficult text prompts. Some exemplary implementations fine-tune the model for 100,000 steps and randomly replace text prompts with padding tokens. During the sampling period, some exemplary implementations sample from logits from a linear combination of unconditional and conditional logits.

[0122] Some exemplary implementations sample 16 images per text prompt. For each output, some exemplary implementations re-rank based on the alignment score of the image and text embeddings of the Contrastive Captioners (CoCa) model. The CoCa base-size model can be trained on the same dataset. It should be noted that re-ranking over small sets of images is computationally inexpensive in image generation from text.

[0123] Exemplary scaling techniques Some exemplary implementations scale models on CloudTPUv4 hardware for both training and inference using the GSPMD and Lingvo frameworks. GSPMD is an XLA compiler-based model partitioning system that treats a cluster of TPUs as a single virtual device and uses sharding annotations on several tensors to enable the compiler to automatically distribute data and instruct it to compute on thousands of devices.

[0124] Training. For both the 350M and 750M models, some exemplary implementations simply train the model with data parallelism. For the 3B model, some exemplary implementations use 4-way in-layer model parallelism and 128-way data parallelism. Partitioning a single dimension in each tensor is sufficient to scale the 3B model. The model weights can be partitioned over the number of feed-forward hidden dimensions and attention head dimensions, and the internal activation tensors of the feed-forward and attention layers can also be partitioned over the hidden and head dimensions, but the difference from Megatron-LM is that some exemplary implementations also fully partition the output activations of the feed-forward and attention layers on different dimensions (they do not have hidden / head dimensions). This approach results in ReduceScatter and AllGather communication patterns instead of AllReduce, which significantly reduces peak active memory.

[0125] As an example, Figure 6 depicts a diagram of an exemplary approach for scaling model training in accordance with an exemplary embodiment of the present disclosure. In particular, Figure 6 is a diagram of 4-way in-layer model parallelism with fully partitioned activations for scaling an exemplary 3B model training. Figure 6 shows a simplified transformer feed-forward layer (with the sequence dimension omitted), where each shading represents data on one device. Some exemplary implementations additionally use 128-way data parallelism.

[0126] One exemplary 20B model has 16 encoder layers and 64 decoder layers. The weight sizes in each layer are of medium magnitude (as opposed to being very large), which results in pipeline parallelism, a good option for scaling. A general pipeline wrapper layer is implemented to enable an exemplary implementation to define a one-stage program. This is later automatically converted into a multi-stage pipeline program. The wrapper layer uses vectorization and shift buffers to reduce pipelining for the tensor partitioning problem, and thus all low-level infrastructure can be reused for pipelining. There are two further advantages from the adoption of GSPMD pipelines. (1) This enables an exemplary implementation to conveniently configure a pipeline within the lower components of the model, thereby simplifying the overall complexity of the encoder-decoder model. And, (2) since pipelining is implemented as tensor partitioning on a vectorized program, an exemplary implementation can reuse the same set of devices for other types of parallelism outside the transformer layer.

[0127] As an example, FIG. 7 depicts a diagram of an exemplary model parallel processing approach in accordance with an exemplary embodiment of the present disclosure. In particular, FIG. 7 depicts a diagram of an exemplary 16-stage GSPMD pipeline for scaling an exemplary 20B model training. FIG. 7 shows how 16 devices are used for data parallelism in the quantizer, embedding, and softmax layers, but are repurposed for pipelining in the encoder and decoder layers. Each shading represents the data or layer assigned to one device. The decoder uses a 4-round circular schedule to further reduce the pipeline bubble ratio. Additionally, the exemplary implementation uses 64-way data parallelism in addition for all layers.

[0128] Some exemplary implementations are configured such that the model has separate encoder and decoder pipelines, each having 16 stages. Some exemplary implementations also use 64-way data parallelism in addition to pipelining to speed up training. However, this results in a smaller per-core batch size and exposes the additional problem of excessive pipeline bubbles. To reduce the ratio of pipeline bubbles, some exemplary implementations adapt a circular schedule in the decoder pipeline such that four layers are executed in round-robin order at each stage. Outside of the encoder and decoder, some exemplary implementations perform data parallelization instead of pipelining for the embedding, softmax, and image tokenizer layers using the same set of devices.

[0129] During training, the Adafactor optimizer can be used to save memory with beta1 = 0.9, beta2 = 0.96, and a decoupled weight decay value of 4.5e-2. The first moment of the optimizer slot variable is additionally quantized from float32 to int8. Some exemplary implementations use a default dropout rate of 0.1 for all models in both the encoder and decoder. A deterministic version of the dropout layer can be used in the 20B model to enable model pipelining. Some exemplary implementations use an exponential learning rate schedule with a default learning rate of 4.5e-5 and a warmup step of 5000. The exponential decay starts at training step 85000 with a total of 450000 steps and a final decay rate of 0.025. Some exemplary implementations do not use the exponential moving average of the model weights to save device memory. In text encoding, some exemplary implementations build a sentence-piece model with a vocabulary size of 16000 on a subset of the training data. Conv-shaped sparse attention can be used in the decoder transformer. Some exemplary implementations additionally clip the gradient norm to a value of 4.0, especially to stabilize training at the start. Some exemplary implementations apply an additional layer regularization layer to the outputs of both the encoder and decoder.

[0130] Inference. One main goal for inference optimization is to speed up small batch image generation. Some exemplary implementations select in-layer model parallelism for both the 3B model and the 20B model. Contrary to training, some exemplary implementations do not fully partition the output activations for the feed-forward and attention layers for inference. This is because (1) each step of autoregressive decoding generates much smaller tensors and AllReduce performs better on small data (when writing), and (2) the activation memory is not relevant during inference and there is no handoff backward.

[0131] Exemplary Devices and Systems FIG. 8A depicts a block diagram of an exemplary computing system 100 in accordance with an exemplary embodiment of the present disclosure. The system 100 includes a user computing device 102, a server computing system 130, and a training computing system 150 communicatively coupled via a network 180.

[0132] The user computing device 102 can be any type of computing device, such as, for example, a personal computing device (e.g., a laptop or desktop), a mobile computing device (e.g., a smartphone or tablet), a gaming console or controller, a wearable computing device, an embedded computing device, or any other type of computing device.

[0133] The user computing device 102 includes one or more processors 112 and a memory 114. The one or more processors 112 can be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.) and can be one processor or multiple processors operably connected. The memory 114 can include one or more non-transitory computer-readable storage media such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, and combinations thereof. The memory 114 can store data 116 and instructions 118. The instructions 118 are executed by the processor 112 to cause the user computing device 102 to perform operations.

[0134] In some implementations, the user computing device 102 can store or include one or more machine learning models 120. For example, the machine learning model 120 can be various machine learning models such as a neural network (e.g., a deep neural network) or other types of machine learning models including non-linear models and / or linear models, or can include them. The neural network can include a feed-forward neural network, a recurrent neural network (e.g., a long short-term memory recurrent neural network), a convolutional neural network, or other forms of neural networks. Some exemplary machine learning models can utilize attention mechanisms such as self-attention. For example, some exemplary machine learning models can include a multi-head self-attention model (e.g., a transformer model). The exemplary machine learning model 120 is discussed with reference to FIGS. 1 and 2.

[0135] In some implementations, one or more machine learning models 120 are received from the server computing system 130 via the network 180, stored in the user computing device memory 114, and then can be used or implemented by one or more processors 112. In some implementations, the user computing device 102 can implement multiple parallel instances of a single machine learning model 120 (e.g., to perform parallel image quantization across multiple instances of an image).

[0136] Additionally or alternatively, one or more machine learning models 140 can be included in, or otherwise stored and implemented by, a server computing system 130 that communicates with a user computing device 102 according to a client-server relationship. For example, the machine learning model 140 can be implemented by the server computing system 140 as part of a web service (e.g., an image quantization, understanding, and / or generation service). Accordingly, one or more models 120 can be stored and implemented on the user computing device 102, and / or one or more models 140 can be stored and implemented on the server computing system 130.

[0137] The user computing device 102 can also include one or more user input components 122 that receive user input. For example, the user input component 122 can be a contact sensing component (a contact sensing display screen or a touch pad) that senses the contact of a user input object (e.g., a finger or a stylus). The contact sensing component can act to implement a virtual keyboard. Other exemplary user input components include a microphone, a traditional keyboard, or other means by which a user can provide user input.

[0138] Server computing system 130 includes one or more processors 132 and a memory 134. The one or more processors 132 may be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.), and may be a single processor or multiple processors operably connected. The memory 134 can include one or more non-transitory computer-readable storage media such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, and combinations thereof. The memory 134 can store data 136 and instructions 138. The instructions 138 are executed by the processor 132 to cause the server computing system 130 to perform operations.

[0139] In some implementations, the server computing system 130 includes one or more server computing devices, or alternatively, is implemented by one or more server computing devices. In cases where the server computing system 130 includes multiple server computing devices, such server computing devices can operate according to a sequential computing architecture, a parallel computing architecture, or some combination thereof.

[0140] As described above, the server computing system 130 may store, or alternatively may include, one or more machine learning models 140. For example, the model 140 may be various machine learning models, or alternatively may include them. Exemplary machine learning models include neural networks or other multi-layer non-linear models. Exemplary neural networks include feedforward neural networks, deep neural networks, regression neural networks, and convolutional neural networks. Some exemplary machine learning models can utilize attention mechanisms such as self-attention. For example, some exemplary machine learning models can include a multi-head self-attention model (e.g., a transformer model). The exemplary model 140 is discussed with reference to FIGS. 1 and 2.

[0141] The user computing device 102 and / or the server computing system 130 can train the model 120 and / or 140 through cross-reference using a training computing system 150 communicatively coupled via the network 180. The training computing system 150 may be separate from the server computing system 130 or may be part of the server computing system 130.

[0142] The training computing system 150 includes one or more processors 152 and a memory 154. The one or more processors 152 may be any suitable processing device (e.g., a processor core, a microprocessor, an ASIC, an FPGA, a controller, a microcontroller, etc.), and may be one processor or multiple processors operably connected. The memory 154 can include one or more non-transitory computer-readable storage media such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, and combinations thereof. The memory 154 can store data 156 and instructions 158. The instructions 158 are executed by the processor 152 to cause the training computing system 150 to perform operations. In some implementations, the training computing system 150 includes one or more server computing devices or, alternatively, is implemented by one or more server computing devices.

[0143] The training computing system 150 can include a model trainer 160 that trains a machine learning model 120 and / or 140 stored in the user computing device 102 and / or the server computing system 130 using various training or learning techniques such as, for example, backpropagation of error. For example, a loss function can be backpropagated through the model to update one or more parameters of the model (e.g., based on the gradient of the loss function). Various loss functions can be used, such as mean squared error, likelihood loss, cross-entropy loss, hinge loss, and / or various other loss functions. Gradient descent techniques can be used to repeatedly update the parameters over several training iterations.

[0144] In some implementations, the step of performing backpropagation of errors can include the step of performing truncated backpropagation over time. The model trainer 160 can implement some generalization techniques (such as weight decay, dropout, etc.) to improve the generalization ability of the model being trained.

[0145] In particular, the model trainer 160 can train the machine learning models 120 and / or 140 based on the set of training data 162. The training data 162 can include, for example, unsupervised and / or supervised training images.

[0146] In some implementations, if the user gives consent, the training examples can be provided by the user's computing device 102. Thus, in such implementations, the model 120 provided to the user's computing device 102 can be trained by the training computing system 150 with user-specific data received from the user computing device 102. In some cases, this process can be referred to as the step of personalizing the model.

[0147] The model trainer 160 includes computer logic utilized to implement the desired functionality. The model trainer 160 can be implemented in hardware, firmware, and / or software that controls a general-purpose processor. For example, in some implementations, the model trainer 160 includes a program file stored in a storage device and loaded into memory for execution by one or more processors. In other implementations, the model trainer 160 includes one or more sets of computer-executable instructions stored on a tangible computer-readable storage medium such as RAM, a hard disk, or an optical or magnetic medium.

[0148] Network 180 can be any type of communication network, such as a local area network (e.g., an intranet), a wide area network (e.g., the Internet), or some combination thereof, and can include any number of wired or wireless links. Generally, communication on network 180 can be carried via any type of wired and / or wireless connection using a variety of communication protocols (e.g., TCP / IP, HTTP, SMTP, FTP), encodings or formats (e.g., HTML, XML), and / or protection schemes (e.g., VPN, secure HTTP, SSL).

[0149] The machine learning models described herein can be used in various tasks, applications, and / or use cases.

[0150] In some implementations, the input to the machine learning model of the present disclosure may be image data. The machine learning model can process the image data to generate an output. By way of example, the machine learning model can process the image data to generate an image recognition output (e.g., recognition of the image data, latent embedding of the image data, encoded representation of the image data, hash of the image data, etc.). As another example, the machine learning model can process the image data to generate an image segmentation output. As another example, the machine learning model can process the image data to generate an image classification output. As another example, the machine learning model can process the image data to generate an image data transformation output (e.g., modification of the image data, etc.). As another example, the machine learning model can process the image data to generate an encoded image data output (e.g., encoded and / or compressed representation of the image data, etc.). As another example, the machine learning model can process the image data to generate a high-quality image data output. As another example, the machine learning model can process the image data to generate a prediction output.

[0151] In some implementations, the input to the machine learning model of the present disclosure may be latent encoded data (e.g., the latent space representation of the input). The machine learning model can process the latent encoded data to generate an output. As an example, the machine learning model can process the latent encoded data to generate a recognition output. As another example, the machine learning model can process the latent encoded data to generate a reconstruction output. As another example, the machine learning model can process the latent encoded data to generate a search output. As another example, the machine learning model can process the latent encoded data to generate a reclustering output. As another example, the machine learning model can process the latent encoded data to generate a prediction output.

[0152] In some cases, the machine learning model can be configured to perform a task that includes encoding the input data for reliable and / or efficient transmission or storage (and / or corresponding decoding). For example, the task may be an audio compression task. The input may include audio data, and the output may include compressed audio data. In another example, the input includes visual data (e.g., one or more images or videos), the output includes compressed visual data, and the task is a visual data compression task. In another example, the task can include generating an embedding for the input data (e.g., input audio or visual data).

[0153] In some cases, the input includes visual data and the task is a computer vision task. In some cases, the input includes pixel data for one or more images and the task is an image processing task. For example, the image processing task may be image classification, the output being a set of scores, each score corresponding to a different object class and representing the likelihood that one or more images depict an object belonging to the object class. The image processing task may be object detection, the image processing output identifying one or more regions in one or more images and, for each region, the likelihood that the region depicts the object of interest. As another example, the image processing task may be image segmentation, the image processing output defining, for each pixel in one or more images, the respective likelihoods for each category in a predefined set of categories. For example, the set of categories may be foreground and background. As another example, the set of categories may be object classes. As another example, the image processing task may be depth estimation, the image processing output defining, for each pixel in one or more images, the respective depth values. As another example, the image processing task may be motion estimation, the network input including a plurality of images, the image processing output defining, for each pixel of one of the input images, the motion of the scene depicted by the pixel between the images in the network input.

[0154] FIG. 8A illustrates one exemplary computing system that can be used to implement the present disclosure. Other computing systems can be used as well. For example, in some implementations, user computing device 102 can include model trainer 160 and training data set 162. In such implementations, model 120 can be both locally trained and used on user computing device 102. In some of such implementations, user computing device 102 can implement model trainer 160 to customize model 120 based on user-specific data.

[0155] FIG. 8B depicts a block diagram of an exemplary computing device 10 implemented in accordance with an exemplary embodiment of the present disclosure. The computing device 10 may be a user computing device or a server computing device.

[0156] The computing device 10 includes several applications (e.g., applications 1 through N). Each application includes its own machine learning library and machine learning model. For example, each application may include a machine learning model. Exemplary applications include a text messaging application, an email application, a dictation application, a virtual keyboard application, a browser application, and the like.

[0157] As illustrated in FIG. 8B, each application can communicate with several other components of the computing device, such as, for example, one or more sensors, a context manager, a device state component, and / or additional components. In some implementations, each application can communicate with each device component using an API (e.g., a public API). In some implementations, the API used by each application is dedicated to that application.

[0158] FIG. 8C depicts a block diagram of an exemplary computing device 50 implemented in accordance with an exemplary embodiment of the present disclosure. The computing device 50 may be a user computing device or a server computing device.

[0159] Computing device 50 includes several applications (e.g., applications 1 through N). Each application communicates with a central intelligence layer. Exemplary applications include a text messaging application, an email application, a dictation application, a virtual keyboard application, a browser application, and the like. In some implementations, each application can communicate with the central intelligence layer (and the models stored therein) using an API (e.g., a common API across all applications).

[0160] The central intelligence layer includes several machine learning models. For example, as shown in FIG. 8C, each machine learning model can be provided for each application and managed by the central intelligence layer. In other implementations, two or more applications may share a single machine learning model. For example, in some implementations, the central intelligence layer can provide a single model for all of the applications. In some implementations, the central intelligence layer is included within the operating system of computing device 50 or is otherwise implemented by the operating system of computing device 50.

[0161] The central intelligence layer can communicate with a central device data layer. The central device data layer can be a centralized repository of data for computing device 50. As shown in FIG. 8C, the central device data layer can communicate with several other components of the computing device, such as, for example, one or more sensors, a context manager, a device state component, and / or additional components. In some implementations, the central device data layer can communicate with each device component using an API (e.g., a private API).

[0162] Further Disclosure The techniques discussed in this specification act and send information between such systems as well as make references to servers, databases, software applications, and other computer-based systems. The inherent flexibility of computer-based systems allows for a wide variety of possible configurations, combinations, and partitions of tasks and functionality among components. For example, the processes discussed in this specification can be implemented using a single device or component, or multiple devices or components working in combination. Databases and applications can be implemented on a single system or distributed across multiple systems. The components being distributed can operate sequentially or in parallel.

[0163] While the subject matter has been described in detail with respect to various specific exemplary implementations, each example is provided as an illustration and not a limitation of the disclosure. Those skilled in the art will readily be able to create modifications, variations, and equivalents to such implementations given the above understanding. Accordingly, the disclosure of the subject matter is not intended to exclude such modifications, variations, and / or additional forms of the subject matter that will readily be apparent to those skilled in the art. For example, features described or depicted as part of one implementation can be used with another implementation to create further implementations. Accordingly, the disclosure is intended to cover such modifications, variations, and equivalents.

Description of the Reference Numerals

[0164] 10 Computing Device 12 Input Image Patch 16 Transformer Encoder, Image Encoder 18 Image Token 19 Quantization Code 22 Codebook, Quantization Codebook, Candidate Code 24 Transformer Decoder, Image Decoder 26 Synthesized Image Patch 28 Synthetic Image 30 Loss Function 50 Computing Device 102 User Computing Device 112 Processor 114 Memory 116 Data 118 Instruction 120 Machine Learning Model 122 User Input Component 130 Server Computing System 132 Processor 134 Memory 136 Data 138 Instruction 140 Machine Learning Model 150 Training Computing System 152 Processor 154 Memory 156 Data 158 Instruction 160 Model Trainer 162 Training Data, Training Dataset 180 Network 212 Input Image Patch 214 Input Image 216 Transformer Encoder, Image Encoder 218 Image Token 222 Codebook 224 Code Prediction Model 226 Predicted Code 228 Transformer Encoder, Image Decoder 230 Synthetic Image Patch 232 Synthetic Image 332 Synthetic Image 334 Natural Language Input, Input Text 402 Text Encoder 404 Autoregressive Code Selector 406 Text Embedding

Claims

1. A computer-implemented method for performing vector quantization of an image, comprising: obtaining, by a computing system comprising one or more computing devices, a plurality of input image patches of the 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 the machine-learned image encoder performs one or more self-attention operations to process the plurality of input image patches and generate the plurality of image tokens in the latent space; mapping, by the computing system, the plurality of image tokens to a plurality of quantization codes included in a quantization codebook including a plurality of candidate codes, wherein the step of mapping, by the computing system, the plurality of image tokens to the plurality of quantization codes included in the quantization codebook including the plurality of candidate codes includes, prior to the step of mapping the plurality of image tokens, projecting, by the computing system, the plurality of image tokens into a lower-dimensional space and / or applying L2 regularization to one or both of the plurality of image tokens and the plurality of candidate codes; providing, by the computing system, the plurality of quantization codes as an encoded version of the image; A computer-implemented method comprising the steps of.

2. The computer-implemented method according to claim 1, wherein the machine-learned image encoder comprises a vision transformer model.

3. The computer-implemented method according to claim 1, wherein the machine-learned image encoder performs one of the one or more self-attention operations on the plurality of input image patches.

4. processing, by the computing system, the plurality of quantization codes with a machine-learned image decoder to generate a plurality of synthetic image patches that form a synthetic image; evaluating, by the computing system, a loss function that provides a loss value based at least in part on the composite image; modifying, by the computing system, one or more of the machine-learned image encoder, the machine-learned image decoder, and one or more of the plurality of candidate codes based at least in part on the loss function; The computer-implemented method of claim 1, further comprising. **Claim 5** The computer-implemented method of claim 4, wherein the machine-learned image decoder comprises a vision transformer model. **Claim 6** The loss function is a logit Laplace loss term, an L2 loss term, a perceptual loss term, and an adversarial generative network loss term The computer-implemented method of claim 4, comprising. **Claim 7** predicting, by the computing system using a machine-learned code prediction model, a plurality of predicted codes autoregressively from the quantization codebook based at least in part on one or more of the plurality of quantization codes; generating, by the computing system, a plurality of synthetic image patches that form a synthetic image by processing the plurality of predicted codes with a machine-learned image decoder; The computer-implemented method of claim 1, further comprising. **Claim 8** evaluating, by the computing system, a code prediction loss function that evaluates a negative log-likelihood based on the plurality of predicted codes; modifying, by the computing system, one or more parameters of the machine-learned code prediction model based on the code prediction loss function; The computer-implemented method of claim 7, further comprising. **Claim 9** The computer-implemented method of claim 7, wherein the step of predicting, by the computing system using a machine-learned code prediction model, the plurality of predicted codes autoregressively comprises conditioning, by the computing system, the machine-learned code prediction model with auxiliary conditioning data that describes one or more desired characteristics of the synthetic image. **Claim 10** The computer-implemented method according to claim 9, wherein the auxiliary conditioned data includes a class label describing a desired class of the synthetic image.

11. The computer-implemented method according to claim 9, wherein the auxiliary conditioned data includes natural language text tokens.

12. The step of conditioning the machine-learned code prediction model by the computing system with the natural language text tokens comprises: processing the natural language text tokens by a text encoder part of the machine-learned code prediction model by the computing system to generate a text embedding; and providing the text embedding as an input to an autoregressive code selection part of the machine-learned code prediction model by the computing system to autoregressively predict the plurality of predicted codes. The computer-implemented method according to claim 11.

13. extracting one or more intermediate features from the machine-learned code prediction model by the computing system; and predicting a class label for the image by the computing system based at least in part on the intermediate features. The computer-implemented method according to claim 7, further comprising:

14. A computer-implemented method for performing vector quantization of an image, comprising: obtaining, by a computing system comprising one or more computing devices, a plurality of quantization codes forming an encoded version of the image, the plurality of quantization codes being selected by mapping a plurality of image tokens generated by a machine-learned image encoder model to the plurality of quantization codes included in a quantization codebook including a plurality of candidate codes, the plurality of image tokens being projected into a lower-dimensional space before the mapping and / or one or both of the plurality of image tokens and the plurality of candidate codes being L2-regularized before the mapping; and processing, by the computing system, the plurality of quantization codes with a machine-learned image decoder to generate a plurality of decoded image patches forming a decoded version of the image. Including A computer-implemented method, wherein one or both of the machine-learned image encoder model and the machine-learned image decoder are configured to perform one or more self-attention operations.

15. A computer system comprising one or more processors and one or more non-transitory computer-readable media that collectively store instructions which, when executed by the one or more processors, cause the one or more processors to perform the method according to any one of claims 1 to 14.

16. One or more processors and one or more non-transitory computer-readable media that collectively store a machine-learned image processing model, wherein the machine-learned image processing model comprises an encoder configured to encode one or more input image patches into one or more image tokens in a latent space, a quantization unit configured to quantize the one or more image tokens into one or more quantization codes selected from a quantization codebook, wherein the quantization unit is configured to project the one or more image tokens into a lower-dimensional space and / or apply L2 regularization to one or both of the one or more image tokens and the one or more candidate codes before mapping the one or more image tokens to one or more quantization codes included in the quantization codebook containing the one or more candidate codes, a quantization unit, a code prediction unit configured to predict one or more predicted quantization codes from the quantization codebook based at least in part on the one or more quantization codes, and a feature prediction unit configured to generate one or more characteristic predictions for the input image patch based at least in part on data extracted from the code prediction unit. A computing system comprising

17. The computing system according to claim 16, further comprising a decoder unit configured to generate an image patch reconstructed based on the one or more quantization codes, or to generate a synthetic image patch based at least in part on the one or more predicted quantization codes.

18. The computing system according to claim 16 or 17, wherein the one or more characteristic predictions include image classification predictions.

19. A computer-implemented method for training a code prediction model, comprising: obtaining, by a computing system comprising one or more computing devices, a training example including a training image and a natural language input describing the 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 quantization codes included in a quantization codebook including a plurality of candidate codes, wherein the step of mapping, by the computing system, the plurality of image tokens to the plurality of quantization codes included in the quantization codebook including the plurality of candidate codes includes, prior to the step of mapping the plurality of image tokens, projecting, by the computing system, the plurality of image tokens into a lower-dimensional space and / or applying L2 regularization to one or both of the plurality of image tokens and the plurality of candidate codes; processing, by the computing system, the natural language input with a text encoder unit of the code prediction model to generate a text embedding; processing, by the computing system, the text embedding with an autoregressive code selection unit of the code prediction model to autoregressively predict a sequence of codes predicted from the quantization codebook; evaluating, by the computing system, a code prediction loss function that compares the predicted sequence of codes to the plurality of quantized codes; changing, by the computing system, one or more parameters of the code prediction model based at least in part on the code prediction loss function; A computer-implemented method comprising: **Claim 20** The computer-implemented method of claim 19, wherein the machine-learned image encoder performs one or more self-attention operations to process a plurality of input image patches and generate the plurality of image tokens in the latent space. **Claim 21** The computer-implemented method of claim 19, wherein the code prediction loss function evaluates the negative log-likelihood of the predicted codes with respect to the quantized codes. **Claim 22** The computer-implemented method according to any one of claims 19 to 21, wherein the machine-learned image encoder and the quantized codebook are pre-trained with an image decoder in an autoencoder architecture. **Claim 23** The computer-implemented method of claim 19, wherein the step of changing, by the computing system, one or more parameters of the code prediction model based at least in part on the code prediction loss function comprises the computing system changing one or more parameters of both the text encoder part and the autoregressive code selection part of the code prediction model based at least in part on the code prediction loss function.

Images