import math from math import pi import typing from typing import Tuple, Optional, Literal from einops import rearrange, repeat import torch import torch.nn as nn import torch.nn.functional as F from torch.amp import autocast from torch.nn import Module, ModuleList from torch import nn, einsum, broadcast_tensors, Tensor ################################################################################# # Rotary Encoding # ################################################################################# # helper functions def exists(val): return val is not None def default(val, d): return val if exists(val) else d def slice_at_dim(t, dim_slice: slice, *, dim): dim += (t.ndim if dim < 0 else 0) colons = [slice(None)] * t.ndim colons[dim] = dim_slice return t[tuple(colons)] # rotary embedding helper functions def rotate_half(x): """Splits the last dimension of a tensor, swaps halves, and negates the first half.""" x = rearrange(x, '... (d r) -> ... d r', r=2) x1, x2 = x.unbind(dim=-1) x = torch.stack((-x2, x1), dim=-1) return rearrange(x, '... d r -> ... (d r)') @autocast('cuda', enabled=False) def apply_rotary_emb( freqs, t, start_index=0, scale=1., seq_dim=-2, freqs_seq_dim=None ): """Applies rotary positional embeddings to a given tensor. Args: freqs (torch.Tensor): The rotary frequencies. t (torch.Tensor): The tensor to apply embeddings to (e.g., queries or keys). start_index (int): The feature dimension index to start applying rotations from. scale (float): A scaling factor, used for xPos. seq_dim (int): The sequence dimension of the input tensor `t`. freqs_seq_dim (Optional[int]): The sequence dimension of the freqs tensor. """ dtype = t.dtype if not exists(freqs_seq_dim): if freqs.ndim == 2 or t.ndim == 3: freqs_seq_dim = 0 if t.ndim == 3 or exists(freqs_seq_dim): seq_len = t.shape[seq_dim] freqs = slice_at_dim(freqs, slice(-seq_len, None), dim=freqs_seq_dim) rot_dim = freqs.shape[-1] end_index = start_index + rot_dim assert rot_dim <= t.shape[ -1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}' # Split t into three parts: left, middle (to be transformed), and right t_left = t[..., :start_index] t_middle = t[..., start_index:end_index] t_right = t[..., end_index:] # Apply rotary embeddings without modifying t in place t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale) out = torch.cat((t_left, t_transformed, t_right), dim=-1) return out.type(dtype) # learned rotation helpers def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None): if exists(freq_ranges): rotations = einsum('..., f -> ... f', rotations, freq_ranges) rotations = rearrange(rotations, '... r f -> ... (r f)') rotations = repeat(rotations, '... n -> ... (n r)', r=2) return apply_rotary_emb(rotations, t, start_index=start_index) # classes class RotaryEmbedding(Module): """ original paper: https://arxiv.org/abs/2104.09864 rescale rotary embeddings to longer sequence length without fine-tuning code source: https://github.com/lucidrains/rotary-embedding-torch """ def __init__( self, dim, custom_freqs: Tensor | None = None, freqs_for: Literal['lang', 'pixel', 'constant'] = 'lang', theta=10000, max_freq=10, num_freqs=1, learned_freq=False, use_xpos=False, xpos_scale_base=512, interpolate_factor=1., theta_rescale_factor=1., seq_before_head_dim=False, cache_if_possible=True, cache_max_seq_len=8192 ): super().__init__() """Initializes the RotaryEmbedding module. Args: dim (int): The feature dimension to apply rotary embeddings to. custom_freqs ([Tensor]): An optional tensor of custom frequencies. freqs_for : The method for generating frequencies. 'lang' is standard for transformers. theta (int): A core hyperparameter for frequency calculation. learned_freq (bool): If True, the frequencies are trainable parameters. use_xpos (bool): If True, enables the xPos (extrapolatable) variant. interpolate_factor (float): A factor for positional interpolation, which can help with length generalization. cache_if_possible (bool): If True, caches calculated frequencies for efficiency. """ theta *= theta_rescale_factor ** (dim / (dim - 2)) self.freqs_for = freqs_for if exists(custom_freqs): freqs = custom_freqs elif freqs_for == 'lang': freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) elif freqs_for == 'pixel': freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi elif freqs_for == 'constant': freqs = torch.ones(num_freqs).float() self.cache_if_possible = cache_if_possible self.cache_max_seq_len = cache_max_seq_len self.register_buffer('cached_freqs', torch.zeros(cache_max_seq_len, dim), persistent=False) self.cached_freqs_seq_len = 0 self.freqs = nn.Parameter(freqs, requires_grad=learned_freq) self.learned_freq = learned_freq # dummy for device self.register_buffer('dummy', torch.tensor(0), persistent=False) # default sequence dimension self.seq_before_head_dim = seq_before_head_dim self.default_seq_dim = -3 if seq_before_head_dim else -2 # interpolation factors assert interpolate_factor >= 1. self.interpolate_factor = interpolate_factor # xpos self.use_xpos = use_xpos if not use_xpos: return scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) self.scale_base = xpos_scale_base self.register_buffer('scale', scale, persistent=False) self.register_buffer('cached_scales', torch.zeros(cache_max_seq_len, dim), persistent=False) self.cached_scales_seq_len = 0 # add apply_rotary_emb as static method self.apply_rotary_emb = staticmethod(apply_rotary_emb) @property def device(self): return self.dummy.device def get_seq_pos(self, seq_len, device, dtype, offset=0): return (torch.arange(seq_len, device=device, dtype=dtype) + offset) / self.interpolate_factor def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None): """Applies rotary embeddings to a single tensor (queries or keys). Args: t (torch.Tensor): The input tensor (queries or keys). seq_dim : The sequence dimension of the tensor. offset (int): An offset for the position sequence, used for caching. scale (Optional[float]): A scaling factor, required if using xPos. Returns: torch.Tensor: The tensor with rotary embeddings applied. """ seq_dim = default(seq_dim, self.default_seq_dim) assert not self.use_xpos or exists( scale), 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings' device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim] seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset) freqs = self.forward(seq, seq_len=seq_len, offset=offset) if seq_dim == -3: freqs = rearrange(freqs, 'n d -> n 1 d') return apply_rotary_emb(freqs, t, scale=default(scale, 1.), seq_dim=seq_dim) def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0): dtype, device, seq_dim = q.dtype, q.device, default(seq_dim, self.default_seq_dim) q_len, k_len = q.shape[seq_dim], k.shape[seq_dim] assert q_len <= k_len q_scale = k_scale = 1. if self.use_xpos: seq = self.get_seq_pos(k_len, dtype=dtype, device=device) q_scale = self.get_scale(seq[-q_len:]).type(dtype) k_scale = self.get_scale(seq).type(dtype) rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset) rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale ** -1) rotated_q = rotated_q.type(q.dtype) rotated_k = rotated_k.type(k.dtype) return rotated_q, rotated_k def rotate_queries_and_keys(self, q, k, seq_dim=None): seq_dim = default(seq_dim, self.default_seq_dim) assert self.use_xpos device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim] seq = self.get_seq_pos(seq_len, dtype=dtype, device=device) freqs = self.forward(seq, seq_len=seq_len) scale = self.get_scale(seq, seq_len=seq_len).to(dtype) if seq_dim == -3: freqs = rearrange(freqs, 'n d -> n 1 d') scale = rearrange(scale, 'n d -> n 1 d') rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim) rotated_k = apply_rotary_emb(freqs, k, scale=scale ** -1, seq_dim=seq_dim) rotated_q = rotated_q.type(q.dtype) rotated_k = rotated_k.type(k.dtype) return rotated_q, rotated_k def get_scale( self, t: Tensor, seq_len = None, offset=0 ): assert self.use_xpos should_cache = ( self.cache_if_possible and exists(seq_len) and (offset + seq_len) <= self.cache_max_seq_len ) if ( should_cache and \ exists(self.cached_scales) and \ (seq_len + offset) <= self.cached_scales_seq_len ): return self.cached_scales[offset:(offset + seq_len)] scale = 1. if self.use_xpos: power = (t - len(t) // 2) / self.scale_base scale = self.scale ** rearrange(power, 'n -> n 1') scale = repeat(scale, 'n d -> n (d r)', r=2) if should_cache and offset == 0: self.cached_scales[:seq_len] = scale.detach() self.cached_scales_seq_len = seq_len return scale def get_axial_freqs(self, *dims): Colon = slice(None) all_freqs = [] for ind, dim in enumerate(dims): if self.freqs_for == 'pixel': pos = torch.linspace(-1, 1, steps=dim, device=self.device) else: pos = torch.arange(dim, device=self.device) freqs = self.forward(pos, seq_len=dim) all_axis = [None] * len(dims) all_axis[ind] = Colon new_axis_slice = (Ellipsis, *all_axis, Colon) all_freqs.append(freqs[new_axis_slice]) all_freqs = broadcast_tensors(*all_freqs) return torch.cat(all_freqs, dim=-1) @autocast('cuda', enabled=False) def forward( self, t: Tensor, seq_len = None, offset=0 ): """Calculates the rotary frequencies for a given sequence of positions. Args: t (torch.Tensor): A tensor of position indices. seq_len (int): The total sequence length, used for caching. offset (int): The starting position offset. Returns: torch.Tensor: A tensor of calculated rotation frequencies. """ should_cache = ( self.cache_if_possible and not self.learned_freq and exists(seq_len) and self.freqs_for != 'pixel' and (offset + seq_len) <= self.cache_max_seq_len ) if ( should_cache and \ exists(self.cached_freqs) and \ (offset + seq_len) <= self.cached_freqs_seq_len ): return self.cached_freqs[offset:(offset + seq_len)].detach() freqs = self.freqs freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs) freqs = repeat(freqs, '... n -> ... (n r)', r=2) if should_cache and offset == 0: self.cached_freqs[:seq_len] = freqs.detach() self.cached_freqs_seq_len = seq_len return freqs ################################################################################# # Multi Head Attention # ################################################################################# class LayerNorm(nn.Module): """Implements a Layer Normalization module.""" def __init__(self, d_model, eps=1e-12): """Initializes the LayerNorm module. Args: d_model (int): The dimension of the model's features. eps (float): A small value added to the variance for numerical stability. """ super(LayerNorm, self).__init__() self.gamma = nn.Parameter(torch.ones(d_model)) self.beta = nn.Parameter(torch.zeros(d_model)) self.eps = eps def forward(self, x): """Applies Layer Normalization to the input tensor along the last dimension. Args: x (torch.Tensor): The input tensor to normalize. Returns: torch.Tensor: The normalized tensor. """ mean = x.mean(-1, keepdim=True) var = x.var(-1, unbiased=False, keepdim=True) # '-1' means last dimension. out = (x - mean) / torch.sqrt(var + self.eps) out = self.gamma * out + self.beta return out class PositionwiseFeedForward(nn.Module): """Implements the Position-wise Feed-Forward network of a Transformer block.""" def __init__(self, d_model, hidden, drop_prob=0.1): """Initializes the PositionwiseFeedForward module. Args: d_model (int): The input and output dimension of the layer. hidden (int): The dimension of the inner hidden layer. drop_prob (float): The probability for the dropout layer. """ super(PositionwiseFeedForward, self).__init__() self.linear1 = nn.Linear(d_model, hidden) self.linear2 = nn.Linear(hidden, d_model) self.relu = nn.ReLU() self.dropout = nn.Dropout(p=drop_prob) def forward(self, x): """Passes the input through the feed-forward network. The process is: Linear -> ReLU -> Dropout -> Linear. Args: x (torch.Tensor): The input tensor. Returns: torch.Tensor: The output tensor. """ x = self.linear1(x) x = self.relu(x) x = self.dropout(x) x = self.linear2(x) return x class ScaleDotProductAttention(nn.Module): def __init__(self): super(ScaleDotProductAttention, self).__init__() self.softmax = nn.Softmax(dim=-1) def forward(self, q, k, v, mask=None, e=1e-12): """ Performs the Scaled Dot-Product Attention calculation. Args: q (torch.Tensor): The query tensor. k (torch.Tensor): The key tensor. v (torch.Tensor): The value tensor. mask (torch.Tensor, optional): A mask to prevent attention to certain positions. Defaults to None. Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing the attention output and the attention scores. """ batch_size, head, length, d_tensor = k.size() k_t = k.transpose(2, 3) # transpose score = (q @ k_t) / math.sqrt(d_tensor) # scaled dot product if mask is not None: score = score.masked_fill(mask == 0, -10000) score = self.softmax(score) v = score @ v return v, score class MultiHeadAttention(nn.Module): """Implements a Multi-Head Attention layer with optional Rotary Position Embeddings.""" def __init__(self, d_model, n_head): """Initializes the MultiHeadAttention layer. Args: d_model (int): The total dimension of the model. n_head (int): The number of attention heads. d_model must be divisible by n_head. """ super(MultiHeadAttention, self).__init__() self.n_head = n_head self.attention = ScaleDotProductAttention() self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_concat = nn.Linear(d_model, d_model) self.rotary_emb = RotaryEmbedding(dim=d_model // n_head) def forward(self, q, k, v, mask=None, apply_rotary=False): """Performs the forward pass for multi-head attention. Args: q (torch.Tensor): The query tensor. k (torch.Tensor): The key tensor. v (torch.Tensor): The value tensor. mask (torch.Tensor, optional): An attention mask. Defaults to None. apply_rotary (bool): If True, applies Rotary Position Embeddings to Q and K before the attention calculation. Defaults to False. Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing the final output tensor and the attention scores. """ q, k, v = self.w_q(q), self.w_k(k), self.w_v(v) q, k, v = self.split(q), self.split(k), self.split(v) if apply_rotary: # add Rotary Positional Embeddings (RoPE) # https://arxiv.org/abs/2104.09864 q = self.rotary_emb.rotate_queries_or_keys(q) k = self.rotary_emb.rotate_queries_or_keys(k) out, attention = self.attention(q, k, v, mask=mask) out = self.concat(out) out = self.w_concat(out) return out, attention def split(self, tensor): """Splits the last dimension of a tensor into multiple heads.""" batch_size, length, d_model = tensor.size() d_tensor = d_model // self.n_head tensor = tensor.view(batch_size, length, self.n_head, d_tensor).transpose(1, 2) return tensor def concat(self, tensor): """Concatenates multiple heads back into a single tensor.""" batch_size, head, length, d_tensor = tensor.size() d_model = head * d_tensor tensor = tensor.transpose(1, 2).contiguous().view(batch_size, length, d_model) return tensor ################################################################################# # Embedding Layers # ################################################################################# class EmbeddingLayer(nn.Module): """A simple lookup-based embedding layer with Kaiming uniform initialization.""" def __init__(self, dim, vocab_dim): super().__init__() self.embedding = nn.Parameter(torch.empty((vocab_dim, dim))) torch.nn.init.kaiming_uniform_(self.embedding, a=math.sqrt(5)) def forward(self, x): """Looks up the embeddings for the given indices. Args: x (torch.Tensor): A tensor of integer indices. Returns: torch.Tensor: The corresponding embedding vectors. """ return self.embedding[x] class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. """ def __init__(self, hidden_size, frequency_embedding_size=256): """Initializes the TimestepEmbedder. Args: hidden_size (int): The final dimension of the timestep embedding. frequency_embedding_size (int): The number of frequencies to use for the sinusoidal embedding. """ super().__init__() self.mlp = nn.Sequential( nn.Linear(frequency_embedding_size, hidden_size, bias=True), nn.SiLU(), nn.Linear(hidden_size, hidden_size, bias=True)) self.frequency_embedding_size = frequency_embedding_size @staticmethod def timestep_embedding(t, dim, max_period=10000): """ Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. :return: an (N, D) Tensor of positional embeddings. """ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py half = dim // 2 freqs = torch.exp( - math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat( [embedding, torch.zeros_like(embedding[:, :1])], dim=-1) return embedding def forward(self, t): t_freq = self.timestep_embedding(t, self.frequency_embedding_size) t_emb = self.mlp(t_freq) return t_emb ################################################################################# # Decoder # ################################################################################# class DecoderLayer(nn.Module): """ code source: https://github.com/hyunwoongko/transformer """ def __init__(self, d_model, ffn_hidden, n_head, drop_prob): """Initializes the DecoderLayer. Args: d_model (int): The dimension of the model. ffn_hidden (int): The dimension of the hidden layer in the feed-forward network. n_head (int): The number of attention heads. drop_prob (float): The dropout probability. """ super(DecoderLayer, self).__init__() self.self_attention = MultiHeadAttention(d_model=d_model, n_head=n_head) self.norm1 = LayerNorm(d_model=d_model) self.dropout1 = nn.Dropout(p=drop_prob) self.enc_dec_attention = MultiHeadAttention(d_model=d_model, n_head=n_head) self.norm2 = LayerNorm(d_model=d_model) self.dropout2 = nn.Dropout(p=drop_prob) self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob) self.norm3 = LayerNorm(d_model=d_model) self.dropout3 = nn.Dropout(p=drop_prob) def forward(self, dec, enc, trg_mask, src_mask, return_attention=False): """Performs one forward pass of the decoder layer. Args: dec (torch.Tensor): The input tensor from the previous decoder layer. enc (torch.Tensor): The output tensor from the encoder (for conditioning). trg_mask (torch.Tensor): The mask for the decoder's self-attention. src_mask (torch.Tensor): The mask for the cross-attention. return_attention (bool): If True, returns the cross-attention weights. Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing the output tensor and the attention weights (or None). """ attention = None _x = dec x, _ = self.self_attention(q=dec, k=dec, v=dec, mask=trg_mask, apply_rotary=True) x = self.dropout1(x) x = self.norm1(x + _x) if enc is not None: _x = x if return_attention: x, attention = self.enc_dec_attention(q=x, k=enc, v=enc, mask=src_mask) else: x, _ = self.enc_dec_attention(q=x, k=enc, v=enc, mask=src_mask) x = self.dropout2(x) x = self.norm2(x + _x) _x = x x = self.ffn(x) x = self.dropout3(x) x = self.norm3(x + _x) return x, attention class Decoder_RoPE(nn.Module): """A decoder that uses Rotary Position Embeddings (RoPE). This model is designed for a diffusion task, taking a ligand sequence, a conditioning protein sequence, and a diffusion timestep (sigma) as input to predict the output logits for the ligand. """ def __init__(self, vocab_size, seq_emb_dim, hidden_size: int=640, nhead: int=8, n_layers: int=4, expand_feedforward: int=3, dropout: float=0.1): """Args: vocab_size (int): The size of the output vocabulary (e.g., ligand tokens). seq_emb_dim (int): The dimension of the input sequence embeddings. hidden_size (int): The main hidden dimension of the Transformer model. nhead (int): The number of attention heads in each DecoderLayer. n_layers (int): The number of DecoderLayers to stack. expand_feedforward (int): The expansion factor for the feed-forward network's hidden layer. dropout (float): The dropout probability. """ super().__init__() self.hidden_size = hidden_size self.vocab_embed = EmbeddingLayer(self.hidden_size, vocab_size) self.linear = nn.Linear(self.hidden_size, vocab_size) self.apply_seq_linear = False if seq_emb_dim != self.hidden_size: self.apply_seq_linear = True self.linear_seq = nn.Linear(seq_emb_dim, self.hidden_size) self.sigma_map = TimestepEmbedder(self.hidden_size) self.layers = nn.ModuleList([DecoderLayer(d_model=self.hidden_size, ffn_hidden=self.hidden_size * expand_feedforward, n_head=nhead, drop_prob=dropout) for _ in range(n_layers)]) def forward(self, ligand: torch.Tensor, sigma: torch.Tensor, sequence: torch.Tensor, sequence_lengths: torch.Tensor, lig_padding_mask: Optional[torch.Tensor]=None, return_attention: bool=False) -> Tuple[torch.Tensor, torch.Tensor]: """Performs the forward pass of the decoder. It processes the ligand sequence conditioned on the protein sequence and the diffusion timestep (sigma). The sigma embedding is prepended to the protein sequence to form a single conditioning context. Args: ligand (torch.Tensor): A batch of ligand token ID tensors. sigma (torch.Tensor): A batch of scalar diffusion timesteps. sequence (torch.Tensor): A batch of conditioning protein sequence embeddings. sequence_lengths (torch.Tensor): The original lengths of the protein sequences. lig_padding_mask (Optional[torch.Tensor]): A padding mask for the ligand. return_attention (bool): If True, returns the cross-attention weights from the last decoder layer. Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple of (output_logits, attention_weights). """ ligand = self.vocab_embed(ligand) sigma = F.silu(self.sigma_map(sigma)).unsqueeze(1) if self.apply_seq_linear: sequence = self.linear_seq(sequence) condition = torch.cat([sigma, sequence], dim=1) sequence_lengths += 1 range_tensor = torch.arange(condition.shape[1], device=sequence.device).unsqueeze(0) condition_mask = range_tensor < sequence_lengths.unsqueeze(1) condition_mask = condition_mask.unsqueeze(1).unsqueeze(2) if lig_padding_mask is not None: lig_padding_mask = lig_padding_mask.unsqueeze(1).unsqueeze(2) for layer in self.layers: ligand, attention = layer(ligand, condition, trg_mask=lig_padding_mask, src_mask=condition_mask, return_attention=return_attention) output = self.linear(ligand) return output, attention