Generation of latent representations of images using a machine learning model
The present disclosure describes techniques for generating latent representations of images using a machine learning model. An image is split and flattened into a series of patches. The series of patches is concatenated with a sequence of latent tokens. The concatenated patches and latent tokens are input into an encoder of the machine learning model. A one-dimensional (1D) latent representation of the image is generated by the encoder. Vector quantization is performed on the 1D latent representation of the image by a vector quantizer of the machine learning model to generate quantized latent tokens. The image is reconstructed based on the quantized latent tokens by a decoder of the machine learning model.
1 . A method of generating latent representations of images using a machine learning model, comprising:
splitting and flattening an image into a series of patches;
concatenating the series of patches with a first sequence of latent tokens;
inputting the concatenated series of patches and the first sequence of latent tokens into an encoder of the machine learning model;
generating, by the encoder, a one-dimensional (1D) latent representation of the image, wherein the 1D latent representation of the image comprises a second sequence of latent tokens generated by the encoder during feature encoding of the concatenated series of patches and the first sequence of latent tokens, and wherein a size of the 1D latent representation is decoupled from a resolution of the image;
performing vector quantization on the 1D latent representation of the image by a vector quantizer of the machine learning model to generate quantized latent tokens; and
reconstructing the image based on the quantized latent tokens by a decoder of the machine learning model.
2 . The method of claim 1 , wherein the second sequence of latent tokens comprises a 1D sequence, wherein the 1D sequence has a same size as the first sequence of latent tokens.
3 . The method of claim 1 , further comprising:
incorporating a sequence of mask tokens into the quantized latent tokens; and
inputting the incorporated sequence of mask tokens and quantized latent tokens into the decoder.
4 . The method of claim 3 , further comprising:
generating the sequence of mask tokens by replicating a single mask token.
5 . The method of claim 1 , wherein the encoder comprises a Vision Transformer (ViT) encoder, and wherein the decoder comprises a ViT decoder.
6 . The method of claim 1 , further comprising:
training the machine learning model using a two-stage training paradigm.
7 . The method of claim 6 , further comprising:
performing a first stage of training on the machine learning model using proxy codes to bypass intricate loss functions and generative adversarial network (GAN) architectures, wherein the proxy codes comprise discrete codes generated by an off-the-shelf Vector Quantized GAN (VQGAN) model.
8 . The method of claim 7 , further comprising:
performing a second stage of training on the machine learning model to fine-tune the decoder, and wherein the second stage of training comprises only training the decoder towards pixel space with VQGAN losses while keeping the encoder and the vector quantizer frozen.
9 . A system of generating latent representations of images using a machine learning model, comprising:
at least one processor; and
at least one memory communicatively coupled to the at least one processor and comprising computer-readable instructions that upon execution by the at least one processor cause the at least one processor to perform operations comprising:
splitting and flattening an image into a series of patches;
concatenating the series of patches with a first sequence of latent tokens;
inputting the concatenated series of patches and the first sequence of latent tokens into an encoder of the machine learning model;
generating, by the encoder, a one-dimensional (1D) latent representation of the image, wherein the 1D latent representation of the image comprises a second sequence of latent tokens generated by the encoder during feature encoding of the concatenated series of patches and the first sequence of latent tokens, and wherein a size of the 1D latent representation is decoupled from a resolution of the image;
performing vector quantization on the 1D latent representation of the image by a vector quantizer of the machine learning model to generate quantized latent tokens; and
reconstructing the image based on the quantized latent tokens by a decoder of the machine learning model.
10 . The system of claim 9 , wherein the second sequence of latent tokens comprises a 1D sequence, wherein the 1D sequence has a same size as the first sequence of latent tokens.
11 . The system of claim 9 , the operations further comprising:
generating a sequence of mask tokens by replicating a single mask token;
incorporating the sequence of mask tokens into the quantized latent tokens; and
inputting the incorporated sequence of mask tokens and quantized latent tokens into the decoder.
12 . The system of claim 9 , wherein the encoder comprises a Vision Transformer (ViT) encoder, and wherein the decoder comprises a ViT decoder.
13 . The system of claim 9 , the operations further comprising:
training the machine learning model using a two-stage training paradigm.
14 . The system of claim 13 , the operations further comprising:
performing a first stage of training on the machine learning model using proxy codes to bypass intricate loss functions and generative adversarial network (GAN) architectures, wherein the proxy codes comprise discrete codes generated by an off-the-shelf Vector Quantized GAN (VQGAN) model; and
performing a second stage of training on the machine learning model to fine-tune the decoder, and wherein the second stage of training comprises only training the decoder towards pixel space with VQGAN losses while keeping the encoder and the vector quantizer frozen.
15 . A non-transitory computer-readable storage medium, storing computer-readable instructions that upon execution by a processor cause the processor to implement operations comprising:
splitting and flattening an image into a series of patches;
concatenating the series of patches with a first sequence of latent tokens;
inputting the concatenated series of patches and the first sequence of latent tokens into an encoder of the machine learning model;
generating, by the encoder, a one-dimensional (1D) latent representation of the image, wherein the 1D latent representation of the image comprises a second sequence of latent tokens generated by the encoder during feature encoding of the concatenated series of patches and the first sequence of latent tokens, and wherein a size of the 1D latent representation is decoupled from a resolution of the image;
performing vector quantization on the 1D latent representation of the image by a vector quantizer of the machine learning model to generate quantized latent tokens; and
reconstructing the image based on the quantized latent tokens by a decoder of the machine learning model.
16 . The non-transitory computer-readable storage medium of claim 15 , wherein the second sequence of latent tokens comprises a 1D sequence, wherein the 1D sequence has a same size as the first sequence of latent tokens.
17 . The non-transitory computer-readable storage medium of claim 15 , the operations further comprising:
generating a sequence of mask tokens by replicating a single mask token;
incorporating the sequence of mask tokens into the quantized latent tokens; and
inputting the incorporated sequence of mask tokens and quantized latent tokens into the decoder.
18 . The non-transitory computer-readable storage medium of claim 15 , wherein the encoder comprises a Vision Transformer (ViT) encoder, and wherein the decoder comprises a ViT decoder.
19 . The non-transitory computer-readable storage medium of claim 15 , the operations further comprising:
training the machine learning model using a two-stage training paradigm.
20 . The non-transitory computer-readable storage medium of claim 19 , the operations further comprising:
performing a first stage of training on the machine learning model using proxy codes to bypass intricate loss functions and generative adversarial network (GAN) architectures, wherein the proxy codes comprise discrete codes generated by an off-the-shelf Vector Quantized GAN (VQGAN) model; and
performing a second stage of training on the machine learning model to fine-tune the decoder, and wherein the second stage of training comprises only training the decoder towards pixel space with VQGAN losses while keeping the encoder and the vector quantizer frozen.