Vector-quantized image modeling

Abstract

This invention provides a vector quantized image modeling method and system. [Solution] The method provides a vector quantized image modeling (VIM) technique that includes the step of pre-training a machine learning model (e.g., a transformer model) to autoregressively predict rasterized image tokens. Individual image tokens are encoded from a trained vision transformer-based VQGAN (ViT-VQGAN). [Effects] This paper proposes multiple improvements to vanilla VQGAN, from architecture to codebook learning, resulting in better performance and reconstruction fidelity. The improved ViT-VQGAN further enhances vector quantized image modeling tasks, including 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 on June 10, 2022, and U.S. Provisional Patent Application No. 63 / 252,452, filed on October 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 focus on applying different random augmentations to images during pre - training for the purpose of differentiating 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 products can only be applied to relatively small maximum resolutions (e.g., 64x64) for image recognition, thereby significantly 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 automatically encode images, and a second-stage CNN or transformer learns to model the density of the input data. While these have proven effective for image generation, the learned representations in characteristic tasks have not been evaluated in research.

[0006] One subfield of image generation is the generation of images from text. Here, images are synthesized based on input text, and the synthesized images depict the content described by the input text. The ability to generate images through language is fascinating, as language is the most natural form of communication, and such generative ability can potentially unlock creative applications in many fields, including art, design, and multimedia content creation. One lineage of research that has recently gained momentum in the space of generating images from text is a technique that leverages diffusion-based image-from-text models that have the ability to generate higher-fidelity images. These models avoid the use of individual image tokens as the core modeling technique for image statics, and instead use diffusion models to achieve better zero-shot Fréchet-start distance (FID) scores for MS-COCO and aesthetically pleasing visual output. Despite these advances, the ability to apply autoregressive modeling to the task of generating images from text remains intriguing, assuming that there is a rich body in conventional works for large-scale language models and general interfaces for individual tokens that are easily applicable to a wider range of modalities. [Overview of the Initiative] [Means for solving the problem]

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

[0008] One exemplary aspect of this disclosure is directed to a computer-implemented method for performing vector quantization of an image. The method includes the step of obtaining a plurality of input image patches of an image by a computing system comprising one or more computing devices. The method includes the step of processing the plurality of input image patches with a machine learning-trained image encoder by the computing system to generate a plurality of image tokens in latent space, where the plurality of image tokens correspond to the plurality of input image patches, and the machine learning-trained 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 latent space. The method includes the step of mapping the plurality of image tokens to a plurality of quantization codes contained in a quantization codebook containing a plurality of candidate codes. The method includes the step of providing the plurality of quantization codes as encoded versions of the image by the computing system.

[0009] In some implementations, machine learning-based image encoders include a vision transformer model.

[0010] In some implementations, a machine learning-based image encoder performs one or more self-attention actions on multiple input image patches.

[0011] In some implementations, the method further includes the steps of: a computing system processing multiple quantization codes with a machine learning image decoder to generate multiple composite image patches that form a composite image; a computing system evaluating a loss function that provides a loss value based at least partially on the composite image; and a computing system modifying a machine learning image encoder, a machine learning image decoder, and one or more of several candidate codes at least partially based on the loss function.

[0012] In some implementations, the machine learning-based image decoder includes 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 a generative adversarial network loss term.

[0014] In some implementations, the steps of a computing system mapping multiple image tokens to multiple quantization codes contained in a quantization codebook containing multiple candidate codes include the steps of a computing system projecting the multiple image tokens into a lower-dimensional space, and, after the step of projecting the image tokens into the lower-dimensional space, the computing system mapping the multiple image tokens to multiple quantization codes contained in the quantization codebook.

[0015] In some implementations, the steps include: a computing system mapping multiple image tokens to multiple quantization codes contained in a quantization codebook containing multiple candidate codes; a computing system applying L2 regularization to one or both of the multiple image tokens and the multiple candidate codes; and, after the step of applying L2 regularization, a computing system mapping multiple image tokens to multiple quantization codes contained in a quantization codebook.

[0016] In some implementations, the method further includes the steps of: using a machine learning-trained code prediction model, a computing system autoregressively predicting multiple predicted codes from a quantization codebook based at least partially on one or more of multiple quantization codes; and the computing system processing the multiple predicted codes with a machine learning-trained image decoder to generate multiple composite image patches that form a composite image.

[0017] In some implementations, the method further includes the steps of: having a computing system evaluate a code prediction loss function that evaluates a negative log-likelihood based on a set of predicted codes; and having the computing system modify one or more parameters of a machine learning code prediction model based on the code prediction loss function.

[0018] In some implementations, the step of a computing system using a machine learning code prediction model autoregressively predicting multiple predicted codes includes the step of the computing system conditioning the machine learning code prediction model with auxiliary conditioning data that describes one or more desired properties of a composite image.

[0019] In some implementations, the auxiliary conditioning data includes class labels that describe the desired class of the composite image.

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

[0021] In some implementations, the method further includes the steps of: using a computing system to extract one or more intermediate features from a machine learning code prediction model; and using a computing system to predict a class label for an image based at least partially 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 that collectively store a machine-learned image processing model. The machine-learned image processing model comprises an encoder portion configured to encode one or more input image patches into one or more image tokens in a latent space, a quantization portion configured to quantize one or more image tokens into one or more quantization codes selected from a codebook, a code prediction portion configured to predict one or more predicted quantization codes from the codebook based at least in part on one or more quantization codes, and a feature prediction portion 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 portion.

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

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

[0028] Another exemplary embodiment is directed to a computer-implemented method for performing image generation from text. The method includes the step of obtaining a natural language input describing desired image content by a computing system comprising one or more computing devices. The method includes the step of processing the natural language input with a text encoder unit of a machine learning code prediction model by the computing system to generate text embeddings. The method includes the step of processing the text embeddings with an autoregressive code selection unit of a machine learning code prediction model by the computing system to autoregressively predict a sequence of codes to be predicted from a quantized codebook containing a plurality of candidate codes. The method includes the step of processing the sequence of quantized codes with a machine learning image decoder by the computing system to generate a plurality of composite image patches that form a composite image. The composite image depicts the desired image content.

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

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

[0031] In some implementations, either or both of the machine learning-trained image decoder and / or codebook were trained together with the image encoder model.

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

