DistilBERT is a smaller, faster, lighter version of BERT produced through knowledge distillation. A student model learns to mimic BERT while retaining much of its language understanding with fewer parameters. section 8.1.1
In a CLIP text encoder, DistilBERT processes captions through these stages:
- Load a pretrained DistilBERT model rather than training the text encoder from scratch. section 8.3.1
- Pass token IDs and an attention mask through the model, whose successive layers produce increasingly contextualized representations. section 8.3.1
- Extract the first token from the final hidden state; this CLS token acts as a learned summary of the entire caption. section 8.3.1
- Use the resulting fixed-size text vector alongside image embeddings in a shared semantic space for contrastive alignment. section 8.1.1
Frozen caption encoder
A DistilBERT text encoder can freeze all model parameters, process a batch of 32 captions, and return an embedding tensor shaped [32, 768]. section 8.3.1
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
Inspect the encoder implementation showing pretrained construction, parameter freezing, final hidden-state extraction, and CLS-token selection. section 8.3.1
