Systems and methods for vision-language distribution alignment
Embodiments described herein a CROss-Modal Distribution Alignment (CROMDA) model for vision-language pretraining, which can be used for retrieval downstream tasks. In the CROMDA mode, global cross-modal representations are aligned on each unimodality. Specifically, a uni-modal global similarity between an image/text and the image/text feature queue are computed. A softmax-normalized distribution is then generated based on the computed similarity. The distribution thus takes advantage of property of the global structure of the queue. CROMDA then aligns the two distributions and learns a modal invariant global representation. In this way, CROMDA is able to obtain invariant property in each modality, where images with similar text representations should be similar and vice versa.
1. A system for vision-language distribution alignment, the system comprising:
a data interface receiving a first batch of image samples and a second batch of text samples;
a memory storing a plurality of processor-executable instructions, an image encoder for encoding the first batch of image samples into a first plurality of image feature representations stored at a first feature queue, and
a text encoder for encoding the second batch of text samples into a second plurality of text feature representations stored at a second feature queue; and
a processor executing the plurality of processor-executable instructions to perform operations comprising:
computing an image-to-image similarity between at least one image feature representation and the first plurality of image feature representations in the first feature queue;
computing a text-to-text similarity between at least one text feature representation and the second plurality of text feature representations in the second feature queue;
computing a cross-modal alignment loss based on the image-to-image similarity and the text-to-text similarity; and
updating the image encoder and the text encoder based at least in part on the cross- modal alignment loss.
2. The system of claim 1 , wherein the image-to-image similarity is computed by:
computing a plurality of momentum image feature representations as the plurality of image feature representations are continuously added to the first feature queue, wherein at least one momentum image feature representation is computed as an exponential moving average of image feature representations that are currently in the first feature queue; and
computing the image-to-image similarity based on vector similarities between the at least one image feature representation and each of the plurality of momentum image feature representations.
3. The system of claim 1 , wherein the text-to-text similarity is computed by:
computing a plurality of momentum text feature representations as the plurality of text feature representations are continuously added to the second feature queue, wherein at least one momentum text feature representation is computed as an exponential moving average of text feature representations that are currently in the second feature queue; and
computing the text-to-text similarity based on vector similarities between the at least one text feature representation and each of the plurality of momentum text feature representations.
4. The system of claim 1 , wherein the cross-modal alignment loss is computed based on a first Kullback-Leibler divergence between the image-to-image similarity and the text-to-text similarity, and a second Kullback-Leibler divergence between the text-to-text similarity and the image-to-image similarity.
5. The system of claim 1 , wherein the memory further stores a cross-modal encoder that applies cross-attention over the first plurality of image feature representations and second plurality of text feature representations.
6. The system of claim 5 , wherein the operations further comprise:
selecting a subset of samples from the first batch of image samples and the second batch of text samples; and
compute an image-text matching loss based on a similarity between an image sample and a text sample in an image-text pair, and a one-hot ground truth corresponding to the image-text pair.
7. The system of claim 6 , wherein the image-text pair is positive or negative, and
wherein the operations further comprise:
building an image-text queue storing a plurality of image-text pairs from the first batch of image samples and the second batch of text samples,
wherein the image-text queue is dynamically updated when new input samples are received;
selecting a negative sample corresponding to a sample in the subset from the image-text queue to form a negative image-text pair.
8. The system of claim 1 , wherein the operations further comprise:
randomly replacing text tokens in the second plurality of text samples with a mask token at a probability;
predicting replaced text tokens based on the cross-attention; and
computing a masked language modeling loss based on the predicted text tokens and original text tokens.
9. The system of claim 1 , wherein the operations further comprise:
computing an image-to-text similarity between the at least one image feature representation to all text feature representations in the second feature queue;
computing a text-to-image similarity between the at least one text feature representation to all image feature representations in the first feature queue; and
computing a cross-modal contrastive loss based on a first cross-entropy between the image-to-text similarity and an image-to-text ground truth, and a second cross-entropy between the text-to-image similarity and a text-to-image ground truth.
10. The system of claim 1 , wherein the operations further comprise:
computing a weighted sum of the cross-modal distribution alignment loss, a cross-modal contrastive loss, a masked language modeling loss and an image-text matching loss; and
updating at least the image encoder and the text encoder based on the weighted sum.
11. A method for vision-language distribution alignment, the method comprising:
receiving, via a data interface, a first batch of image samples and a second batch of text samples;
encoding, by an image encoder, the first batch of image samples into a first plurality of image feature representations stored at a first feature queue;
encoding, by a text encoder, the second batch of text samples into a second plurality of text feature representations stored at a second feature queue;
computing an image-to-image similarity between at least one image feature representation and the first plurality of image feature representations in the first feature queue;
computing a text-to-text similarity between at least one text feature representation all and the second plurality of text feature representations in the second feature queue;
computing a cross-modal alignment loss based on the image-to-image similarity and the text-to-text similarity; and
updating the image encoder and the text encoder based at least in part on the cross-modal alignment loss.
12. The method of claim 11 , wherein the image-to-image similarity is computed by:
computing a plurality of momentum image feature representations as the plurality of image feature representations are continuously added to the first feature queue, wherein at least one momentum image feature representation is computed as an exponential moving average of image feature representations that are currently in the first feature queue; and
computing the image-to-image similarity based on vector similarities between the at least one image feature representation and each of the plurality of momentum image feature representations.
13. The method of claim 11 , wherein the text-to-text similarity is computed by:
computing a plurality of momentum text feature representations as the plurality of text feature representations are continuously added to the second feature queue, wherein at least one momentum text feature representation is computed as an exponential moving average of text feature representations that are currently in the second feature queue; and
computing the text-to-text similarity based on vector similarities between the at least one text feature representation and each of the plurality of momentum text feature representations.
14. The method of claim 11 , wherein the cross-modal alignment loss is computed based on a first Kullback-Leibler divergence between the image-to-image similarity and the text- to-text similarity, and a second Kullback-Leibler divergence between the text-to-text similarity and the image-to-image similarity.
15. The method of claim 11 , further comprising:
applying, by a cross-modal encoder, cross-attention over the first plurality of image feature representations and second plurality of text feature representations.
16. The method of claim 11 , further comprising:
selecting a subset of samples from the first batch of image samples and the second batch of text samples; and
compute an image-text matching loss based on a similarity between an image sample and a text sample in an image-text pair, and a one-hot ground truth corresponding to the image- text pair.
17. The method of claim 16 , wherein the image-text pair is positive or negative, and
wherein the method further comprises:
building an image-text queue storing a plurality of image-text pairs from the first batch of image samples and the second batch of text samples, wherein the image-text queue is dynamically updated when new input samples are received;
selecting a negative sample corresponding to a sample in the subset from the image- text queue to form a negative image-text pair.
18. The method of claim 11 , further comprising:
randomly replacing text tokens in the second plurality of text samples with a mask token at a probability;
predicting replaced text tokens based on the cross-attention; and
computing a masked language modeling loss based on the predicted text tokens and original text tokens.
19. The method of claim 11 , further comprising:
computing an image-to-text similarity between the at least one image feature representation to all text feature representations in the second feature queue;
computing a text-to-image similarity between the at least one text feature representation to all image feature representations in the first feature queue;
computing a cross-modal contrastive loss based on a first cross-entropy between the image-to-text similarity and an image-to-text ground truth, and a second cross-entropy between the text-to-image similarity and a text-to-image ground truth;
computing a weighted sum of the cross-modal distribution alignment loss, the cross-modal contrastive loss, a masked language modeling loss and an image-text matching loss; and
jointly updating at least the image encoder and the text encoder based on the weighted sum.
20. A non-transitory processor-readable storage medium storing a plurality of processor-executable instructions for vision-language distribution alignment, the instructions being executed by a processor to perform operations comprising:
receiving, via a data interface, a first batch of image samples and a second batch of text samples;
encoding, by an image encoder, the first batch of image samples into a first plurality of image feature representations stored at a first feature queue;
encoding, by a text encoder, the second batch of text samples into a second plurality of text feature representations stored at a second feature queue;
computing an image-to-image similarity between at least one image feature representation and the first plurality of image feature representations in the first feature queue;
computing a text-to-text similarity between at least one text feature representation and the second plurality of text feature representations in the second feature queue;
computing a cross-modal alignment loss based on the image-to-image similarity and the text-to-text similarity; and
updating the image encoder and the text encoder based at least in part on the cross- modal alignment loss.