[0033] Another exemplary embodiment is directed to a computer-implemented method for training a code prediction model. The method includes the step of obtaining a natural language input describing training examples, including training images, and the content of the training images, by a computing system comprising one or more computing devices. The method includes the step of processing multiple image patches from the training images with a machine learning-trained image encoder by the computing system to generate multiple image tokens in a latent space, where the multiple image tokens correspond to multiple image patches. The method includes the step of mapping the multiple image tokens to multiple quantized codes contained in a quantized codebook, which contains multiple candidate codes. The method includes the step of processing the natural language input with a text encoder unit of the code prediction model by the computing system to generate text embeddings. The method includes the step of processing the text embeddings with an autoregressive code selection unit of the code prediction model by the computing system to autoregressively predict a sequence of codes to be predicted from the quantized codebook. The method includes the step of evaluating a code prediction loss function by the computing system that compares the predicted sequence of codes with multiple quantized codes. The method includes the step of modifying one or more parameters of the code prediction model by the computing system, at least in part, based on the code prediction loss function.

[0034] In some implementations, a machine learning-based image encoder processes multiple input image patches by performing one or more self-attention actions, generating multiple image tokens in the latent space.

[0035] In some implementations, the code prediction loss function evaluates the negative log-likelihood of the predicted code for the quantized code.

[0036] In some implementations, machine learning-trained image encoders and codebooks are pre-trained on image decoders in an autoencoder architecture.

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

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

[0039] These and other features, aspects, and advantages of various embodiments of this disclosure will be better understood by referring to the following description and the appended claims. The appended drawings, incorporated herein and forming part thereof, illustrate exemplary embodiments of this disclosure, along with helpful descriptions illustrating the relevant principles.

[0040] A detailed discussion of embodiments for those skilled in the art is provided herein, with reference to the accompanying drawings. [Brief explanation of the drawing]

[0041] [Figure 1] This figure illustrates an exemplary machine learning model for performing vector quantization of an image, according to an exemplary embodiment of the present disclosure. [Figure 2] This figure illustrates an exemplary machine learning model for performing image generation, according to an exemplary embodiment of the present disclosure. [Figure 3] This figure illustrates an exemplary machine learning model for performing image generation, according to an exemplary embodiment of the present disclosure. [Figure 4] This figure illustrates an exemplary machine learning model for performing image generation, according to an exemplary embodiment of the present disclosure. [Figure 5]This is a schematic diagram illustrating a method for performing image generation from text, according to an exemplary embodiment of the present disclosure. [Figure 6] This figure illustrates an exemplary method for scaling model training, according to an exemplary embodiment of the present disclosure. [Figure 7] This figure illustrates an exemplary model parallel processing method according to an exemplary embodiment of the present disclosure. [Figure 8A] This is a block diagram illustrating an exemplary computing system according to an exemplary embodiment of the present disclosure. [Figure 8B] This is a block diagram illustrating an exemplary computing device according to an exemplary embodiment of the present disclosure. [Figure 8C] This is a block diagram illustrating an exemplary computing device according to an exemplary embodiment of the present disclosure. [Modes for carrying out the invention]

[0042] The reference numbers that appear repeatedly across multiple diagrams are intended to identify the same features in various implementations.

[0043] overview Generally, this disclosure is directed toward vector quantized image modeling using vision transformers and improved codebook handling. In particular, this disclosure provides a vector quantized image modeling (VIM) technique that includes a step of 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 trained vision transformer-based VQGAN (an exemplary implementation of which may be called ViT-VQGAN). This disclosure proposes several improvements over vanilla VQGAN, from architecture to codebook training, resulting in better efficiency and reconstruction fidelity. The improved ViT-VQGAN further enhances vector quantized image modeling tasks, including unconditional image generation, conditional image generation (e.g., class-conditional image generation), and unsupervised representation learning.

[0044] The illustrative experiments demonstrate the technical effectiveness of the proposed technique. In particular, when trained on ImageNet at a 256x256 resolution, the illustrative implementation of this disclosure achieves a starting score (IS) of 175.1 and a Fréchet starting distance (FID) of 4.17, which is a dramatic improvement over vanilla VQGAN, which obtains 70.6 and 17.04 for the starting score (IS) and Fréchet starting distance (FID), respectively. Based on ViT-VQGAN and unsupervised pre-training, the illustrative experiments further evaluate the pre-trained transformer by averaging intermediate features, similar to Image GPT (iGPT). This ImageNet pre-trained VIM-L significantly outperforms iGPT-L in linear probe accuracy, from 60.3% to 72.1% at similar model sizes. ViM-L also outperforms iGPT-XL, which is trained with additional web image data and a larger model size.

[0045] Another exemplary aspect of this disclosure is directed toward the application of the proposed vector quantization image model to the task of generating images from text. Specifically, this disclosure provides an autoregressive model that can generate photorealistic images from text descriptions, which has advantages over diffusion-based methods. Exemplary implementations of these models may be referred to as "Babeldraw". In some implementations, Babeldraw encodes images as sequences of individual tokens using an improved image tokenizer, the transformer-based ViT-VQGAN described herein. This naturally reduces the task at hand to the well-known problem of machine conversion from text to image tokens. A second, consistent quality improvement can be obtained by scaling the encoder-decoder transformer parameters from 350M and 750M to 3B and 20B. Compared to recent diffusion-based models, with fine-tuning using MS-COCO training splits, the 20B Babeldraw model achieves a comparable zero-shot MS-COCO Flechet-start distance (FID) score of 7.31 and an unprecedented FID score of 4.03. Using a post-super-resolution upsampler trained on a frozen 256x256 image tokenizer, Babeldraw reliably generates high-resolution, photographically realistic images.

[0046] Therefore, the exemplary implementations of this disclosure relate to vector quantized image modeling (VIM) methods for both image generation and image understanding tasks. Specifically, several exemplary implementations adhere to the following two-step approach.

[0047] Stage 1: Vector Quantization. Assuming an image (e.g., with a resolution of 256x256), a 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). This disclosure proposes several improvements to VQGAN (Esser et al., 2021), from architecture to codebook training. The resulting ViT-VQGAN has better efficiency and reconstruction fidelity in terms of pixel-level reconstruction measures, start score (IS), and Fréchet start distance (FID). ViT-VQGAN can be thoroughly trained on image-only data using an objective function that combines logit-Laplace loss, L2 loss, adversarial loss, and / or perceptual loss.

[0048] Stage 2: Vector Quantized Image Modeling. Several exemplary implementations may include a 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 decoder-only transformer model 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 added 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 the class logit (also known as linear probe accuracy).

[0050] One key component for improving both image generation and image understanding in VIM is having a better image quantizer in terms of computer efficiency and reconstruction quality. An efficient quantizer can speed up Stage 2 training, where random augmentation is first applied to the image and then to the image quantizer's encoder as input tokens. Furthermore, an image quantizer with better reconstruction quality can potentially reduce information loss compared to the original image in pixel space, which is important for image understanding tasks.

