A label embedding is a learned vector or tensor representation of a class label. In a conditional denoising U-Net, it converts clothing-item class information into feature maps that guide image denoising. section 6.0 section 6.3
The mechanism proceeds through four implementation stages:
- Encode the integer label as a one-hot vector with one entry for each class. section 6.2.1
- Apply an embedding block containing two linear layers with a GELU activation between them. section 6.2.1
- Create embeddings at multiple U-Net feature sizes, such as (512, 1, 1) and (256, 1, 1), so they can enter different up blocks. section 6.2.1
- Add the label embeddings to upsampling features alongside time embeddings, conditioning predicted noise on both class and diffusion step. section 6.3.1
Embedding Fashion MNIST labels
With ten classes and labels [2, 9, 0], one-hot encoding produces three class vectors. Two EmbedLayer instances transform them into tensors shaped [3, 512, 1, 1] and [3, 256, 1, 1]. A randomly applied context mask can suppress the class information for unconditional training. section 6.2.1
Listing 6.1 Defining an embedding block
from torch import nn
class EmbedLayer(nn.Module):
def __init__(self, input_dim, emb_dim): #1
super().__init__()
self.input_dim = input_dim
layers = [nn.Linear(input_dim, emb_dim),
nn.GELU(),
nn.Linear(emb_dim, emb_dim),] #2
self.model = nn.Sequential(*layers)
def forward(self, x):
x = x.view(-1, self.input_dim)
return self.model(x)
Use the embedding-block listing to trace the label transformation from input vector through GELU-activated linear layers. section 6.2.1
