![\"Image](\"https://pytorch.org/docs/stable/_images/Softplus.png\")
\n Softplus function
\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/ribonanza_diagram.png\")
\n\n### References \n\n- [RibonanzaNet code][1]\n- [How does DeepMind AlphaFold2 work?][2]\n- [AlphaFold v2 Github Repository][3]\n- [Ribonanza paper][4]\n- [AlphaFold2 complementary paper][5]\n\n[1]: https://github.com/Shujun-He/RibonanzaNet/blob/main/Network.py\n[2]: https://borisburkov.net/2021-12-25-1/\n[3]: https://github.com/google-deepmind/alphafold\n[4]: https://www.biorxiv.org/content/10.1101/2024.02.24.581671v1.full.pdf\n[5]: https://static-content.springer.com/esm/art%3A10.1038%2Fs41586-021-03819-2/MediaObjects/41586_2021_3819_MOESM1_ESM.pdf","metadata":{}},{"cell_type":"markdown","source":"# | Install Libraries [β](#top) \n\n***\n\nInstall libraries.","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Import Libraries [β](#top) \n\n***\n\nImport libraries.","metadata":{}},{"cell_type":"code","source":"import math\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport yaml\n\n\nfrom einops import rearrange, repeat, reduce\nfrom einops.layers.torch import Rearrange\nfrom functools import partialmethod\nfrom torch import einsum\nfrom torch.nn.parameter import Parameter\nfrom typing import Any, Dict, List, Literal, Optional, Tuple, Union","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:53.002530Z","iopub.execute_input":"2025-03-13T23:30:53.002896Z","iopub.status.idle":"2025-03-13T23:30:58.205640Z","shell.execute_reply.started":"2025-03-13T23:30:53.002868Z","shell.execute_reply":"2025-03-13T23:30:58.204952Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Configuration [β](#top) \n\n***\n\nThe next cell writes a YAML file with the different model parameters as well as other training and data configurations needed for the training stage.\n\nWe will use this file when instantiating the model and it is required to define the model's architecture.","metadata":{}},{"cell_type":"code","source":"%%writefile config.yaml\n\nlearning_rate: 0.001 # The learning rate for the optimizer\nbatch_size: 4 # Number of samples per batch\ntest_batch_size: 8 # Number of samples per batch\nepochs: 30 # Total training epochs\noptimizer: \"ranger\" # Optimization algorithm\ndropout: 0.05 # Dropout regularization rate\nweight_decay: 0.0001\nk: 5\nninp: 256\nnlayers: 9\nnclass: 2\nntoken: 5 # AUGC + padding/N token\nnhead: 8\nuse_bpp: False\nuse_flip_aug: true\nbpp_file_folder: \"../../input/bpp_files/\"\ngradient_accumulation_steps: 1\nuse_triangular_attention: false\npairwise_dimension: 64\nuse_bpp: False\n\n#Data scaling\nuse_data_percentage: 1\nuse_dirty_data: true # turn off for data scaling and data dropout experiments\n\n# Other configurations\nfold: 0\nnfolds: 6\ninput_dir: \"../../input/\"\ngpu_id: \"0\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.206711Z","iopub.execute_input":"2025-03-13T23:30:58.207124Z","iopub.status.idle":"2025-03-13T23:30:58.212521Z","shell.execute_reply.started":"2025-03-13T23:30:58.207094Z","shell.execute_reply":"2025-03-13T23:30:58.211895Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Utils [β](#top) \n\n***\n\nUtility functions.","metadata":{}},{"cell_type":"code","source":"class Config:\n def __init__(self, **entries):\n self.__dict__.update(entries)\n self.entries=entries\n\n def print(self):\n print(self.entries)\n \n\ndef default(val: Any, d: Any) -> Any:\n \"\"\"\n Returns `val` if it is not None, otherwise returns the default value `d`.\n :param val: The primary value.\n :param d: The default value to return if `val` is None.\n :return: `val` if it is not None, otherwise `d`.\n \"\"\"\n return val if exists(val) else d\n\n\ndef exists(val: Any) -> bool:\n \"\"\"\n Checks whether a given value is not None.\n :param val: The value to check.\n :return: True if `val` is not None, otherwise False.\n \"\"\"\n return val is not None\n\n\ndef init_weights(m: torch.nn.Module) -> None:\n \"\"\"\n Initializes the weights of a given module if it is an instance of `torch.nn.Linear`. \n Currently, the function does not apply any initialization but has commented-out \n Xavier initialization methods.\n :param m: The module to initialize, expected to be a `torch.nn.Linear` instance.\n :return: None\n \"\"\"\n if m is not None and isinstance(m, nn.Linear):\n pass\n\n\ndef load_config_from_yaml(file_path):\n \"\"\"Load YAML file\"\"\"\n with open(file_path, 'r') as file:\n config = yaml.safe_load(file)\n return Config(**config)\n\n\ndef sep():\n print(\"β\"*100)","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.214247Z","iopub.execute_input":"2025-03-13T23:30:58.214499Z","iopub.status.idle":"2025-03-13T23:30:58.243120Z","shell.execute_reply.started":"2025-03-13T23:30:58.214475Z","shell.execute_reply":"2025-03-13T23:30:58.242101Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Dropout [β](#top) \n\n***\n\nDropout functions.","metadata":{}},{"cell_type":"code","source":"class Dropout(nn.Module):\n \"\"\"\n Implementation of dropout with the ability to share the dropout mask\n along a particular dimension.\n\n If not in training mode, this module computes the identity function.\n \"\"\"\n\n def __init__(self, r: float, batch_dim: Union[int, List[int]]):\n \"\"\"\n Args:\n r:\n Dropout rate\n batch_dim:\n Dimension(s) along which the dropout mask is shared\n \"\"\"\n super(Dropout, self).__init__()\n\n self.r = r\n if type(batch_dim) == int:\n batch_dim = [batch_dim]\n self.batch_dim = batch_dim\n self.dropout = nn.Dropout(self.r)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Args:\n x:\n Tensor to which dropout is applied. Can have any shape\n compatible with self.batch_dim\n \"\"\"\n shape = list(x.shape)\n if self.batch_dim is not None:\n for bd in self.batch_dim:\n shape[bd] = 1\n mask = x.new_ones(shape)\n mask = self.dropout(mask)\n x = x * mask\n return x\n\n\nclass DropoutRowwise(Dropout):\n \"\"\"\n Convenience class for rowwise dropout as described in subsection\n 1.11.6.\n \"\"\"\n\n __init__ = partialmethod(Dropout.__init__, batch_dim=-3)\n\n\nclass DropoutColumnwise(Dropout):\n \"\"\"\n Convenience class for columnwise dropout as described in subsection\n 1.11.6.\n \"\"\"\n\n __init__ = partialmethod(Dropout.__init__, batch_dim=-2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.244599Z","iopub.execute_input":"2025-03-13T23:30:58.244855Z","iopub.status.idle":"2025-03-13T23:30:58.267484Z","shell.execute_reply.started":"2025-03-13T23:30:58.244835Z","shell.execute_reply":"2025-03-13T23:30:58.266869Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Mish [β](#top) \n\n***\n\nMish is a self-regularized non-monotonic activation function proposed by Diganta Misra in 2019. Its purpose is to serve as an alternative to more common activation functions like ReLU, Leaky ReLU, or Swish/SiLU in neural networks.\n\nMish has several properties that can be beneficial in neural networks:\n\n1. It's smooth, unlike ReLU which has a non-differentiable point at 0\n2. It's non-monotonic, which allows for better gradient flow in some contexts\n3. It has a slight regularization effect due to its bounded nature at large negative values\n4. It often provides better performance on various deep learning tasks compared to ReLU\n\n\n Softplus function
\n![\"Image](\"https://pytorch.org/docs/stable/_images/Tanh.png\")
\n Hyperbolic Tangent function
\n![](\"https://i.ibb.co/R7fwdF6/bert3.png\")
","metadata":{}},{"cell_type":"code","source":"class ScaledDotProductAttention(nn.Module):\n '''\n Scaled Dot-Product Attention module, computing attention scores based on query and key similarity.\n '''\n \n def __init__(self, temperature: float, attn_dropout: float = 0.1) -> None:\n \"\"\"\n Initializes the Scaled Dot-Product Attention module.\n \n :param temperature: Scaling factor for the dot product attention scores.\n :param attn_dropout: Dropout rate applied to attention weights.\n \"\"\"\n super().__init__()\n self.temperature: float = temperature\n self.dropout: nn.Dropout = nn.Dropout(attn_dropout)\n\n def forward(\n self, \n q: torch.Tensor, # (B, nhead, L, d_k)\n k: torch.Tensor, # (B, nhead, L, d_k)\n v: torch.Tensor, # (B, nhead, L, d_v)\n mask: torch.Tensor | None = None, # (B, 1, L, L) or None\n attn_mask: torch.Tensor | None = None # (B, 1, L, L) or None\n ) -> tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Forward pass of the Scaled Dot-Product Attention.\n \n :param q: Query tensor of shape (B, nhead, L, d_k), where B is batch size, nhead is the number of attention heads,\n L is the sequence length, and d_k is the key/query dimension.\n :param k: Key tensor of shape (B, nhead, L, d_k).\n :param v: Value tensor of shape (B, nhead, L, d_v), where d_v is the value dimension.\n :param mask: Optional bias mask tensor of shape (B, 1, L, L), used for causal masking or padding.\n :param attn_mask: Optional attention mask tensor of shape (B, 1, L, L), where -1 values indicate positions to mask.\n :return: Tuple containing:\n - output (torch.Tensor): The result of the attention mechanism, shape (B, nhead, L, d_v).\n - attn (torch.Tensor): Attention weights after softmax and dropout, shape (B, nhead, L, L).\n \"\"\"\n \n attn = torch.matmul(q, k.transpose(2, 3)) / self.temperature # (B, nhead, L, L)\n \n if mask is not None:\n attn = attn + mask # Apply bias mask (B, nhead, L, L)\n \n if attn_mask is not None:\n attn = attn.float().masked_fill(attn_mask == -1, float('-1e-9')) # Apply attention mask (B, nhead, L, L)\n \n attn = self.dropout(F.softmax(attn, dim=-1)) # (B, nhead, L, L)\n output = torch.matmul(attn, v) # (B, nhead, L, d_v)\n \n return output, attn","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.308236Z","iopub.execute_input":"2025-03-13T23:30:58.308418Z","iopub.status.idle":"2025-03-13T23:30:58.328594Z","shell.execute_reply.started":"2025-03-13T23:30:58.308401Z","shell.execute_reply":"2025-03-13T23:30:58.327885Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | MultiHead Attention [β](#top) \n\n***\n\nThe same process as before can be repeated many times with different Key, Query, and Value projections, forming what is called a multi-head attention. Each head can focus on different projections of the input embeddings. Multihead attention extends self-attention by applying multiple attention mechanisms (or \"heads\") in parallel. Each head learns different attention patterns, which are then combined to produce a more expressive representation.\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/introduction-to-transformers/multihead_attention.png\")
\n\n### Input shapes \n\nThe inputs `q`, `k`, and `v` (query, key, value) have the following shapes:\n- `q`: [bs, len_q, d_model]\n- `k`: [bs, len_k, d_model]\n- `v`: [bs, len_v, d_model]\n\nWhere:\n- `bs` is the batch size (first dimension)\n- `len_q` is the sequence length of the query\n- `len_k` is the sequence length of the key\n- `len_v` is the sequence length of the value (typically equal to `len_k`)\n- `d_model` is the model's embedding dimension\n\nThe module then projects these inputs into multiple heads:\n- Each head has dimension `d_k` for queries and keys\n- Each head has dimension `d_v` for values\n- There are `n_head` different attention heads\n\nThe attention calculations happen in a shape of `[bs, n_head, len_q, len_k]` and the output has the same dimensionality as the input query: `[bs, len_q, d_model]`.\n\nThis is a standard multi-head attention implementation where vectors are projected into multiple subspaces, attention is calculated separately in each subspace, and then the results are concatenated and projected back to the original dimension.\n","metadata":{}},{"cell_type":"code","source":"class MultiHeadAttention(nn.Module):\n \"\"\"\n Multi-Head Attention module\n :param d_model: The number of input features.\n :param n_head: The number of heads to use.\n :param d_k: The dimensionality of the keys.\n :param d_v: The dimensionality of the values.\n :param dropout: The dropout rate to apply to the attention weights.\n \"\"\"\n def __init__(\n self,\n d_model: int,\n n_head: int,\n d_k: int,\n d_v: int,\n dropout: float = 0.1\n ):\n super().__init__()\n\n self.n_head = n_head\n self.d_k = d_k\n self.d_v = d_v\n\n self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False) # (d_model) -> (n_head * d_k)\n self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False) # (d_model) -> (n_head * d_k)\n self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False) # (d_model) -> (n_head * d_v)\n self.fc = nn.Linear(n_head * d_v, d_model, bias=False) # (n_head * d_v) -> (d_model)\n\n self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)\n\n self.dropout = nn.Dropout(dropout)\n self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)\n\n\n def forward(\n self, \n q: torch.Tensor, # Shape: [batch_size, len_q, d_model]\n k: torch.Tensor, # Shape: [batch_size, len_k, d_model]\n v: torch.Tensor, # Shape: [batch_size, len_v, d_model]\n mask: Optional[torch.Tensor] = None, # Optional attention mask\n src_mask: Optional[torch.Tensor] = None # Optional source mask\n \n ) -> Tuple[torch.Tensor, torch.Tensor]: # Returns (output, attention)\n\n d_k, d_v, n_head = self.d_k, self.d_v, self.n_head\n bs, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)\n\n residual = q # Shape: [bs, len_q, d_model]\n\n # Linear projections and reshape to multiple heads\n q = self.w_qs(q).view(bs, len_q, n_head, d_k) # Shape: [bs, len_q, n_head, d_k]\n k = self.w_ks(k).view(bs, len_k, n_head, d_k) # Shape: [bs, len_k, n_head, d_k]\n v = self.w_vs(v).view(bs, len_v, n_head, d_v) # Shape: [bs, len_v, n_head, d_v]\n\n # Transpose for multi-head attention computation\n q, k, v = (\n q.transpose(1, 2), # Shape: [bs, n_head, len_q, d_k]\n k.transpose(1, 2), # Shape: [bs, n_head, len_k, d_k]\n v.transpose(1, 2)\n ) # Shape: [bs, n_head, len_v, d_v]\n\n if mask is not None:\n mask = mask # Shape remains unchanged\n\n if src_mask is not None:\n src_mask[src_mask == 0] = -1\n src_mask = src_mask.unsqueeze(-1).float() # Shape: [bs, len_k, 1]\n attn_mask = torch.matmul(src_mask, src_mask.permute(0, 2, 1)).unsqueeze(1) \n # Shape: [bs, 1, len_k, len_k]\n q, attn = self.attention(q, k, v, mask=mask, attn_mask=attn_mask)\n else:\n q, attn = self.attention(q, k, v, mask=mask)\n\n # Reshape back to original format\n q = q.transpose(1, 2).contiguous().view(bs, len_q, -1) # Shape: [bs, len_q, n_head * d_v]\n q = self.dropout(self.fc(q)) # Shape: [bs, len_q, d_model]\n q += residual # Shape: [bs, len_q, d_model]\n\n q = self.layer_norm(q) # Shape: [bs, len_q, d_model]\n\n return q, attn # Output: [bs, len_q, d_model], Attention: [bs, n_head, len_q, len_k]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.330676Z","iopub.execute_input":"2025-03-13T23:30:58.330942Z","iopub.status.idle":"2025-03-13T23:30:58.351533Z","shell.execute_reply.started":"2025-03-13T23:30:58.330921Z","shell.execute_reply":"2025-03-13T23:30:58.350837Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Positional Encoding [β](#top) \n\n***\n\nSince self-attention mechanisms do not have inherent order awareness, this encoding helps the model distinguish between different positions in a sequence. Positional embeddings are vectors that contain information about a position in the sequence. This adds information about the sequence even before attention is applied, and it allows attention to calculate relationships knowing the relative order.\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/introduction-to-transformers/positional_encoding.png\")
\n\n \nA detailed explanation of how it works can be found [here](https://kazemnejad.com/blog/transformer_architecture_positional_encoding/), but a quick explanation is that we create a vector for each element representing its position with regard to every other element in the sequence. Positional encoding follows this formula which, in practice, we wonβt really need to understand: \n\n$$\\begin{equation}\n p_{i,j} = \\left\\{\n \\begin{array}{@{}ll@{}}\n \\sin \\left(\\frac{1}{10000^{\\frac{j}{dim\\:embed}}} \\right), & \\text{if}\\ j=even \\\\\n \\cos \\left(\\frac{1}{10000^{\\frac{j}{dim\\:embed}}} \\right), & \\text{if}\\ j=odd \\\\\n \\end{array}\\right.\n\\end{equation} \n$$","metadata":{}},{"cell_type":"code","source":"class PositionalEncoding(nn.Module):\n def __init__(\n self,\n d_model: int,\n dropout: float = 0.1,\n max_len: int = 200\n ):\n super(PositionalEncoding, self).__init__()\n self.dropout = nn.Dropout(p=dropout)\n pe = torch.zeros(max_len, d_model, dtype=torch.float32)\n position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)\n div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model))\n pe[:, 0::2] = torch.sin(position * div_term)\n pe[:, 1::2] = torch.cos(position * div_term)\n pe = pe.unsqueeze(0).transpose(0, 1) # Shape: (max_len, 1, d_model)\n\n self.register_buffer('pe', pe)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n :param x: tensor of shape (seq_len, batch_size, d_model)\n :return: tensor of shape (seq_len, batch_size, d_model) with positional encodings added\n \"\"\"\n x = x + self.pe[:x.size(0), :]\n return self.dropout(x)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.352717Z","iopub.execute_input":"2025-03-13T23:30:58.352950Z","iopub.status.idle":"2025-03-13T23:30:58.373379Z","shell.execute_reply.started":"2025-03-13T23:30:58.352930Z","shell.execute_reply":"2025-03-13T23:30:58.372636Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Outer Product Mean [β](#top) \n\n***\n\nThe `OuterProductMean` class was proposed in the paper [Highly accurate protein structure prediction with AlphaFold][1]. It computes pairwise interactions between elements in a sequence representation. It is designed to capture interactions between all pairs of positions in a sequence by computing their outer product. This is particularly useful in models that need to understand relationships between any two elements in a sequence, such as protein structure prediction models.\n\n> The MSA representation updates the pair representation through an element-wise outer product that is summed over the MSA sequence dimension. In contrast to previous work, this operation is applied within every block rather than once in the network, which enables the continuous communication from the evolving MSA representation to the pair representation.\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/outer_product_mean.png\")
\n\n[1]: https://www.nature.com/articles/s41586-021-03819-2","metadata":{}},{"cell_type":"code","source":"class OuterProductMean(nn.Module):\n \"\"\"\n Outer Product Mean class.\n :param in_dim: Dimensionality of the input sequence representations (default: 256).\n :param dim_msa: Intermediate lower-dimensional representation (default: 32).\n :param pairwise_dim: Final dimensionality of the pairwise output (default: 64).\n \"\"\"\n def __init__(\n self,\n in_dim: int = 256,\n dim_msa: int = 32,\n pairwise_dim: int = 64\n ):\n super(OuterProductMean, self).__init__()\n self.proj_down1 = nn.Linear(in_dim, dim_msa) # projects the input sequence representation into a lower dimensional space\n self.proj_down2 = nn.Linear(dim_msa ** 2, pairwise_dim) # projects the outer product representation (reshaped) to the final pairwise_dim.\n\n def forward(\n self,\n seq_rep: torch.Tensor, # shape: (batch_size, seq_length, in_dim)\n pair_rep: torch.Tensor = None # shape: (batch_size, seq_length, seq_length, pairwise_dim)\n ):\n seq_rep = self.proj_down1(seq_rep) # output shape: (batch_size, seq_length, dim_msa)\n outer_product = torch.einsum('bid,bjc -> bijcd', seq_rep, seq_rep) # output shape: (batch_size, seq_length, seq_length, dim_msa, dim_msa)\n outer_product = rearrange(outer_product, 'b i j c d -> b i j (c d)') # flattens the last two dimensions: (batch_size, seq_length, seq_length, dim_msa ** 2).\n outer_product = self.proj_down2(outer_product) # output shape: (batch_size, seq_length, seq_length, pairwise_dim)\n\n if pair_rep is not None:\n outer_product = outer_product + pair_rep\n\n return outer_product ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.374333Z","iopub.execute_input":"2025-03-13T23:30:58.374623Z","iopub.status.idle":"2025-03-13T23:30:58.392402Z","shell.execute_reply.started":"2025-03-13T23:30:58.374592Z","shell.execute_reply":"2025-03-13T23:30:58.391804Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Triangle Multiplicative Module [β](#top) \n\n***\n\nIn AlphaFold 2, the Triangle Multiplicative Module is a crucial component designed to capture geometric constraints inherent in protein structures. This module operates on pairwise residue representations, ensuring that the predicted distances between residues adhere to the triangle inequality principle, a fundamental property in Euclidean space.\n\nThe module updates the pair representation $z_{ij}$ between residues $i$ and $j$ by considering their relationships with all other residues $k$ in the protein sequence. This is achieved through two symmetric operations:\n\n1. **Outgoing Update:** Aggregates information from all columns corresponding to residue $i$ and $j$, effectively capturing how these residues jointly interact with others.\n\n2. **Incoming Update:** Aggregates information from all rows corresponding to residue $i$ and $j$, focusing on how other residues jointly influence this pair.\n\nThese operations are analogous to nodes in graph theory, where the update of an edge depends on the nodes it connects and their shared neighbors. By incorporating these updates, the module enforces a form of structural consistency, ensuring that the predicted distances between residues are geometrically plausible.\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/triangular_multiplicative_update.png\")
","metadata":{}},{"cell_type":"code","source":"def exists(val):\n return val is not None\n\ndef default(val, d):\n return val if val is not None else d\n\nclass TriangleMultiplicativeModule(nn.Module):\n \"\"\"\n This class is applied to the pairwise residue representations, ensuring that the predicted distances \n between residues adhere to the triangle inequality principle.\n \"\"\"\n def __init__(\n self,\n *,\n dim: int,\n hidden_dim: Optional[int] = None,\n mix: str = 'ingoing'\n ):\n super().__init__()\n assert mix in {'ingoing', 'outgoing'}, 'mix must be either ingoing or outgoing'\n\n hidden_dim = default(hidden_dim, dim)\n self.norm = nn.LayerNorm(dim)\n\n self.left_proj = nn.Linear(dim, hidden_dim)\n self.right_proj = nn.Linear(dim, hidden_dim)\n self.left_gate = nn.Linear(dim, hidden_dim)\n self.right_gate = nn.Linear(dim, hidden_dim)\n self.out_gate = nn.Linear(dim, hidden_dim)\n\n # Initialize all gating to identity\n for gate in (self.left_gate, self.right_gate, self.out_gate):\n nn.init.constant_(gate.weight, 0.)\n nn.init.constant_(gate.bias, 1.)\n\n if mix == 'outgoing':\n self.mix_einsum_eq = '... i k d, ... j k d -> ... i j d'\n elif mix == 'ingoing':\n self.mix_einsum_eq = '... k j d, ... k i d -> ... i j d'\n\n self.to_out_norm = nn.LayerNorm(hidden_dim)\n self.to_out = nn.Linear(hidden_dim, dim)\n\n def forward(\n self,\n x: torch.Tensor, # (batch_size, seq_len, seq_len, dim)\n src_mask: Optional[torch.Tensor] = None # (batch_size, seq_len)\n ) -> torch.Tensor: # Output: (batch_size, seq_len, seq_len, dim)\n if exists(src_mask):\n src_mask = src_mask.unsqueeze(-1).float() # (batch_size, seq_len, 1)\n mask = torch.matmul(src_mask, src_mask.permute(0, 2, 1)) # (batch_size, seq_len, seq_len)\n mask = rearrange(mask, 'b i j -> b i j ()') # (batch_size, seq_len, seq_len, 1)\n\n assert x.shape[1] == x.shape[2], 'feature map must be symmetrical'\n \n x = self.norm(x) # (batch_size, seq_len, seq_len, dim)\n\n left = self.left_proj(x) # (batch_size, seq_len, seq_len, hidden_dim)\n right = self.right_proj(x) # (batch_size, seq_len, seq_len, hidden_dim)\n\n if exists(src_mask):\n left = left * mask # (batch_size, seq_len, seq_len, hidden_dim)\n right = right * mask # (batch_size, seq_len, seq_len, hidden_dim)\n\n left_gate = self.left_gate(x).sigmoid() # (batch_size, seq_len, seq_len, hidden_dim)\n right_gate = self.right_gate(x).sigmoid() # (batch_size, seq_len, seq_len, hidden_dim)\n out_gate = self.out_gate(x).sigmoid() # (batch_size, seq_len, seq_len, hidden_dim)\n\n left = left * left_gate # (batch_size, seq_len, seq_len, hidden_dim)\n right = right * right_gate # (batch_size, seq_len, seq_len, hidden_dim)\n\n out = einsum(self.mix_einsum_eq, left, right) # (batch_size, seq_len, seq_len, hidden_dim)\n\n out = self.to_out_norm(out) # (batch_size, seq_len, seq_len, hidden_dim)\n out = out * out_gate # (batch_size, seq_len, seq_len, hidden_dim)\n return self.to_out(out) # (batch_size, seq_len, seq_len, dim)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.393103Z","iopub.execute_input":"2025-03-13T23:30:58.393317Z","iopub.status.idle":"2025-03-13T23:30:58.423215Z","shell.execute_reply.started":"2025-03-13T23:30:58.393270Z","shell.execute_reply":"2025-03-13T23:30:58.422408Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Triangle Attention [β](#top) \n\n***\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/triangular_self_attention.png\")
\n","metadata":{}},{"cell_type":"code","source":"class TriangleAttention(nn.Module):\n def __init__(\n self,\n in_dim: int = 128,\n dim: int = 32,\n n_heads: int = 4,\n wise: Literal['row', 'col'] = 'row'\n ):\n \"\"\"\n Implements Triangle Attention Mechanism.\n :param in_dim: Input feature dimension.\n :param dim: Dimension of query, key, and value per head.\n :param n_heads: Number of attention heads.\n :param wise: Whether to apply row-wise or column-wise attention.\n \"\"\"\n super(TriangleAttention, self).__init__()\n self.n_heads = n_heads\n self.wise = wise\n self.norm = nn.LayerNorm(in_dim)\n self.to_qkv = nn.Linear(in_dim, dim * 3 * n_heads, bias=False)\n self.linear_for_pair = nn.Linear(in_dim, n_heads, bias=False)\n self.to_gate = nn.Sequential(\n nn.Linear(in_dim, in_dim),\n nn.Sigmoid()\n )\n self.to_out = nn.Linear(n_heads * dim, in_dim)\n\n def forward(self, z: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for TriangleAttention.\n :param z: Input tensor of shape (B, I, J, in_dim).\n :param src_mask: Source mask of shape (B, I, J).\n :return: Output tensor of shape (B, I, J, in_dim).\n \"\"\"\n # Spawn pair mask\n src_mask = src_mask.clone()\n src_mask[src_mask == 0] = -1\n src_mask = src_mask.unsqueeze(-1).float() # (B, I, J, 1)\n attn_mask = torch.matmul(src_mask, src_mask.permute(0, 2, 1)) # (B, I, J, I)\n\n wise = self.wise\n z = self.norm(z) # (B, I, J, in_dim)\n\n # Compute bias and gate\n gate = self.to_gate(z) # [1] (B, I, J, in_dim)\n b = self.linear_for_pair(z) # [5] (B, I, J, n_heads) \n\n # Compute Q, K, V\n q, k, v = torch.chunk(self.to_qkv(z), 3, -1) # [2], [3], [4]: each (B, I, J, n_heads * dim)\n q, k, v = map(lambda x: rearrange(x, 'b i j (h d)->b i j h d', h=self.n_heads), (q, k, v)) \n # Each: (B, I, J, n_heads, dim)\n scale = q.size(-1) ** 0.5 # Scalar\n\n if wise == 'row':\n eq_attn = 'brihd,brjhd->brijh'\n eq_multi = 'brijh,brjhd->brihd'\n b = rearrange(b, 'b i j (r h)->b r i j h', r=1) # (B, 1, I, J, n_heads)\n softmax_dim = 3\n attn_mask = rearrange(attn_mask, 'b i j->b 1 i j 1') # (B, 1, I, J, 1)\n elif wise == 'col':\n eq_attn = 'bilhd,bjlhd->bijlh'\n eq_multi = 'bijlh,bjlhd->bilhd'\n b = rearrange(b, 'b i j (l h)->b i j l h', l=1) # (B, I, J, 1, n_heads)\n softmax_dim = 2\n attn_mask = rearrange(attn_mask, 'b i j->b i j 1 1') # (B, I, J, 1, 1)\n else:\n raise ValueError('wise should be col or row!')\n\n # Compute attention logits\n logits = (torch.einsum(eq_attn, q, k) / scale + b) # [6], [7] (B, I, J, I, n_heads) or (B, I, J, J, n_heads)\n logits = logits.masked_fill(attn_mask == -1, float('-1e-9')) # Apply mask\n\n # Compute attention weights\n attn = logits.softmax(softmax_dim) # [8] (B, I, J, I, n_heads) or (B, I, J, J, n_heads)\n\n # Compute attention output\n out = torch.einsum(eq_multi, attn, v) # [9] (B, I, J, n_heads, dim)\n out = gate * rearrange(out, 'b i j h d-> b i j (h d)') # [10] (B, I, J, in_dim)\n\n # Final projection\n z_ = self.to_out(out) # (B, I, J, in_dim)\n\n return z_\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.423984Z","iopub.execute_input":"2025-03-13T23:30:58.424254Z","iopub.status.idle":"2025-03-13T23:30:58.443011Z","shell.execute_reply.started":"2025-03-13T23:30:58.424232Z","shell.execute_reply":"2025-03-13T23:30:58.442201Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | ConvTransformer Encoder [β](#top) \n\n***\n\n### The Evoformer block \n\nThe `ConvTransformerEncoderLayer` is pretty similar to the `EvoformerBlock` in the AlphaFold2 model. Since they are similar I will discuss the `EvoformerBlock` first. The Evoformer module of the neural network iteratively updates MSA embedding and pair representation, essentially, detecting patterns of interaction between aminoacids. The Evoformer module consists of 48 identical blocks that take MSA embedding and pair representation on input and produce their refined versions as output. In the `RibonanzaNet`, the number of `ConvTransformerEncoderLayer` layers is different, by default it is set to 9.\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/evoformer_block.png\")
\n\n### The ConvTransformerEncoderLayer \n\nThe `ConvTransformerEncoderLayer` (or Evoformer) is composed of the building blocks we have seen before, such as `MultiHeadAttention`, `TriangleMultiplicativeMdule` and `TriangleAttention`. This class is similar to the Transformer Encoder from the RNAdegformer model: \n\n> RNAdegformer involves the Transformer encoder, whose blocks process a one-dimensional representation of the sequence. In prior work, best predictions from RNAdegformer came from supplementing standard Transformer operations with one dimensional convolutional operations, which are effective in capturing information on sequence-local motifs, and biasing the pairwise attention matrix with terms encoding sequence distance as well as the base pair probability (BPP) matrix computed by conventional secondary structure prediction methods like EternaFold.\n\n\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/transformer_encoder_block.png\"){kind=spacer}
\n\nEvoformer block displayed in the diagram, even though it is similar, is not the same as the ConvTransformerEncoderLayer\n![](\"https://kaggle-images.s3.us-west-2.amazonaws.com/ribonanza-3d/ribonanzanet.png\")
\n","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\nclass RibonanzaNet(nn.Module):\n \"\"\"\n A transformer-based neural network for sequence processing, incorporating convolutional transformer encoder layers,\n outer product mean operations, and relative positional encoding.\n \"\"\"\n def __init__(self, config: object):\n \"\"\"\n Initializes the RibonanzaNet model.\n \n :param config: Configuration object containing model hyperparameters.\n - ninp (int): Input embedding dimension.\n - ntoken (int): Vocabulary size for embedding layer.\n - nclass (int): Number of output classes.\n - nhead (int): Number of attention heads.\n - nlayers (int): Number of transformer encoder layers.\n - dropout (float): Dropout probability.\n - pairwise_dimension (int): Dimension of pairwise features.\n - use_triangular_attention (bool): Whether to use triangular attention.\n - use_bpp (bool): Whether to use base-pairing probability features.\n - k (int): Kernel size for convolutions in transformer layers.\n \"\"\"\n super(RibonanzaNet, self).__init__()\n self.config = config\n nhid = config.ninp * 4\n \n self.transformer_encoder = []\n print(f\"Constructing {config.nlayers} ConvTransformerEncoderLayers\")\n for i in range(config.nlayers):\n k = config.k if i != config.nlayers - 1 else 1\n self.transformer_encoder.append(\n ConvTransformerEncoderLayer(\n d_model=config.ninp, nhead=config.nhead,\n dim_feedforward=nhid,\n pairwise_dimension=config.pairwise_dimension,\n use_triangular_attention=config.use_triangular_attention,\n dropout=config.dropout, k=k)\n )\n self.transformer_encoder = nn.ModuleList(self.transformer_encoder)\n \n self.encoder = nn.Embedding(config.ntoken, config.ninp, padding_idx=4)\n self.decoder = nn.Linear(config.ninp, config.nclass)\n \n if config.use_bpp:\n self.mask_dense = nn.Conv2d(2, config.nhead // 4, 1)\n else:\n self.mask_dense = nn.Conv2d(1, config.nhead // 4, 1)\n \n self.OuterProductMean = OuterProductMean(in_dim=config.ninp, pairwise_dim=config.pairwise_dimension)\n self.pos_encoder = RelativePositionalEncoding(config.pairwise_dimension)\n\n def forward(self, src: torch.Tensor, src_mask: torch.Tensor = None, return_aw: bool = False):\n \"\"\"\n Forward pass of the RibonanzaNet model.\n \n :param src: Input tensor of shape (B, L), where B is the batch size and L is the sequence length.\n :param src_mask: Optional mask tensor of shape (B, L, L), used for attention masking.\n :param return_aw: Boolean flag indicating whether to return attention weights.\n :return: Output tensor of shape (B, L, nclass) if return_aw is False, or a tuple (output, attention_weights).\n \"\"\"\n B, L = src.shape # (Batch size, Sequence length)\n src = self.encoder(src).reshape(B, L, -1) # (B, L, ninp)\n \n pairwise_features = self.OuterProductMean(src) # (B, L, L, pairwise_dimension)\n pairwise_features = pairwise_features + self.pos_encoder(src) # (B, L, L, pairwise_dimension)\n \n attention_weights = []\n for i, layer in enumerate(self.transformer_encoder):\n if src_mask is not None:\n if return_aw:\n src, aw = layer(src, pairwise_features, src_mask, return_aw=return_aw)\n attention_weights.append(aw)\n else:\n src, pairwise_features = layer(src, pairwise_features, src_mask, return_aw=return_aw)\n else:\n if return_aw:\n src, aw = layer(src, pairwise_features, return_aw=return_aw)\n attention_weights.append(aw)\n else:\n src, pairwise_features = layer(src, pairwise_features, return_aw=return_aw)\n \n output = self.decoder(src).squeeze(-1) + pairwise_features.mean() * 0 # (B, L, nclass)\n \n if return_aw:\n return output, attention_weights\n else:\n return output\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.486796Z","iopub.execute_input":"2025-03-13T23:30:58.486985Z","iopub.status.idle":"2025-03-13T23:30:58.509752Z","shell.execute_reply.started":"2025-03-13T23:30:58.486967Z","shell.execute_reply":"2025-03-13T23:30:58.508902Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Build Model [β](#top) \n\n***\n\nLet's build the model to check everything is OK and print the output shape.","metadata":{}},{"cell_type":"code","source":"config = load_config_from_yaml(\"config.yaml\")\nmodel = RibonanzaNet(config).cuda()\nx = torch.ones(4, 128).long().cuda()\nmask = torch.ones(4, 128).long().cuda()\nmask[:,120:] = 0\nprint(f\"Output shape: {model(x,src_mask=mask).shape}\"), sep()\nmodel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-13T23:30:58.510689Z","iopub.execute_input":"2025-03-13T23:30:58.510961Z","iopub.status.idle":"2025-03-13T23:31:00.161521Z","shell.execute_reply.started":"2025-03-13T23:30:58.510932Z","shell.execute_reply":"2025-03-13T23:31:00.160671Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# | Conclusions [β](#top) \n\n***\n\nThat's it! We have gone through all the blocks involved in the `RibonanzaNet`. As we have seen, it holds many similarities with some building block of AlphaFold 2, as well as the [RNAdegformer][1] and other top Kaggle solutions from the Stanford Ribonanza RNA Folding competition.\n\nShujun has also provided both a [finetuning][2] code and an [inference][3] code for starters! Go check them out.\n\nI hope this tutorial has been useful to you to better understand the `RibonanzaNet` network used in this competition.\n\nPlease leave any comments below and improvements for the notebook.\n\nBest luck! π\n\n[1]: https://www.kaggle.com/competitions/stanford-ribonanza-rna-folding/discussion/460301\n[2]: https://www.kaggle.com/code/shujun717/ribonanzanet-3d-finetune\n[3]: https://www.kaggle.com/code/shujun717/ribonanzanet-3d-inference","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}