[0051] Further exemplary embodiments of this disclosure are directed toward an autoregressive image generation model capable of generating photorealistic images from text descriptions. The autoregressive models provided herein can benefit from better image tokenizers and model scaling. More specifically, this disclosure provides a transformer-based sequence-to-sequence model for image generation (an exemplary implementation of which may be called "Babeldraw"). The image generation model can take tokenized text tokens as input to an encoder and predict individual image tokens using a decoder in an autoregressive manner. In some examples, the image tokens can be generated by the transformer-based ViT-VQGAN image tokenizer described herein, which can produce a model that is better than VQGAN in terms of efficiency and image reconstruction fidelity.

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

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

[0054] Thus, this disclosure makes several contributions to image generation from text. First, it is understood that the image tokenizer is one of the key factors for a two-stage image generation from text model and is directly related to the image quality of the generated image. 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 with BERT alone provides a mild assistance to natural language comprehension for a text-to-image system. Text encoders fine-tuned for image generation from text perform worse on language comprehension tasks, suggesting a potential gap between general language comprehension and vision-based language comprehension.

[0055] The systems and methods of this disclosure offer several technical effects and advantages. As one example, the proposed technique can achieve equivalent results to current state-of-the-art technical models (e.g., iGPT) with smaller models and less data for certain tasks (e.g., image recognition with generative pretraining). Enabling the use of smaller models with less data (while achieving equivalent results) reduces the consumption of computing resources such as processor usage, memory usage, and network bandwidth. Thus, the proposed technique conserves computing resources and improves the functionality of computers.

[0056] As another example, exemplary implementations of this disclosure exhibit superior efficiency and reconstruction fidelity compared to existing models in terms of pixel-level reconstruction measurement, start score (IS), and Fréchet start distance (FID). Thus, this disclosure represents an improvement in the ability to encode image data into encoded representations. Encoded data offers numerous advantages, including savings in memory space and network bandwidth. This disclosure enables a system to encode images to achieve such advantages while simultaneously reducing the data loss experienced when decoding the images back into decoded images.

[0057] As another exemplary technical effect and benefit, this disclosure enables improved image generation and image understanding using VIM by providing an improved image quantizer in terms of computer efficiency and reconstruction quality. An efficient quantizer can speed up Stage 2 training, where random augmentation is first applied to the image and then to the image quantizer's encoder as input tokens. Faster training speed results in savings of computing resources such as processor usage, memory usage, and / or network bandwidth usage. Furthermore, an image quantizer with better reconstruction quality can potentially reduce information loss compared to the original image in pixel space, which is important for image understanding tasks. Thus, the computer performance for image understanding tasks can be improved.

[0058] Exemplary embodiments of this disclosure will be discussed in further detail here with reference to the figures.

[0059] Exemplary Stage 1 technique Figure 1 illustrates an exemplary machine learning model for performing vector quantization of an image, according to an exemplary embodiment of the present disclosure.

[0060] Specifically, as shown in Figure 1, the computing system can obtain multiple input image patches 12 of image 14.

[0061] The computing system can process multiple input image patches 12 with a machine learning-trained image encoder 16 to generate multiple image tokens 18 in the latent space. The multiple image tokens 18 may correspond to multiple input image patches 12. The machine learning-trained image encoder 16 can process the multiple input image patches 12 by performing one or more self-attention operations to generate multiple image tokens 18 in the latent space.

[0062] The computing system can map multiple image tokens 18 to multiple quantization codes 20 contained in a quantization codebook 22 which contains multiple candidate codes. The computing system can provide multiple quantization codes 20 as encoded versions of the image.

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

[0064] The computing system can evaluate a loss function 30 that provides loss values ​​based at least partially on a synthesized image. The loss function includes a logit-Laplace loss term, an L2 loss term, a perceptual loss term, and / or a generative adversarial network loss term (for example, the GAN evaluation is specifically illustrated in 30, while others are not).

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

[0066] In some implementations, the step of mapping multiple image tokens 18 to multiple quantization codes 20 may include the step of projecting the multiple image tokens 18 into a lower-dimensional space, and, after the step of projecting the image tokens 18 into the lower-dimensional space, the step of mapping the multiple image tokens to multiple quantization codes contained in a 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 multiple image tokens 18 to multiple quantization codes 20 may include the step of applying L2 regularization to one or both of the multiple image tokens 18 and the multiple candidate codes in the codebook 22, and, after the step of applying L2 regularization, the step of mapping the multiple image tokens 18 to the multiple quantization codes 20 contained in the quantization codebook 22.

[0068] Exemplary Stage 2 Technique Figure 2 illustrates an exemplary machine learning model for performing vector quantization of an image, according to an exemplary embodiment of the present disclosure.

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

[0070] A computing system can process one or more input image patches 212 with a machine learning-trained image encoder 216 to generate one or more image tokens 218 in latent space. Each of the one or more image tokens 218 can correspond to one or more input image patches 212. The machine learning-trained image encoder 216 can process one or more input image patches 212 by performing one or more self-attention operations to generate one or more image tokens 218 in latent space. For example, the machine learning-trained image encoder 216 may be trained or learned according to the process shown in Figure 1.

[0071] The computing system can map one or more image tokens 218 to one or more quantization codes 220 contained in a quantization codebook 222 which contains multiple candidate codes. For example, the quantization codebook 222 may be learned according to the process shown in Figure 1.

[0072] Furthermore, in some implementations, a computing system (i.e., a computing device of the same or different components) may process one or more quantization codes 220 with a machine learning code prediction model 224 to autoregressively predict a number of predicted codes 226 from a quantization codebook 222, based at least partially on one or more quantization codes 220.

[0073] The computing system can process multiple predicted codes 226 with a machine learning-trained image decoder 228 to generate multiple composite image patches 230 that form a composite image 232.

[0074] In some implementations or examples, the computing system can further evaluate a code prediction loss function that assesses the negative log-likelihood based on multiple predicted codes 226. The computing system can modify one or more parameters of a machine learning 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 across an image corpus, for example, using the code prediction loss function.

[0075] In some implementations, during the training and / or inference period, the machine learning-trained code prediction model 224 can be conditioned with auxiliary conditioning data that describes one or more desired characteristics of the synthesized image. As one example, the auxiliary conditioning data may include class labels that describe a desired class of the synthesized image 232. For example, in Figure 2, the model 224 is conditioned with the label [Cat], and as a result, the synthesized image 232 depicts a cat, but the synthesized image 232 shares visual characteristics with image 214. As another example, the auxiliary conditioning data may include natural language text tokens and / or any other conditioning data that gives priority to image synthesis.

[0076] Figure 2 illustrates an exemplary method in which the image generation process is conditioned both by a patch 212 from the input image 214 and by text input (e.g., a class label). However, in some implementations, image generation can be conditioned by text alone. For example, Figure 3 illustrates an example in which the generation of a composite image 332 is conditioned by the natural language input 334, "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] In some embodiments of this disclosure, in certain implementations that include text conditioning and / or perform image generation tasks from text, the code prediction model 224 may include a text encoder that generates the text to be embedded from the text input.

