A transformer block combines attention sublayers with a feed-forward network to process a sequence. The same generic block can support either an encoder or a decoder. section 4.3.1
Trace the block by its role in each architecture:
- Process encoder inputs through self-attention and a feed-forward network. section 4.3.1
- Process decoder inputs through causal self-attention, cross-attention to encoder outputs, and a feed-forward network. section 4.3.1
- Stack multiple blocks to form larger encoders and decoders; the image-to-text model uses six of each. section 4.3.1
Vision encoder
Convert an image into patch embeddings, add positional encodings, and pass the resulting sequence through repeated TransformerBlock instances configured with decoder=False and masking=False. section 4.3.1
Listing 4.4 Defining a generic transformer block
class TransformerBlock(nn.Module):
def __init__(self, hidden_size=128, num_heads=4,
decoder=False, masking=True):
super(TransformerBlock, self).__init__()
self.decoder = decoder
self.norm1 = nn.LayerNorm(hidden_size)
self.attn1 = AttentionBlock(hidden_size=hidden_size,
num_heads=num_heads,
masking=masking)
if self.decoder:
self.norm2 = nn.LayerNorm(hidden_size)
self.attn2 = AttentionBlock(hidden_size=hidden_size,
num_heads=num_heads, masking=False)
self.norm_mlp = nn.LayerNorm(hidden_size)
self.mlp = nn.Sequential(nn.Linear(hidden_size, hidden_size * 4),
nn.ELU(),
nn.Linear(hidden_size * 4, hidden_size))
def forward(self, x, input_key_mask=None,
cross_key_mask=None, kv_cross=None):
x = self.attn1(x, x, key_mask=input_key_mask) + x #1
x = self.norm1(x)
if self.decoder: #2
x = self.attn2(x, kv_cross,key_mask=cross_key_mask) + x
x = self.norm2(x)
x = self.mlp(x) + x #3
return self.norm_mlp(x)
Inspect the generic PyTorch TransformerBlock implementation, including layer normalization, attention modules, decoder-specific cross-attention, and the feed-forward network. section 4.3.1
