Label masking probabilistically removes class-label conditioning during diffusion training, allowing one model to learn both conditional and unconditional predictions for classifier-free guidance (CFG). section 6.2.1
Implement the mechanism in stages:
- Represent each class as a one-hot vector; the clothing example uses 10 classes. section 6.2.1
- Sample a Bernoulli context mask and use it to remove label information probabilistically. section 6.2.1
- Pass the masked vectors through two embedding blocks, producing label embeddings shaped (batch, 512, 1, 1) and (batch, 256, 1, 1). section 6.2.1
PyTorch label embeddings
For labels [2, 9, 0], create two 10-input embedding blocks, sample the context mask, convert labels to one-hot tensors, apply the mask, and reshape the outputs to 512- and 256-channel conditioning tensors. An exercise repeats this process for labels [1, 3, 5, 7]. section 6.2.1
Listing 6.3 Generating label embeddings
n_classes=10
contextembed1=EmbedLayer(n_classes, 512) #1
contextembed2=EmbedLayer(n_classes, 256) #1
label=torch.tensor([2, 9, 0]).long() #2
context_mask = torch.bernoulli(
torch.zeros_like(label)+0.1) #3
onehot = torch.nn.functional.one_hot(label,
num_classes=n_classes).type(torch.float) #4
context_mask = context_mask[:, None]
context_mask = context_mask.repeat(1,n_classes)
context_mask = (-1*(1-context_mask))
c = onehot * context_mask
cemb1 = contextembed1(c).view(-1, 512, 1, 1)
cemb2 = contextembed2(c).view(-1, 256, 1, 1)
print("the shape of the first label embedding is", cemb1.shape)
print("the shape of the second label embedding is", cemb2.shape)
Use the listing to trace mask sampling, one-hot conversion, label removal, embedding, and shape inspection. section 6.2.1