[0078] As an example, Figure 4 shows an exemplary image generation method in which the code prediction model 224 includes a text encoder 402 and an autoregressive code selector 404. The text encoder 402 can receive 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 may be a transformer model or other model that performs self-attention, or may include a transformer model or other model that performs self-attention. In some implementations, the text encoder can be pre-trained (e.g., using masked language modeling pre-training techniques, contrast 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 distribution of tokens over sets of text-and-image training examples, each containing a text-and-image pair.

[0080] Exemplary vector quantization 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 which is thoroughly trained by straight-through estimation. VQGAN is a model that improves upon VQVAE by introducing adversarial loss generated by a discriminator. Further improvements to VQGAN are achieved in this specification to improve efficiency and reconstruction quality.

[0081] VQGAN as an example in Vision Transformer The core network architecture used by both VQVAE and VQGAN for encoding and reconstructing images is a CNN. VQGAN introduces transformer-like elements in the form of nonlocalized attention blocks, which allows it to incorporate distance interactions in fewer layers. An exemplary implementation of this disclosure comes one step closer by replacing the CNN encoder and decoder with a vision transformer (ViT). Given sufficient data (a large amount of unlabeled image data), ViT-VQGAN has been found to be less constrained by the inductive prior imposed by convolution. Furthermore, ViT-VQGAN results in better computer efficiency for accelerators and produces higher quality reconstructions, as shown in Table 1 below.

[0082] [Table 1]

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

[0084] The exemplary ViT-VQGAN encoder first maps 8x8 non-overlapping image patches to image tokens, and then the transformer block encodes the 256x256 resolution image into a sequence of 32x32=1024 tokens. The exemplary decoder performs the reverse operation, mapping each image token back from the latent variables to the 8x8 image patches and regrouping them into a 256x256 image. In some implementations, a two-layer feedforward network with an intermediate tanh active layer is applied to 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 mean prediction due to the logit-Laplace loss). In some implementations, sigmoid activation is applied to the mean prediction of the decoder due to the logit-Laplace loss. This exemplary technique results in high-quality reconstruction without visible grid artifacts.

[0085] Exemplary Codebook Learning Vanilla VQVAE typically suffers from low codebook usage due to insufficient codebook initialization. Consequently, a significant portion of the code is hardly used or invalid during the training period. A reduced effective codebook size results in worse reconstruction during Stage 1 quantizer training and insufficient versatility during Stage 2 image synthesis. As a result, VQGAN relies on a top-k and top-p (nucleus) sampling discovery method with a default codebook size of 1024 to achieve the best results in image synthesis. Exemplary implementations of this disclosure include two improvements that can significantly improve codebook usage, even with a larger codebook size of 8192. During the image synthesis period, the exemplary implementations perform simple sampling at a temperature of 1.0 without the top-k and top-p discovery method.

