OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data

Represent words and images as vectors

Released in 2022, the OpenCLIP (Contrastive Language–Image Pre–training) family of transformer-based neural nets is a collection of models trained as pure feature extractors learning joint text and image representations from scratch. Utilizing the LAION-5B dataset, which contains five billion image-text pairs, the authors examine the effects of scaling on the performance of OpenCLIP models. The study reveals that performance consistently improves with the scaling of model size, data and computational resources, adhering to a power law. Interestingly, OpenCLIP outperforms in zero-shot retrieval tasks, while the original OpenAI CLIP models excel in zero-shot classification. The authors hypothesize that the training dataset significantly influences these task-specific scaling differences.

Training Set Information

LAION-5B, containing 5.85 billion CLIP-filtered image-text pairs, 14 times bigger than LAION-400M, previously the biggest openly accessible image-text dataset in the world.

Model Information

Examples

Download Example Notebook

Open in Wolfram Cloud

Resource retrieval

Get the pre-trained net:

In[1]:=

Out[1]=

NetModel parameters

This model consists of a family of individual nets, each identified by a specific parameter combination. Inspect the available parameters:

In[2]:=

Out[2]=

Pick a non-default net by specifying the parameters:

In[3]:=

Out[3]=

Pick a non-default uninitialized net:

In[4]:=

Out[4]=

Basic usage

Use the OpenCLIP text encoder to obtain the feature representation of a piece of text:

In[5]:=

textEmbedding = NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text"}]["Challenge accepted!"];

The default OpenCLIP text encoder embeds the input text into a vector of size 512:

In[6]:=

Out[6]=

Use the OpenCLIP image encoder to obtain the feature representation of an image:

In[7]:=

(* Evaluate this cell to get the example input *) CloudGet["https://www.wolframcloud.com/obj/f96193b5-d0d1-4daa-9e6a-d94b25133718"]

The default OpenCLIP image encoder embeds the input text into a vector of size 512:

In[8]:=

Out[8]=

Feature space visualization

Get a set of images:

In[9]:=

(* Evaluate this cell to get the example input *) CloudGet["https://www.wolframcloud.com/obj/c73d854f-5ce6-4b76-a3c4-46c6bcd941d1"]

Visualize the features of a set of images:

In[10]:=

