Text-based image retrieval
A method, apparatus, non-transitory computer readable medium, and system for media processing include obtaining a text prompt describing content, generating, using a multi-modal encoder, a text embedding based on the text prompt, and obtaining an image depicting the content based on the text embedding. The multi-modal encoder is trained to encode image descriptions based on a similarity between a caption of a training image and a paraphrase of the caption.
1 . A method for media processing, comprising:
obtaining a text prompt describing content;
generating, using a text encoder of a multi-modal encoder, a text embedding in a multi-modal embedding space based on the text prompt, wherein the text encoder is trained to encode image descriptions in the multi-modal embedding space by computing a paraphrase-caption loss based on a similarity between a caption of a training image and a paraphrase of the caption and updating parameters of the text encoder based on the paraphrase-caption loss, wherein the caption comprises a first text string describing the training image and the paraphrase comprises a second text string describing the training image, and wherein the first text string is different from the second text string; and
obtaining an image depicting the content based on the text embedding.
2 . The method of claim 1 , wherein obtaining the image comprises:
identifying an image embedding of the image; and
retrieving the image from a database based on a comparison of the text embedding and the image embedding.
3 . The method of claim 2 , wherein:
the text embedding and the image embedding comprise vectors in the multi-modal embedding space.
4 . The method of claim 1 , further comprising:
retrieving a plurality of images from a database based on the text embedding.
5 . The method of claim 1 , wherein obtaining the image comprises:
generating the image using an image generation model conditioned on the text embedding.
6 . The method of claim 1 , further comprising:
tokenizing the text prompt to obtain a sequence of tokens representing the content, wherein the text embedding is generated based on the sequence of tokens.
7 . A method for training a machine learning model, comprising:
obtaining a training set comprising a training image, a caption of the training image, and a paraphrase of the caption, wherein the caption comprises a first text string describing the training image and the paraphrase comprises a second text string describing the training image, and wherein the first text string is different from the second text string;
encoding, using an image encoder of a multi-modal encoder, the training image to obtain an image embedding in a multi-modal embedding space;
encoding, using a text encoder of the multi-modal encoder, the caption and the paraphrase to obtain a caption embedding and a paraphrase embedding, respectively, in the multi-modal embedding space; and
training the text encoder of the multi-modal encoder by computing an image-caption loss based on a first similarity between the image embedding and the caption embedding, computing a paraphrase-caption loss based on a second similarity between the caption embedding and the paraphrase embedding, and updating parameters of the text encoder based on the image-caption loss and the paraphrase-caption loss.
8 . The method of claim 7 , wherein training the text encoder comprises:
computing a paraphrase-paraphrase loss based on a third similarity between the paraphrase embedding and an additional paraphrase embedding of an additional paraphrase of the paraphrase and updating the parameters of the text encoder based on the paraphrase-paraphrase loss.
9 . The method of claim 7 , wherein obtaining the training set comprises:
generating the caption based on the training image.
10 . The method of claim 7 , wherein obtaining the training set comprises:
generating the paraphrase based on the caption.
11 . The method of claim 10 , wherein generating the paraphrase comprises:
generating a prompt requesting a variant of the caption using different language; and
providing the prompt to a large language model.
12 . The method of claim 7 , wherein obtaining the training set comprises:
generating an additional paraphrase based on the paraphrase.
13 . The method of claim 7 , wherein training the text encoder comprises:
fine-tuning a pre-trained multi-modal encoder.
14 . The method of claim 7 , wherein training the text encoder comprises:
freezing the image encoder.
15 . A system for media processing, comprising:
at least one processor;
at least one memory storing instructions executable by the at least one processor; and
a multi-modal encoder comprising a text encoder comprising encoding parameters stored in the at least one memory, the text encoder configured to generate a text embedding in a multi-modal embedding space based on a text prompt, wherein the text encoder is trained to encode image descriptions in the multi-modal embedding space by computing a paraphrase-caption loss based on a similarity between a caption of a training image and a paraphrase of the caption and updating parameters of the text encoder based on the paraphrase-caption loss, wherein the caption comprises a first text string describing the training image and the paraphrase comprises a second text string describing the training image, and wherein the first text string is different from the second text string.
16 . The system of claim 15 , the system further comprising:
a language generation model comprising text generation parameters stored in the at least one memory, the language generation model trained to generate the paraphrase.
17 . The system of claim 15 , the system further comprising:
a database storing an image embedding; and
a retrieval component configured to retrieve an image from the database based on the text embedding and the image embedding.
18 . The system of claim 15 , the system further comprising:
an image generation model comprising image generation parameters stored in the at least one memory, the image generation model trained to generate an image based on the text embedding.