[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 experiments here, and e is the codebook vector. In some implementations, the quantized codebook index is the input feature z in terms of Euclidean distance. e Determined by searching for the codebook vector closest to (x),

[0091]

number

[0092] This will result.

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

[0094] L2 regularization code. Several exemplary implementations involve encoding the latent variable z e(x) and the codebook latent variable e are also to be applied to l2 regularization. In some implementations, the codebook variable starts from a normal distribution. The l2 regularized latent variable is obtained by mapping all latent variables onto a sphere.

[0095]

number

[0096] The Euclidean distance between z is z e This can be extended to the cosine similarity of two vectors between (x) and e, further improving the training stability and reconstruction quality shown in our experiments.

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

[0098] One preliminary application for VGG-based perceptual loss states that the VGG network is pre-trained with supervised classification loss, and therefore, for linear probe accuracy measurement, surveillance may slip into stage 2. Consequently, for all reported unsupervised learning results, several exemplary implementations exclude the perceptual loss during the ViT-VQGAN training period. For all unconditional and class-conditional image synthesis, several exemplary implementations use a ViT-VQGAN quantizer trained with perceptual loss, resulting in higher fidelity reconstructions.

[0099] Exemplary vector quantization image modeling Using the trained ViT-VQGAN, images can be encoded into individual latent codes (id) flattened in raster order, similar to image GPT. Using a transformer model with only a decoder, the density P(x) of image data can be autoregressively modeled as shown in the following equation.

[0100]

number

[0101] Here, θ is a learnable weight. The training objective is to train the negative log-likelihood of the data.

[0102]

number

[0103] The goal is to minimize it.

[0104] Table 2 summarizes exemplary architectural configurations for the transformer. Several exemplary implementations first embed individual image token IDs at each position into a learnable embedding space, along with additional learnable 2D position embeddings. Both embedding dimensions are the same as the model dimensions. Several exemplary implementations apply a stack of transformer blocks to the input with causal attention across the entire sequence. A dropout rate of 0.1 can be used for all remaining outputs, activation outputs, and attention outputs. At the final layer of all transformer blocks, several exemplary implementations apply additional layer regularization.

[0105] [Table 2]

[0106] Exemplary image synthesis A pre-trained generative transformer model can achieve unconditional image generation by simply sampling token by token from an output softmax distribution. All samples used in both qualitative and quantitative results can be obtained without temperature degradation. The sampled tokens are then fed into a ViT-VQGAN decoder to decode the output image. An exemplary default Stage 1 ViT-VQGAN encodes a 256x256 input image into a 32x32 latent code with a codebook size of 8192, while a 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 a model for image synthesis. Several exemplary implementations extend unconditional generation to class-conditional generation by adding class ID tokens before image tokens. A separate embedding layer can learn about the class ID tokens and image tokens from the beginning, with the embedding dimension being the same as the transformer model dimension. During the sampling period, the class ID token can be placed in the first position, and the remaining image tokens can be decoded autoregressively.

[0108] Exemplary unsupervised learning In image understanding tasks, several exemplary implementations feed all input image tokens into a pre-trained transformer, yielding a sequence of 1024 token features. Some exemplary implementations take the layer output at a specific block l across a total of L blocks, average it over the (frozen) sequence of token features, and insert a (trainable) softmax layer to project the averaged features onto a class logit. Some exemplary implementations simply take one specific transformer block output instead of concatenated different block outputs, as in iGPT. Often, the most discriminative feature for a linear probe is typically located 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 spanning both language and image tokens.

[0110] Example image tokenizer Autoregressive text-image models rely on several forms of linearization of a 2D image to a 1D sequence of patch representations. In the limit, these are merely pixels, but they require modeling very long sequences, even for relatively small images (for example, a 256x256 image yields 65536 rasterized pixels). Worse still, it relies on a very low-level representation of the input rather than something richer, as indicated by the pixel locations in the image context. Many constructs have solved the above problem by using individual variational autoencoders to learn quantized representations of image patches across a set of raw images. Instead of learning representations that can take arbitrary values ​​in latent space, a visual codebook is learned to embed patches into its nearest codebook entry, which is a learned and indexable location throughout the entire latent space. These entries can be thought of as visual word types, and any representation of any of these words in a given image patch is therefore 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 balanced use of its entries across a wide range of images, and (b) support the reconstruction of sequences of visual tokens as high-quality output images. Several exemplary implementations use ViT-VQGAN, an example described herein, which addresses both requirements, namely, regularized and factorized code that contribute to training stability, reconstruction quality, and codebook usage, by employing techniques to improve codebook learning.

[0112] In Babeldraw's image tokenizer, several exemplary implementations train ViT-VQGAN as described above, but use images from 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, 512 model dimensions, and 2048 hidden dimensions) to learn 8192 image token classes for the codebook. Note that only the second stage encoder-decoder training relies on the encoder and codebook of the learned quantizer. 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 larger size decoder (32 blocks, 16 heads, 1280 model dimensions, and 5120 hidden dimensions). Both the input and output of the image tokenizer may be 256x256 resolution.

[0113] Some exemplary implementations, when zoomed in on a portion of an image, present a pixelated pattern (saturated pixel values) in the output of ViT-VQGAN. To address this issue, some exemplary implementations remove the final sigmoid layer, exposing 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, while 256x256 resolution images capture most of the content, structure, and texture, higher resolution images such as 512x512 or 1024x1024 provide greater visual satisfaction. For this purpose, several exemplary implementations leverage a simple super-resolution module on top of an image tokenizer. Stacked convolutional layers with the remaining connections can be used as a super-resolution network module. It can be trained with ViT-VQGAN with the same losses (perceptual loss, StyleGAN loss, and L2 loss) to map from reconstructed images to higher resolution reconstructed images. Note that the diffusion model can also be used here as an iterative refined super-resolution module, with or without conditioning, on text input.

[0115] Exemplary image from text in encoder decoder In some exemplary implementations, a standard encoder-decoder transformer model is trained in a second stage by treating text-to-image as a sequence-to-sequence modeling problem. The model takes text as input and is trained using next-token predictions of image latent codes generated from the image tokenizer in the first stage. During inference, the model can autoregressively sample image tokens, which are later decoded into pixels using the decoder of the previously trained image tokenizer.

[0116] Text prompts can be truncated to a maximum of 128, and images with a resolution of 256x256 can be encoded into tokens of length 1024 (32x32). Several exemplary implementations train with cross-entropy loss and use an int8 quantized Adafactor optimizer with a constantly decaying second moment factor (beta1=0.9, beta2=0.96). Data types can be downsampled to bfloat16 or attention projection and feedforward transformer layers, while all layer norms and model outputs can remain float32. All models can use conv-shaped mask sparse attention, and several exemplary implementations train four size variants ranging from 350 million to 20 billion parameters.

[0117] Most existing two-stage text-to-image generation models are decoder-only models. Early investigations on a scale of 350 million to 750 million parameters showed that Babeldraw's encoder-decoder variant outperformed decoder-only models in both training loss and text-to-image generation quality.

[0118] Exemplary text encoder pre-training The encoder-decoder architecture separates text encoding from image token generation. Therefore, several exemplary implementations also include a direct method for warming up a model with a pre-trained text encoder. Intuitively, a good text encoder for visual synthesis should be capable of general language comprehension in addition to visually robust prompts. Several exemplary implementations pre-train the text encoder on two datasets: the Colossal Clean Crawled Corpus (C4) with BERT pre-training objectives and an image-text dataset with contrasting learning objectives. After pre-training, several exemplary implementations continue training both the encoder and decoder for image generation from text with a softmax cross-entropy loss over a vocabulary of 8192 individual image tokens.

[0119] Example guidance and re-ranking without a classifier Classifier-less guidance is typically used in situations where the sample quality of a diffusion model is improved without a pre-trained classifier. In this setup, the generative model G is trained to perform unconditional generation G(z) (where z represents random noise) and conditional generation G(z,c) (where c represents several conditions, such as language description in this case). It is simply implemented as a conditional vector that is randomly omitted (hidden with a mask or switched to a learned embedding) with a certain probability. During the inference process period, the output I is sampled 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 guidance without a classifier. Intuitively, it can be seen that it decreases the unconditional likelihood of the sample while increasing the conditional likelihood, which improves the fit of the generated sample to the condition.

[0120] Classifier-less guidance has proven highly effective and is similarly applied in the context of autoregressive models for image generation from text. One exemplary method involves randomly swapping text prompts with padding tokens while fine-tuning the model. During the inference period, a linear combination of logits sampled from the unconditional model (conditioned with padding tokens) and the model conditioned with the original text prompts is taken.

[0121] Several exemplary implementations also apply guidance without a classifier to Babeldraw, and it has been found that this has a positive impact on output quality, particularly on difficult text prompts. Some exemplary implementations fine-tune the model over 100,000 steps, randomly swapping 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] Several exemplary implementations sample 16 images for each text prompt. For each output, several exemplary implementations re-rank it based on the image and text embedding match scores of a Contrastive Captioners (CoCa) model. A CoCa-based size model can be trained on the same dataset. It is worth noting that re-ranking across small sets of images is computationally inexpensive for image generation sampling from text.

[0123] Exemplary scaling techniques Several exemplary implementations use GSPMD and the Lingvo framework to scale models on CloudTPUv4 hardware for both training and inference. 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 a few tensors to instruct the compiler to automatically distribute the data and compute it across thousands of devices.

[0124] Training. In both 350M and 750M models, several exemplary implementations simply train the model with data parallelism. In 3B models, several exemplary implementations use 4-way in-layer model parallelism and 128-way data parallelism. Scaling 3B models is sufficient by partitioning a single dimension in each tensor. Model weights can be partitioned on a number of feedforward hidden dimensions and attention head dimensions, and the internal activation tensors of the feedforward and attention layers can also be partitioned on the hidden and head dimensions, but the difference from Megatron-LM is that some exemplary implementations also completely partition the output activations of the feedforward 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 activation memory.

[0125] As one example, Figure 6 illustrates an exemplary technique for scaling model training according to an exemplary embodiment of the present disclosure. In particular, Figure 6 illustrates four-way in-layer model parallelism with fully partitioned activity scaling exemplary 3B model training. Figure 6 shows a simplified transformer feedforward layer (with the sequence dimension omitted), where each shade represents data on one device. Some exemplary implementations additionally utilize 128-way data parallelism.

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

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

