A text encoder turns a caption or prompt into numerical embeddings. In CLIP, it works alongside an image encoder: both representations are aligned in a shared space so matching captions and images can be compared. The text encoder does not generate an image; a text-conditioned diffusion model can use its output to guide generation. section 8.1 section 8.0
To implement a CLIP-style caption encoder:
- Tokenize captions into input IDs and supply an attention mask to DistilBERT. section 8.3.1
- Take the first token’s representation from the final hidden layer. This CLS representation summarizes the caption through self-attention; a batch of 32 produces a raw output of shape [32, 768]. section 8.3.1
- Align text and image representations in a shared space. The illustrated CLIP setup uses 256-value embeddings and contrastive training that favors matching pairs over nonmatching pairs. section 8.1.1 section 8.5
Listing 8.7 Creating a text encoder
class TextEncoder(nn.Module):
def __init__(self, model_name=CFG.text_encoder_model,
pretrained=CFG.pretrained,
trainable=CFG.trainable):
super().__init__()
if pretrained:
self.model=DistilBertModel.from_pretrained(model_name)
else:
self.model=DistilBertModel(config=DistilBertConfig())
for p in self.model.parameters():
p.requires_grad = trainable
self.target_token_idx = 0 #1
def forward(self, input_ids, attention_mask):
output = self.model(input_ids=input_ids,
attention_mask=attention_mask)
last_hidden_state = output.last_hidden_state #2
return last_hidden_state[:,
self.target_token_idx,:] #3
The TextEncoder implementation shows pretrained or newly configured DistilBERT, an optional trainable setting, attention-mask input, and selection of the final-layer representation at token index 0. section 8.3.1
Retrieval versus generation
For retrieval, encode “A dog plays on the beach,” compare its embedding by cosine similarity with embeddings of 8,091 Flickr 8k images, and rank the matches. For Stable Diffusion, tokenize prompts such as “a lion” to a fixed 77-token length; the CLIP text encoder produces a [1, 77, 768] conditioning tensor. An empty prompt supplies the unconditional embedding used in classifier-free guidance. These are different uses and output forms, not interchangeable tensor shapes. section 8.1.2 section 10.3
The architecture depends on the application: the illustrated CLIP build uses DistilBERT for captions, Stable Diffusion uses a CLIP text encoder, and Imagen uses pretrained T5-XXL. In each case, inspect the expected token length and embedding shape before connecting the encoder to the next component. section 8.1.1 section 13.1.2 section 10.2.1
