Lesson · 40 min · Free
Multi-Head Attention in Generative Models
Multi-Head Attention in Generative Models body { font-family: sans-serif; line-height: 1.6; margin: 20px; } h1, h2 { color: #2c3e50; } h1 { border-bottom: 2px solid #3498db; padding-bottom: 10px; } h2 { color: #34495e; }
Multi-Head Attention in Generative Models
Welcome to this lesson on Multi-Head Attention, a pivotal mechanism in modern generative models, particularly those based on the Transformer architecture. While initially popularized in natural language processing (NLP), its principles have found profound applications in various scientific domains, including drug discovery. For pharmacy and biotech students, understanding this concept is crucial for appreciating the capabilities of advanced AI in tasks like de novo molecular design, protein structure prediction, and reaction pathway generation. Generative models aim to learn the underlying distribution of a dataset to produce new, similar data points. In drug discovery, this could mean generating novel small molecules with desired pharmacological properties, or predicting the 3D structure of a protein from its amino acid sequence. Traditional recurrent neural networks (RNNs) and convolutional neural networks (CNNs) have limitations in capturing long-range dependencies in sequential data, which is where the attention mechanism excels.
The Power of Multi-Head Attention
At its core, the attention mechanism allows a model to weigh the importance of different parts of the input sequence when processing a particular element. Instead of processing information sequentially, attention enables the model to "look at" and focus on relevant parts of the input simultaneously. This is particularly powerful for complex biological sequences (e.g., DNA, RNA, protein sequences) or molecular graphs where relationships between distant elements can be critical. Single-head attention computes a weighted sum of "value" vectors, where the weights are derived from the similarity between a "query" vector and "key" vectors. The "query," "key," and "value" vectors are typically linear transformations of the input embeddings. The multi-head aspect extends this by running several attention mechanisms in parallel. Each "head" learns to attend to different parts of the input or different types of relationships. The outputs from these multiple heads are then concatenated and linearly transformed to produce the final output. Why is "multi-head" important? Imagine you are trying to understand a complex drug molecule. One head might learn to focus on the pharmacophore, another on potential metabolic liabilities, and yet another on structural rigidity. By having multiple heads, the model can capture a richer and more diverse set of relationships and dependencies within the data, leading to a more comprehensive understanding and better generative capabilities. This parallel processing also helps the model focus on different "representation subspaces" at different positions. In the context of drug discovery, multi-head attention in generative models like Transformers (e.g., in models like AlphaFold for protein folding or various molecular generation models) allows for: Capturing long-range dependencies: Essential for understanding how distant amino acid residues interact in a protein, or how different functional groups in a large molecule influence each other. Parallel processing: Unlike RNNs, Transformers can process all elements in a sequence simultaneously, significantly speeding up training and inference, especially for long sequences. Learning diverse relationships: Each attention head can specialize in identifying different types of patterns or relationships within the input, leading to a more robust and nuanced representation. Improved interpretability: While not perfectly transparent, analyzing attention weights can sometimes offer insights into which parts of a molecule or sequence the model considers important for a given prediction or generation task. Let's look at a simplified conceptual implementation of a single attention head, which forms the building block of multi-head attention: import torch import torch.nn as nn import torch.nn.functional as F class SingleHeadAttention(nn.Module): def __init__(self, embed_dim): super().__init__() self.query = nn.Linear(embed_dim, embed_dim) self.key = nn.Linear(embed_dim, embed_dim) self.value = nn.Linear(embed_dim, embed_dim) self.scale = embed_dim ** -0.5 # Scaling factor for dot product attention def forward(self, x): # x is typically (batch_size, sequence_length, embed_dim) Q = self.query(x) # (batch_size, sequence_length, embed_dim) K = self.key(x) # (batch_size, sequence_length, embed_dim) V = self.value(x) # (batch_size, sequence_length, embed_dim) # Compute attention scores: Q * K_transpose # (batch_size, sequence_length, embed_dim) @ (batch_size, embed_dim, sequence_length) # Result: (batch_size, sequence_length, sequence_length) attention_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale # Apply softmax to get attention weights attention_weights = F.softmax(attention_scores, dim=-1) # Weighted sum of values: attention_weights * V # (batch_size, sequence_length, sequence_length) @ (batch_size, sequence_length, embed_dim) # Result: (batch_size, sequence_length, embed_dim) output = torch.matmul(attention_weights, V) return output, attention_weights # Example usage: # embed_dim = 128 # sequence_length = 10 # batch_size = 4 # # model = SingleHeadAttention(embed_dim) # input_data = torch.randn(batch_size, sequence_length, embed_dim) # output, weights = model(input_data) # print("Output shape:", output.shape) # Expected: (4, 10, 128) # print("Attention weights shape:", weights.shape) # Expected: (4, 10, 10) Now, let's see how multiple such heads are combined to form Multi-Head Attention: import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" self.num_heads = num_heads self.head_dim = embed_dim // num_heads # Dimension of each head # Linear layers for Query, Key, Value for all heads combined self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3) self.output_proj = nn.Linear(embed_dim, embed_dim) self.scale = self.head_dim ** -0.5 def forward(self, x): batch_size, sequence_length, embed_dim = x.shape # Project input to Q, K, V for all heads # qkv_proj(x) -> (batch_size, sequence_length, embed_dim * 3) # Then split into Q, K, V qkv = self.qkv_proj(x).chunk(3, dim=-1) # Returns a tuple of 3 tensors Q, K, V = [t.view(batch_size, sequence_length, self.num_heads, self.head_dim).transpose(1, 2) for t in qkv] # Q, K, V now have shape (batch_size, num_heads, sequence_length, head_dim) # Compute attention scores # (batch_size, num_heads, sequence_length, head_dim) @ (batch_size, num_heads, head_dim, sequence_length) attention_scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale # Apply softmax to get attention weights attention_weights = F.softmax(attention_scores, dim=-1) # Weighted sum of values # (batch_size, num_heads, sequence_length, sequence_length) @ (batch_size, num_heads, sequence_length, head_dim) # Result: (batch_size, num_heads, sequence_length, head_dim) attended_values = torch.matmul(attention_weights, V) # Concatenate heads and project back to original embedding dimension # First, transpose to (batch_size, sequence_length, num_heads, head_dim) # Then, reshape to (batch_size, sequence_length, embed_dim) attended_values = attended_values.transpose(1, 2).contiguous().view(batch_size, sequence_length, embed_dim) # Final linear projection output = self.output_proj(attended_values) return output, attention_weights # Example usage: # embed_dim = 256 # num_heads = 8 # sequence_length = 50 # E.g., length of a protein sequence or SMILES string # batch_size = 16 # # model = MultiHeadAttention(embed_dim, num_heads) # input_data = torch.randn(batch_size, sequence_length, embed_dim) # output, weights = model(input_data) # print("Multi-Head Attention Output shape:", output.shape) # Expected: (16, 50, 256) # print("Multi-Head Attention weights shape:", weights.shape) # Expected: (16, 8, 50, 50) In these examples, embed_dim refers to the dimensionality of the input embedding for each token/element in the sequence. For drug discovery, these embeddings could represent chemical substructures, amino acid residues, or even atoms in a molecular graph. The output of the multi-head attention layer is a new sequence of embeddings, where each element's embedding has been informed by all other elements in the sequence, weighted by their learned importance.
Key Takeaways:
Multi-Head Attention allows generative models to weigh the importance of different input elements when generating or processing information. It operates by running multiple "attention heads" in parallel, each focusing on different aspects or relationships within the data. This mechanism is crucial for capturing long-range dependencies, which are prevalent in biological sequences and molecular structures. It enables parallel processing, overcoming the sequential limitations of traditional RNNs. In drug discovery, Multi-Head Attention underpins advanced generative models for tasks like de novo molecular design, protein structure prediction, and reaction pathway synthesis.
Practice Exercise:
Consider a generative model designed to propose novel small molecules in the SMILES format. If a molecule has 50 atoms (represented as tokens in the SMILES string), and your Multi-Head Attention layer uses 4 heads with an embedding dimension of 128, describe the shape of the query, key, and value tensors after they have been split for individual heads but before the matrix multiplication for attention scores. Explain why having 4 heads might be more beneficial than a single head of dimension 128 for this task, specifically considering the diverse chemical features of a small molecule.
Watch the full lesson — free
This topic is part of AI in Drug Discovery, a complete AI-narrated video course. Press play once and watch the entire lecture like a movie.
Start the course free →