[0128] Several exemplary implementations configure the model to have separate encoder and decoder pipelines, each with 16 stages. Some exemplary implementations also use 64-way data parallelism in addition to pipelined training to speed up training. However, this reduces the batch size per core and exposes it to the further challenge of excessive pipeline bubbles. To reduce the ratio of pipeline bubbles, some exemplary implementations fit a circular schedule into the decoder pipeline, with four layers running in round-robin order at each stage. Apart from the encoder and decoder, some exemplary implementations use the same set of devices to perform data parallelism instead of pipelined training for the embedding, softmax, and image tokenizer layers.

[0129] During the training period, the Adafactor optimizer can be used to conserve memory with beta1=0.9, beta2=0.96, and a separation weight decay value of 4.5e-2. The first moment of the optimizer slot variable is additionally quantized from float32 to int8. Several 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 20B models to enable model pipelining. Several exemplary implementations use an exponential learning rate schedule with a default learning rate of 4.5e-5 and 5000 warm-up steps. Exponential decay starts at training step 85000 with a total of 450000 steps, resulting in a final decay rate of 0.025. Several exemplary implementations do not use an exponential moving average of model weights to conserve device memory. For text coding, several exemplary implementations build sentence-piece models 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. On both the encoder and decoder outputs, some exemplary implementations apply an additional regularization layer.

[0130] Inference. One primary goal for inference optimization is to speed up small batch image generation. Several exemplary implementations opt for in-layer model parallelism for both 3B and 20B models. In contrast to training, some exemplary implementations do not fully partition output activity for feedforward and attention layers for inference. This is because (1) each step of autoregressive decoding generates a much smaller tensor, and AllReduce (when writing) performs better on smaller data, and (2) active memory is independent of the inference period and there is no pass-through.

[0131] Exemplary devices and systems Figure 8A shows a block diagram of an exemplary computing system 100 according to 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, all of which are connected communicably via a network 180.

[0132] The user computing device 102 may be any type of computing device, such as a personal computing device (e.g., a laptop or desktop), a mobile computing device (e.g., a smartphone or tablet), a game 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 memory 114. The one or more processors 112 may be any suitable processing device (e.g., a processor core, microprocessor, ASIC, FPGA, controller, microcontroller, etc.) and may be one processor or multiple processors operably connected. The memory 114 may include one or more non-temporary 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 an operation.

[0134] In some implementations, the user computing device 102 may store or contain one or more machine learning models 120. For example, the machine learning models 120 may be, or may not contain, a variety of machine learning models, such as neural networks (e.g., deep neural networks) or other types of machine learning models including nonlinear and / or linear models. Neural networks may include feedforward neural networks, recurrent neural networks (e.g., long-short-term memory recurrent neural networks), convolutional neural networks, or other forms of neural networks. Some exemplary machine learning models may leverage attention mechanisms, such as self-attention. For example, some exemplary machine learning models may include multi-head self-attention models (e.g., transformer models). Exemplary machine learning models 120 are discussed with reference to Figures 1 and 2.

[0135] In some implementations, one or more machine learning models 120 can be received from a server computing system 130 via a network 180, stored in user computing device memory 114, and then 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 (for example, to perform parallel image quantization across multiple instances of an image).

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

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

[0138] The server computing system 130 includes one or more processors 132 and memory 134. The one or more processors 132 may be any suitable processing device (e.g., a processor core, microprocessor, ASIC, FPGA, controller, microcontroller, etc.) and may be one processor or multiple processors connected in an operable manner. The memory 134 may include one or more non-temporary 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 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 any combination thereof.

[0140] As described above, the server computing system 130 may store, or otherwise include, one or more machine learning models 140. For example, the models 140 may or may not include a variety of machine learning models. Exemplary machine learning models include neural networks or other multilayer nonlinear models. Exemplary neural networks include feedforward neural networks, deep neural networks, recurrent neural networks, and convolutional neural networks. Some exemplary machine learning models may leverage attention mechanisms such as self-attention. For example, some exemplary machine learning models may include multi-head self-attention models (e.g., transformer models). Exemplary models 140 are discussed with reference to Figures 1 and 2.

[0141] The user computing device 102 and / or the server computing system 130 can train models 120 and / or 140 through cross-referencing using a training computing system 150 which is communicably connected via a 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 memory 154. The one or more processors 152 may be any suitable processing device (e.g., a processor core, microprocessor, ASIC, FPGA, controller, microcontroller, etc.) and may be one processor or multiple processors operably connected. The memory 154 may include one or more non-temporary computer-readable storage media and combinations thereof, such as RAM, ROM, EEPROM, EPROM, flash memory devices, magnetic disks, etc. 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 is implemented by one or more server computing devices.

