Momentum contrastive autoencoder
The embodiments are directed to training a momentum contrastive autoencoder using a contrastive learning framework. The contrastive learning framework learns a latent space distribution by matching latent representations of the momentum contrastive autoencoder to a pre-specified distribution, such as a distribution over a unit hyper-sphere. Once the latent space distribution is learned, samples for a new data set may be obtained from the latent space distribution. This results in a simple and scalable algorithm that avoids many of the optimization challenges of existing generative models, while retaining the advantage of efficient sampling.
1 . A system for training a contrastive momentum autoencoder, the system comprising:
one or more non-transitory memories;
one or more processors coupled the one or more memories and configured to execute instructions that cause the one or more processors to perform operations, the operations comprising:
receive an input data set;
encode, using a first encoder neural network comprising first encoder parameters, a first latent representation from the input data set;
encode, using a second encoder neural network having a same structure as the first encoder neural network and comprising second encoder parameters having values that are different from values of the first encoder parameters, a second latent representation from a same input data set as for the first encoder neural network;
determine a contrastive loss using the first latent representation and second latent representation, the contrastive loss based at least on a similarity between the first latent representation and the second latent representation;
decode, using a decoder neural network comprising decoder parameters, the first latent representation into an output data set;
determine a reconstruction loss based on the output data set and the input data set; and
update at least one parameter in the first encoder parameters of the first encoder neural network and at least one parameter in the decoder parameters of the decoder neural network based on the contrastive loss and the reconstruction loss.
2 . The system of claim 1 , the operations further comprising:
after the updating of the at least one parameter in the first encoder parameters, update at least one parameter in the second encoder parameters of the second encoder neural network based on a moving average of the at least one parameter in the first encoder parameters.
3 . The system of claim 1 , wherein the first encoder is configured to normalize the first latent representation using l 2 normalization.
4 . The system of claim 3 , the operations further comprising:
learning a distribution of a latent space by mapping a distribution of the normalized first latent representation to a prior distribution, wherein the prior distribution is a uniform distribution over a unit hyper-sphere, and wherein the mapping minimizes the contrastive loss based on the first latent representation, the second latent representation, and the prior distribution.
5 . The system of claim 4 , wherein the contrastive loss includes a positive component associated with a loss based at least on a similarity between the first latent representation and the second latent representation and a negative component associated with a loss that corresponds to mapping the first latent representation to the prior distribution that is the uniform distribution over the unit hyper-sphere.
6 . The system of claim 5 , wherein the prior distribution is based on prior latent representations encoded by the second encoder neural network.
7 . The system of claim 6 , the operations further comprising:
updating the distribution of a latent space using the second latent representation encoded by the second encoder neural network.
8 . The system of claim 4 , the operations further comprising:
retrieving samples from the learned distribution over the latent space; and
decoding, using the decoder neural network, the samples into a new output data set.
9 . The system of claim 8 , wherein the new output data set is a new data set that is different from the input data set and the output data set, wherein the input data set comprises a first image, the output data set comprises a second image, and the new output dataset comprises a third image.
10 . The system of claim 1 , wherein the first encoder neural network and the second encoder neural network are configured to receive the input data set as multiple mini-batches; and
the operations further comprising:
updating the at least one parameter in the first encoder parameters of the first encoder neural network and at least one parameter in the decoder parameters of the decoder neural network based on the contrastive loss and the reconstruction loss associated with a data sample in each mini-batch in the multiple mini-batches.
11 . The system of claim 1 , further comprising, generating a first output image via the trained decoder neural network and the trained first encoder neural network without the trained second encoder neural network, in response to a first input image.
12 . The system of claim 1 , wherein the second encoder neural network trains the first encoder neural network such that the first encoder parameters resemble the second encoder parameters after the updating of the at least one parameter in the first encoder parameters and at least one parameter in the decoder parameters.
13 . A method for training a contrastive momentum autoencoder, the method comprising:
receiving a training dataset including images;
encoding, using a first encoder neural network comprising first encoder parameters, a first latent representation from an image in the training dataset input data set;
encoding, using a second encoder neural network having a same structure as the first encoder neural network and comprising second encoder parameters having values that are different from values of the first encoder parameters, a second latent representation from a same image as for the first encoder neural network;
determining a contrastive loss using the first latent representation and second latent representation, the contrastive loss based at least on a similarity between the first latent representation and the second latent representation;
decoding, using a decoder neural network comprising decoder parameters, the first latent representation into an output image;
determining a reconstruction loss using the output image and the image in the training dataset;
updating at least one parameter in the first encoder parameters and at least one parameter in the decoder parameters based on the contrastive loss and the reconstruction loss; and
after the first encoder neural network and the decoder neural network is trained:
receiving an input image; and
generating, by the trained first encoder neural network and the trained decoder neural network, a reconstructed image in response to the input image.
14 . The method of claim 13 , further comprising:
after the updating of the at least one parameter in the first encoder parameters, updating at least one parameter in the second encoder parameters of the second encoder neural network based on a moving average of the at least one parameter in the first encoder parameters.
15 . The method of claim 13 , wherein normalizing using l 2 normalization and the first encoder neural network, the first latent representation.
16 . The method of claim 15 , further comprising:
learning a distribution of a latent space by mapping a distribution of the normalized first latent representation to a prior distribution, wherein the prior distribution is a uniform distribution over a unit hyper-sphere, and wherein the mapping minimizes the contrastive loss based on the first latent representation, the second latent representation, and the prior distribution.
17 . The method of claim 16 , wherein the contrastive loss includes a positive component associated with a loss between the first latent representation and the second latent representation and a negative component associated with a loss that corresponds to mapping the first latent representation to the prior distribution that is the uniform distribution over the unit hyper-sphere.
18 . The method of claim 17 , further comprising:
updating the distribution of a latent space using the second latent representation encoded by the second encoder neural network.
19 . The method of claim 13 , further comprising:
receiving the input data set as multiple mini-batches; and
updating the at least one parameter in the first encoder parameters of the first encoder neural network and at least one parameter in the decoder parameters of the decoder neural network based on the contrastive loss and the reconstruction loss associated with each mini-batch in the multiple mini-batches.
20 . A non-transitory computer readable medium having instructions thereon, that when executed by a processor cause the processor to perform operations that train a contrastive momentum autoencoder, the operations comprising:
receiving an input data set;
encoding, using a first encoder neural network comprising first encoder parameters, a first latent representation from the input data set;
encoding, using a second encoder neural network having a same structure as the first encoder neural network and comprising second encoder parameters having values that are different from values of the first encoder parameters, a second latent representation from a same input data set as for the first encoder neural network;
determining a contrastive loss using the first latent representation and second latent representation, the contrastive loss based at least on a similarity between the first latent representation and the second latent representation;
decoding, using a decoder neural network comprising decoder parameters, the first latent representation into an output data set;
determining a reconstruction loss based on the output data set and the input data set; and
updating at least one parameter in the first encoder parameters and at least one parameter in the decoder parameters based on the contrastive loss and the reconstruction loss.