Systems and methods for semi-supervised learning with contrastive graph regularization
Embodiments described herein provide an approach (referred to as “Co-training” mechanism throughout this disclosure) that jointly learns two representations of the training data, their class probabilities and low-dimensional embeddings. Specifically, two representations of each image sample are generated: a class probability produced by the classification head and a low-dimensional embedding produced by the projection head. The classification head is trained using memory-smoothed pseudo-labels, where pseudo-labels are smoothed by aggregating information from nearby samples in the embedding space. The projection head is trained using contrastive learning on a pseudo-label graph, where samples with similar pseudo-labels are encouraged to have similar embeddings.
1. A method for semi-supervised learning with contrastive graph regularization, the method comprising:
receiving a batch of unlabeled image samples;
generating, from an unlabeled image sample, a weakly augmented image sample, a first strongly augmented image sample, and a second strongly augmented image sample;
generating, by an encoder and a classification of a neural model, a pseudo-label corresponding to the weakly augmented image sample, wherein the pseudo-label is generated by:
generating, by the encoder and the classification of the neural model, a classification probability for the weakly augmented image sample,
generating, by the encoder and a projection head of the neural model, an embedding for the weakly augmented image sample,
storing, at a memory bank, the generated classification probability and the generated embedding, and
aggregating class probabilities from neighboring samples in the memory bank to compute the pseudo-label;
generating, by the neural model, a first embedding corresponding to the first strongly augmented image sample and a second embedding corresponding to the second strongly augmented image sample;
building an embedding graph by comparing pairwise similarity between the first embedding and the second embedding;
generating a pseudo-label graph, comprising: constructing a similarity matrix among generated pseudo-labels corresponding to the batch of unlabeled image samples, wherein the similarity matrix contains a first similarity comparing each generated pseudo-label to itself, and a second similarity comparing each generated pseudo-label to another pseudo-label;
computing a contrastive loss based on a cross-entropy between the embedding graph and the pseudo-label graph;
computing an unsupervised classification loss based on a cross-entropy between the pseudo label and the classification probability for the weakly augmented image sample;
updating the neural model based at least in part on the contrastive loss and the unsupervised classification loss via backpropagation;
receiving, via a user data interface, an input image for classification; and
generating, via the neural model, an output classification for the input image.
2. The method of claim 1 , wherein the first embedding and the second embedding are generated by an encoder and a projection head of the neural model.
3. The method of claim 1 , further comprising:
storing class probabilities or embeddings generated from unlabeled image samples and labelled image samples at the memory bank according to a first-in-first-out strategy.
4. The method of claim 1 , wherein the generated classification probability is smoothed by a moving-average over a plurality of generated classification probabilities during training.
5. The method of claim 1 , further comprising: unsupervised classification loss receiving a batch of labeled image samples with ground-truth labels;
generating, by an encoder and a classification head of the neural model, a third classification prediction of a labeled image sample; and
computing a supervised classification loss between a ground-truth label corresponding to the labeled image sample and the third classification prediction.
6. The method of claim 5 , further comprising:
computing a weighted sum of the supervised classification loss, the unsupervised classification loss and the contrastive loss; and
jointly training the encoder, the classification head, and the projection head of the neural model based on the weighted sum.
7. The method of claim 1 , further comprising:
updating parameters of an encoder, classification head and a projection head of the neural model with a moving average of the parameters.
8. A system for semi-supervised learning with contrastive graph regularization, the system comprising:
a data interface that receives a batch of unlabeled image samples;
a memory that stores a neural model;
a processor that reads instructions from the memory to perform:
generating, from an unlabeled image sample, a weakly augmented image sample, a first strongly augmented image sample, and a second strongly augmented image sample;
generating, by an encoder and a classification of the neural model, a pseudo-label corresponding to the weakly augmented image sample, wherein the pseudo-label is generated by:
generating, by the encoder and the classification of the neural model, a classification probability for the weakly augmented image sample,
generating, by the encoder and a projection head of the neural model, an embedding for the weakly augmented image sample,
storing, at a memory bank, the generated classification probability and the generated embedding, and
aggregating class probabilities from neighboring samples in the memory bank to compute the pseudo-label;
generating, by the neural model, a first embedding corresponding to the first strongly augmented image sample and a second embedding corresponding to the second strongly augmented image sample;
building an embedding graph by comparing pairwise similarity between the first embedding and the second embedding;
generating a pseudo-label graph, comprising: constructing a similarity matrix among generated pseudo-labels corresponding to the batch of unlabeled image samples, wherein the similarity matrix contains a first similarity comparing each generated pseudo-label to itself, and a second similarity comparing each generated pseudo-label to another pseudo-label;
computing a contrastive loss based on a cross-entropy between the embedding graph and the pseudo-label graph;
computing an unsupervised classification loss based on a cross-entropy between the pseudo label and the classification probability for the weakly augmented image sample;
updating the neural model based at least in part on the contrastive loss and the unsupervised classification loss via backpropagation;
receiving, via a user data interface, an input image for classification; and
generating, via the neural model, an output classification for the input image.
9. The system of claim 8 , wherein the first embedding and the second embedding are generated by an encoder and a projection head of the neural model.
10. The system of claim 8 , wherein the processor further reads instructions from the memory to perform:
storing class probabilities or embeddings generated from unlabeled image samples and labeled image samples at the memory bank according to a first-in-first-out strategy.
11. The system of claim 8 , wherein the generated classification probability is smoothed by a moving-average over a plurality of generated classification probabilities during training.
12. The system of claim 8 , wherein the processor further reads instructions from the memory to perform:
receiving a batch of labeled image samples with ground-truth labels;
generating, by an encoder and a classification head of the neural model, a third classification prediction of a labeled image sample; and
computing a supervised classification loss between a ground-truth label corresponding to the labeled image sample and the third classification prediction.
13. The system of claim 12 , wherein the processor further reads instructions from the memory to perform:
computing a weighted sum of the supervised classification loss, the unsupervised classification loss and the contrastive loss; and
jointly training the encoder, the classification head, and the projection head of the neural model based on the weighted sum.
14. A non-transitory processor-readable medium storing processor-executable instructions for semi-supervised learning with contrastive graph regularization, the instructions being executed by a processor to perform:
receiving a batch of unlabeled image samples;
generating, from an unlabeled image sample, a weakly augmented image sample, a first strongly augmented image sample, and a second strongly augmented image sample;
generating, by an encoder and a classification of a neural model, a pseudo-label corresponding to the weakly augmented image sample, wherein the pseudo-label is generated by:
generating, by the encoder and the classification of the neural model, a classification probability for the weakly augmented image sample,
generating, by the encoder and a projection head of the neural model, an embedding for the weakly augmented image sample,
storing, at a memory bank, the generated classification probability and the generated embedding, and
aggregating class probabilities from neighboring samples in the memory bank to compute the pseudo-label;
generating, by the neural model, a first embedding corresponding to the first strongly augmented image sample and a second embedding corresponding to the second strongly augmented image sample;
building an embedding graph by comparing pairwise similarity between the first embedding and the second embedding;
generating a pseudo-label graph, comprising: constructing a similarity matrix among generated pseudo-labels corresponding to the batch of unlabeled image samples, wherein the similarity matrix contains a first similarity comparing each generated pseudo-label to itself, and a second similarity comparing each generated pseudo-label to another pseudo-label;
computing a contrastive loss based on a cross-entropy between the embedding graph and the pseudo-label graph;
computing an unsupervised classification loss based on a cross-entropy between the pseudo label and the classification probability for the weakly augmented image sample;
updating the neural model based at least in part on the contrastive loss and the unsupervised classification loss via backpropagation;
receiving, via a user data interface, an input image for classification; and
generating, via the neural model, an output classification for the input image.