[0143] The training computing system 150 may include a model trainer 160 that trains machine learning models 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 error backpropagation. For example, a loss function can be backpropagated through the model to update one or more parameters of the model (for example, 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. The gradient descent technique can be used to iteratively update parameters over several training iterations.

[0144] In some implementations, the step of performing error backpropagation may include a step of performing backpropagation with time-varying truncation. The model trainer 160 can improve the generalization ability of the trained model by implementing several generalization techniques (e.g., weight decay, dropout, etc.).

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

[0146] In some implementations, if the user consents, training examples can be provided by the user's computing device 102. Therefore, 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's computing device 102. In some cases, this process can be called a model personalization step.

[0147] The model trainer 160 includes computer logic used to achieve 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, loaded into memory, and executed by one or more processors. In other implementations, the model trainer 160 includes one or more sets of computer executable instructions stored in a tangible computer-readable storage medium such as RAM, a hard disk, or an optical or magnetic medium.

[0148] Network 180 may 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 any combination thereof, and may include any number of wired or wireless links. Generally, communications on Network 180 can be carried over any type of wired and / or wireless connection using a variety of communication protocols (e.g., TCP / IP, HTTP, SMTP, FTP), encoding or formatting (e.g., HTML, XML), and / or protection methods (e.g., VPN, Secure HTTP, SSL).

[0149] The machine learning models described herein can be used for a variety of tasks, applications, and / or use cases.

[0150] In some implementations, the input to the machine learning model of this disclosure may be image data. The machine learning model can process the image data to generate an output. For example, the machine learning model can process the image data to generate an image recognition output (e.g., recognition of image data, latent embedding of image data, encoded representation of image data, hash of image data, etc.). Another example is that the machine learning model can process the image data to generate an image segmentation output. Another example is that the machine learning model can process the image data to generate an image classification output. Another example is that the machine learning model can process the image data to generate an image data transformation output (e.g., modification of image data, etc.). Another example is that the machine learning model can process the image data to generate an encoded image data output (e.g., encoded and / or compressed representation of image data, etc.). Another example is that the machine learning model can process the image data to generate a high-quality image data output. Another example is that 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 this disclosure may be latent coded data (e.g., a latent spatial representation of the input). The machine learning model can process the latent coded data to generate an output. For example, the machine learning model can process the latent coded data to generate a recognition output. As another example, the machine learning model can process the latent coded data to generate a reconstruction output. As yet another example, the machine learning model can process the latent coded data to generate a search output. As yet another example, the machine learning model can process the latent coded data to generate a reclustered output. As yet another example, the machine learning model can process the latent coded data to generate a prediction output.

[0152] In some cases, a machine learning model can be configured to perform a task that includes a step of encoding input data for reliable and / or efficient transmission or memory (and / or corresponding decoding). For example, the task may be an audio compression task. The input may contain audio data, and the output may contain compressed audio data. In another example, the input may contain visual data (e.g., one or more images or videos), and the output may contain compressed visual data, and the task is a visual data compression task. In yet another example, the task may include a step of generating embeddings 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, where the output is a set of scores, each corresponding to a different object class, representing the likelihood that one or more images depict an object belonging to that object class. The image processing task may be object detection, where the image processing output identifies one or more regions in one or more images, and for each region, the likelihood that the region depicts the object in question. As another example, the image processing task may be image segmentation, where the image processing output defines the likelihood for each pixel in one or more images for each category in a predetermined set of categories. For example, the set of categories may be foreground and background. As yet another example, the set of categories may be object classes. As yet another example, the image processing task may be depth estimation, where the image processing output defines the depth value for each pixel in one or more images. As another example, the image processing task may be motion estimation, where the network input includes multiple images, and the image processing output defines, for each pixel of one of the input images, the motion of the scene depicted by the pixels between the images in the network input.

[0154] Figure 8A illustrates one exemplary computing system that can be used to implement the present disclosure. Other computing systems can be used in a similar manner. For example, in some implementations, the user computing device 102 may include a model trainer 160 and a training dataset 162. In such implementations, the model 120 can be both trained and used locally on the user computing device 102. In some such implementations, the user computing device 102 may implement the model trainer 160 to personalize the model 120 based on user-specific data.

[0155] Figure 8B shows a block diagram of an exemplary computing device 10 implemented according to 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 (for example, 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, and a browser application.

[0157] As illustrated in Figure 8B, each application can communicate with several other components of the computing device, such as one or more sensors, a context manager, device state components, and / or further 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 specific to that application.

[0158] Figure 8C shows a block diagram of an exemplary computing device 50 implemented according to an exemplary embodiment of the present disclosure. The computing device 50 may be a user computing device or a server computing device.

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

[0160] The central intelligence layer includes several machine learning models. For example, as illustrated in Figure 8C, each machine learning model can be dedicated to a specific 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 may have a single model for all applications. In some implementations, the central intelligence layer is contained within the operating system of the computing device 50, or is implemented by the operating system of the computing device 50.

[0161] The central intelligence layer can communicate with the central device data layer. The central device data layer may be a centralized repository of data for the computing device 50. As illustrated in Figure 8C, the central device data layer can communicate with several other components of the computing device, such as one or more sensors, a context manager, a device state component, and / or further 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 technologies discussed herein involve referencing servers, databases, software applications, and other computer-based systems, as well as performing actions and transmitting information to and from such systems. The inherent flexibility of computer-based systems allows for a wide variety of possible configurations, combinations, and divisions of tasks and functionalities between their components. For example, the processes discussed herein 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. Distributed components can operate sequentially or in parallel.

[0163] While the subject matter is described in detail with respect to various specific examples of implementation, each example is provided as an example and not as an limitation of the disclosure. Those skilled in the art, having attained the above understanding, will readily be able to create modifications, variations, and equivalents to such embodiments. Therefore, the disclosure of the subject matter does not preclude the inclusion of such modifications, variations, and / or additions to the subject matter that will readily become apparent to those skilled in the art. For example, features described or mentioned as part of one embodiment can be used in conjunction with another embodiment to create further embodiments. Thus, the disclosure is intended to cover such modifications, variations, and equivalents. [Explanation of Symbols]

[0164] 10 Computing Devices 12 Input Image Patches 16 Transformer encoders, image encoders 18 Image Tokens 20 Quantization Code 22 Codebooks, Quantization Codebooks, Candidate Codes 24 Transformer Decoders, Image Decoders 26 Composite Image Patch 28 Composite Images 30 Loss Function 50 Computing Devices 102 User Computing Devices 112 processors 114 memory 116 data 118 Command 120 Machine Learning Models 122 User Input Components 130 Server Computing Systems 132 processors 134 memory 136 data 138 Command 140 Machine Learning Models 150 Training Computing Systems 152 processors 154 memory 156 data 158 Command 160 Model Trainer 162 training data, training dataset 180 Networks 212 Input Image Patch 214 Input Images 216 Transformer encoders, image encoders 218 image tokens 222 Codebook 224 Code Prediction Models 226 Predicted Codes 228 Transformer Encoders, Image Decoders 230 Composite Image Patches 232 Composite Images 332 Composite Images 334 Natural language input, input text 402 Text iOS 404 Autoregressive Code Selector 406 Text Embedding

Claims

1. A computer-implemented method for performing vector quantization of an image, A computing system comprising one or more computing devices obtains multiple input image patches of an image, The computing system performs the steps of processing the plurality of input image patches with a machine learning-trained image encoder and generating 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 learning-trained 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 computing system performs the steps of mapping the plurality of image tokens to a plurality of quantization codes included in a quantization codebook containing a plurality of candidate codes, The computing system provides the plurality of quantization codes as encoded versions of the image. A computer implementation method that includes this.

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

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

4. The computing system performs the steps of processing the multiple quantization codes with a machine learning-trained image decoder to generate multiple composite image patches that form a composite image, The computing system provides a loss function that provides loss values ​​based at least partially on the composite image, The computing system performs the steps of modifying the machine learning image encoder, the machine learning image decoder, and one or more of the multiple candidate codes, at least partially based on the loss function. A computer-implemented method according to any one of claims 1 to 3, further comprising:

5. The computer-implemented method according to claim 4, wherein the machine learning-trained image decoder comprises a vision transformer model.

6. The aforementioned loss function is, Logit-Laplace loss term, L2 loss term, Perceptual loss term, and Adversarial generative network loss term A computer-implemented method according to claim 4 or 5, including the method described in claim 4 or 5.

7. The computing system performs the step of mapping the plurality of image tokens to the plurality of quantization codes included in the quantization codebook which includes the plurality of candidate codes, The computing system performs the steps of projecting the plurality of image tokens into a lower-dimensional space, The steps include: Projecting the image tokens into the lower-dimensional space; and, using the computing system, mapping the plurality of image tokens to the plurality of quantization codes contained in the quantization codebook. A computer-implemented method according to any one of claims 1 to 6, including the method described in any one of claims 1 to 6.

8. The computing system performs the step of mapping the plurality of image tokens to the plurality of quantization codes included in the quantization codebook which includes the plurality of candidate codes, The computing system includes the step of applying L2 regularization to one or both of the plurality of image tokens and the plurality of candidate codes, Following the step of applying the L2 regularization, the computing system performs the step of mapping the plurality of image tokens to the plurality of quantization codes contained in the quantization codebook. A computer-implemented method according to any one of claims 1 to 7, including the method described in any one of claims 1 to 7.

9. The computing system, using a machine learning-based code prediction model, autoregressively predicts a plurality of predicted codes from the quantized codebook based at least partially on one or more of the plurality of quantized codes; The computing system performs the steps of processing the multiple predicted codes with a machine learning image decoder to generate multiple composite image patches that form a composite image. A computer-implemented method according to any one of claims 1 to 8, further comprising:

10. The computing system performs the steps of evaluating a code prediction loss function that evaluates the negative log-likelihood based on the plurality of predicted codes, The computing system performs the steps of modifying one or more parameters of the machine learning code prediction model based on the code prediction loss function. The computer-implemented method according to claim 9, further comprising:

11. The computer-implemented method according to claim 9 or 10, wherein the step of autoregressively predicting the plurality of predicted codes by the computing system using the machine-learned code prediction model includes the step of conditioning the machine-learned code prediction model with auxiliary conditioning data that describes one or more desired properties of the composite image by the computing system.

12. The computer-implemented method according to claim 11, wherein the auxiliary conditioning data includes class labels describing a desired class of the composite image.

13. The computer-implemented method according to claim 11 or 12, wherein the auxiliary conditional data includes natural language text tokens.

14. The computing system performs the step of conditioning the machine learning code prediction model with the natural language text tokens, The computing system performs the steps of: processing the natural language text tokens with the text encoder unit of the machine learning code prediction model to generate text embeddings; The computing system provides the text embedding as input to the autoregressive code selection unit of the machine learning code prediction model, thereby autoregressively predicting the plurality of predicted codes. A computer-implemented method according to claim 13, including the method described in claim 13.

15. The computing system performs the steps of: extracting one or more intermediate features from the machine learning code prediction model; The computing system performs the steps of predicting a class label for the image based at least partially on the intermediate features, and A computer-implemented method according to any one of claims 9 to 14, further comprising:

16. A computer-implemented method for performing vector quantization of an image, A computing system comprising one or more computing devices obtains a plurality of quantization codes that form an encoded version of an image, wherein the plurality of quantization codes are selected by mapping a plurality of image tokens generated by a machine learning image encoder model to the plurality of quantization codes included in a quantization codebook containing a plurality of candidate codes. The computing system performs the steps of processing the multiple quantization codes with a machine learning-trained image decoder to generate multiple decoded image patches that form a decoded version of the image. Includes, A computer implementation method in which 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 actions.

17. A computer system comprising one or more processors and one or more non-temporary computer-readable media that, when executed by the one or more processors, collectively store instructions causing the one or more processors to perform the method according to any one of claims 1 to 16.

18. One or more processors, A non-temporary computer-readable medium that collectively stores machine learning-trained image processing models. The machine learning image processing model is equipped with, An encoder unit 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 partially based on the one or more quantization codes, A feature prediction unit is configured to generate one or more feature predictions for the input image patch based at least partially on the data extracted from the code prediction unit. A computing system equipped with the following features.

19. The computing system according to claim 18, further comprising a decoder unit configured to generate an image patch reconstructed based on one or more quantization codes, or to generate a composite image patch at least partially based on one or more predicted quantization codes.

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

21. A computer-implemented method for generating images from text, A computing system comprising one or more computing devices obtains natural language input describing desired image content; The computing system performs the steps of: processing the natural language input with the text encoder section of a machine learning-based code prediction model to generate text embeddings; The computing system performs the following steps: it processes the text embedding in the autoregressive code selection unit of the machine learning code prediction model to autoregressively predict a sequence of codes to be predicted from a quantized codebook containing multiple candidate codes; The computing system performs the steps of processing the sequence of quantization codes with a machine learning-trained image decoder to generate a plurality of composite image patches that form a composite image. Includes, A computer-implemented method wherein the composite image depicts the desired image content.

22. The computer-implemented method according to claim 21, wherein one or more of the text encoder unit of the machine learning code prediction model, the autoregressive code selection unit of the machine learning code prediction model, and the machine learning image decoder are configured to perform one or more self-attention operations.

23. The computer-implemented method according to claim 21 or 22, wherein one or more of the text encoder unit of the machine-learned code prediction model, the autoregressive code selection unit of the machine-learned code prediction model, and the machine-learned image decoder comprises a transform neural network.

24. The computer-implemented method according to any one of claims 21 to 23, wherein one or both of the machine-learned image decoder and the codebook are trained together in an image encoder model.

25. The computer implementation method according to any one of claims 21 to 23, wherein the text encoder portion of the machine learning code prediction model is pre-trained in a pre-training task.

26. A computer-implemented method for training a code prediction model, A computing system comprising one or more computing devices obtains training examples including training images and natural language input describing the content of the training images, The computing system processes multiple image patches from the training image using a machine learning-trained image encoder to generate multiple image tokens in a latent space, wherein the multiple image tokens correspond to the multiple image patches. The computing system performs the steps of mapping the plurality of image tokens to a plurality of quantization codes included in a quantization codebook containing a plurality of candidate codes, The computing system performs the steps of processing the natural language input in the text encoder section of the code prediction model to generate text embeddings, The computing system performs the steps of processing the text embedding in the autoregressive code selection unit of the code prediction model to autoregressively predict the sequence of codes to be predicted from the quantized codebook, The computing system includes the step of evaluating a code prediction loss function that compares the sequence of predicted codes with the plurality of quantized codes, The computing system performs the steps of modifying one or more parameters of the code prediction model, at least partially based on the code prediction loss function. A computer implementation method that includes this.

27. The computer-implemented method according to claim 26, wherein the machine-learned image encoder processes the plurality of input image patches by performing one or more self-attention operations and generates the plurality of image tokens in the latent space.

28. The computer-implemented method according to claim 26 or 27, wherein the code prediction loss function evaluates the negative log-likelihood of the predicted code for the quantized code.

29. The computer-implemented method according to any one of claims 26 to 28, wherein the machine-learned image encoder and the codebook are pre-trained in an image decoder in an autoencoder architecture.

30. The computer-implemented method according to claim 26, wherein the computing system modifies one or more parameters of the code prediction model at least partially based on the code prediction loss function, the computing system modifies one or more parameters of both the text encoder portion and the autoregressive code selection portion of the code prediction model at least partially based on the code prediction loss function.

Images