FeatureSpacePlot[
Thread[NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Image"}][imgs, NetPort[{"embed", "Output"}]] -> imgs],
LabelingSize -> 70,
LabelingFunction -> Callout,
ImageSize -> 600,
AspectRatio -> 0.9,
RandomSeeding -> 123456
]

Out[10]=

Define a list of sentences in two broad categories:

In[11]:=

sentences = {
"The Empire State Building's observation deck in New York is a must-visit for its iconic skyline views.",
"The Charging Bull in the financial district of New York has become a symbol of market optimism.",
"Times Square in New York is best known for its bright billboards and bustling atmosphere.",
"The Statue of Liberty in New York stands as a universal symbol of freedom and opportunity.",
"Central Park in New York is an urban oasis, providing a natural escape amidst the city's skyscrapers.",
"Sacré-Cœur in Paris offers both spiritual solace and panoramic views from its hilltop location.",
"The Eiffel Tower's light in Paris show adds a romantic touch to the city's engineering marvel.",
"Bridges over the Seine in Paris are scenic spots that often host art and book vendors.",
"The Louvre's glass pyramid in Paris modernizes the entrance to a museum filled with historical art.",
"The Panthéon in Paris serves as a tribute to national heroes, complete with educational exhibits."
};

Visualize the similarity between the sentences using the net as a feature extractor:

In[12]:=

FeatureSpacePlot[
Thread[NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text"}][sentences, NetPort[{"embed", "Output"}]] -> sentences]
, LabelingFunction -> Callout, LabelingSize -> {120, 70}, ImageSize -> 600, PlotMarkers -> {Automatic, Scaled[0.01]}, RandomSeeding -> 123465]

Out[12]=

Connecting text and images

Define a test image:

In[13]:=

(* Evaluate this cell to get the example input *) CloudGet["https://www.wolframcloud.com/obj/a8cc9424-a808-494c-9878-45322e441c91"]

Define a list of text descriptions:

In[14]:=

descriptions = {
"Blossoming rose on textbook among delicate petals",
"Photo of foggy forest",
"A portrait of a man with red hair taking a picture with an analog camera",
"Yellow flower in tilt shift lens",
"Woman in black leather jacket and white pants",
"A portrait of a woman with red hair taking a picture with an analog camera",
"Calm body of lake between mountains",
"Close up shot of a baby girl",
"Smiling man surfing on wave in ocean",
"A portrait of a woman with red hair taking a picture with a phone",
"A portrait of a woman with red hair taking a picture with a digital camera",
"A woman with eyeglasses smiling",
"Elderly woman carrying a baby",
"A portrait of a woman with blue hair taking a picture with an analog camera"
};

Embed the test image and text descriptions into the same feature space:

In[15]:=

textFeatures = NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text", "Architecture" -> "ViT-B/16"}][descriptions];
imgFeatures = NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Image", "Architecture" -> "ViT-B/16"}][img];

Rank the text description with respect to the correspondence to the input image according to the CosineDistance. Smaller distances mean higher correspondence between the text and the image:

In[16]:=

Dataset@SortBy[
Thread[{descriptions, First@DistanceMatrix[{imgFeatures}, textFeatures, DistanceFunction -> CosineDistance]}], Last]

Out[16]=

Zero-shot image classification

By using the text and image feature extractors together, it's possible to perform generic image classification between any set of classes without having to explicitly train any model for those particular classes (zero-shot classification). Obtain the FashionMNIST test data, which contains ten thousand test images and 10 classes:

In[17]:=

Display a few random examples from the set:

In[18]:=

Out[18]=

Get a mapping between class IDs and labels:

In[19]:=

Out[19]=

Generate the text templates for the FashionMNIST labels and embed them. The text templates will effectively act as classification labels:

In[20]:=

Out[20]=

In[21]:=

textEmbeddings = NetModel[
"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text"][labelTemplates];

In[22]:=

Out[22]=

Classify an image from the test set. Obtain its embedding:

In[23]:=

Out[23]=

In[24]:=

imgFeatures = NetModel[
"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Image"][img];

In[25]:=

Out[25]=

The result of the classification is the description of the embedding that is closest to the image embedding:

In[26]:=

Out[26]=

Find the top 10 description nearest to the image embedding:

In[27]:=

SortBy[Rule @@@ MapAt[labelTemplates[[#]] &, Nearest[textEmbeddings -> {"Index", "Distance"}, imgFeatures, 10, DistanceFunction -> CosineDistance], {All, 1}], Last]

Out[27]=

Obtain the accuracy of this procedure on the entire test set. Extract the features for all the images (if a GPU is available, setting TargetDevice -> "GPU" is recommended as the computation will take several minutes on CPU):

In[28]:=

imageEmbeddings = NetModel[
"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Image"][testData[[All, 1]], TargetDevice -> "CPU"];

In[29]:=

Out[29]=

Calculate the distance matrix between the computed text and image embeddings:

In[30]:=

Obtain the top-1 predictions:

In[31]:=

Obtain the final classification results:

In[32]:=

Out[32]=

Attention visualization for images

Just like the original Vision Transformer (see the model "Vision Transformer Trained on ImageNet Competition Data"), the image feature extractor divides the input images in 7x7 patches and performs self-attention on a set of 50 vectors: 49 vectors, or "tokens," representing the 7x7 patches and an additional one, a "feature extraction token," that is eventually used to produce the final feature representation of the image. Thus the attention procedure for this model can be visualized by inspecting the attention weights between the feature extraction token and the patch tokens. Define a test image:

In[33]:=

(* Evaluate this cell to get the example input *) CloudGet["https://www.wolframcloud.com/obj/80b0126c-3262-4c5e-8195-fc94034a0021"]

Extract the attention weights used for the last block of self-attention:

In[34]:=

attentionMatrix = Transpose@
NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Image"}][testImage, NetPort[{"transformer", -1, "self-attention", "attention", "AttentionWeights"}]];

In[35]:=

Out[35]=

Extract the attention weights between the feature extraction token and the input patches. These weights can be interpreted as which patches in the original image the net is "looking at" in order to perform the feature extraction:

In[36]:=

Out[37]=

Reshape the weights as a 3D array of 12 7x7 matrices. Each matrix corresponds to an attention head, while each element of the matrices corresponds to a patch in the original image:

In[38]:=

In[39]:=

Out[39]=

Visualize the attention weight matrices. Patches with higher values (red) are what is mostly being "looked at" for each attention head:

In[40]:=

Out[40]=

Define a function to visualize the attention matrix on an image:

In[41]:=

$visualizeAttention[img_Image, attentionMatrix_] := Block[{heatmap, wh}, wh = ImageDimensions[img]; heatmap = ImageApply[{#, 1 - #, 1 - #} &, ImageAdjust@Image[attentionMatrix]]; heatmap = ImageResize[heatmap, wh*256/Min[wh]]; ImageCrop[ImageCompose[img, {ColorConvert[heatmap, "RGB"], 0.4}], ImageDimensions[heatmap]] ]$

Visualize the mean attention across all the attention heads:

In[42]:=

Out[42]=

Visualize each attention head separately:

In[43]:=

Out[43]=

Attention visualization for text

The text feature extractor tokenizes the input string prepending and appending the special tokens StartOfString and EndOfString and then performs causal self-attention on the token embedding vectors. After the self-attention stack, the last vector (corresponding to the token EndOfString) is used to obtain the final feature representation of the text. Thus the attention procedure for this model can be visualized by inspecting the attention weights between the last vector and the previous ones. Define a test string:

In[44]:=

Extract the NetEncoder of the net to encode the string:

In[45]:=

netEnc = NetExtract[
NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text"}], "Input"]

Out[45]=

In[46]:=

Out[46]=

Extract the list of available tokens and inspect how the input string was tokenized. Even though the BPE tokenization generally segments the input into subwords, it's common to observe that all tokens correspond to full words. Also observe that the StartOfString and EndOfString tokens are added automatically:

In[47]:=

In[48]:=

Out[48]=

In[49]:=

Out[49]=

Feed the string to the net and extract the attention weights used for the last block of self-attention:

In[50]:=

attentionMatrix = Transpose@
NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text"}][text, NetPort[{"transformer", -1, "self-attention", "attention", "AttentionWeights"}]];
Dimensions[attentionMatrix]

Out[51]=

Extract the attention weights between the last vector and the previous ones, leaving the initial vector corresponding to StartOfString out. These weights can be interpreted as which tokens in the original sentence the net is "looking at" in order to perform the feature extraction:

In[52]:=

Out[53]=

Inspect the average attention weights for each token across the attention heads. Observe that the token the net is mostly focused on is "hair":

In[54]:=

Out[54]=

Visualize each head separately:

In[55]:=

Out[55]=

Extract the attention weights for all 12 attention layers:

In[56]:=

In[57]:=

allAttentionWeights = Transpose[
Values@NetModel[{"OpenCLIP Multi-domain Feature Extractor Trained on LAION-2B Data", "InputDomain" -> "Text"}][text, spec], 2 <-> 3];
Dimensions[allAttentionWeights]

Out[58]=

Compute the average across all heads, leaving the StartOfString token out:

In[59]:=

Out[60]=

Define a function to visualize the attention weights:

In[61]:=

$visualizeTokenAttention[attnMatrix_] := Block[{g, style}, g = WeightedAdjacencyGraph[attnMatrix]; style = Thread@Directive[ Arrowheads[.02], Thickness /@ (Rescale[AnnotationValue[g, EdgeWeight]]/200), Opacity /@ Rescale@AnnotationValue[g, EdgeWeight] ]; Graph[g, GraphLayout -> "LinearEmbedding", EdgeStyle -> Thread[EdgeList[g] -> style], VertexLabels -> Thread[Range[Length@Rest[tokens]] -> Map[Rotate[Style[Text[#], 12, Bold], 60 Degree] &, Rest[tokens]]], VertexCoordinates -> Table[{i, 0}, {i, Length[attnMatrix]}], ImageSize -> Large] ]$

Explore the attention weights for every layer. A thicker arrow pointing from token A to token B indicates that the layer is paying attention to token B when generating the vector corresponding to token A:

In[62]:=