Added switching off of column attention in evoformer when using sequence embeddings.

- Added flag `no_column_attention` in evoformer config.
- Added check in `evoformer.py` to switch off `MSAColumnAttention` when the config flag `no_column_attention` is `True`.

openfold/config.py

@@ -514,6 +514,7 @@ config = mlc.ConfigDict(
                 "transition_n": 4,
                 "msa_dropout": 0.15,
                 "pair_dropout": 0.25,
+                "no_column_attention": False,
                 "blocks_per_ckpt": blocks_per_ckpt,
                 "clear_cache_between_blocks": False,
                 "tune_chunk_size": tune_chunk_size,

openfold/model/evoformer.py

@@ -323,6 +323,7 @@ class EvoformerBlock(nn.Module):
         transition_n: int,
         msa_dropout: float,
         pair_dropout: float,
+        no_column_attention: bool,
         inf: float,
         eps: float,
     ):
@@ -336,12 +337,14 @@ class EvoformerBlock(nn.Module):
             inf=inf,
         )
 
-        self.msa_att_col = MSAColumnAttention(
-            c_m,
-            c_hidden_msa_att,
-            no_heads_msa,
-            inf=inf,
-        )
+        self.no_column_attention = no_column_attention
+        if self.no_column_attention == False:
+            self.msa_att_col = MSAColumnAttention(
+                c_m,
+                c_hidden_msa_att,
+                no_heads_msa,
+                inf=inf,
+            )
 
         self.msa_dropout_layer = DropoutRowwise(msa_dropout)
 
@@ -397,17 +400,18 @@ class EvoformerBlock(nn.Module):
             ),
             inplace=inplace_safe,
         )
-        
-        m = add(m, 
-            self.msa_att_col(
-                m, 
-                mask=msa_mask, 
-                chunk_size=chunk_size,
-                use_lma=use_lma,
-                use_flash=use_flash,
-            ),
-            inplace=inplace_safe,
-        )
+
+        if self.no_column_attention == False:
+            m = add(m,
+                self.msa_att_col(
+                    m,
+                    mask=msa_mask,
+                    chunk_size=chunk_size,
+                    use_lma=use_lma,
+                    use_flash=use_flash,
+                ),
+                inplace=inplace_safe,
+            )
 
         if(not inplace_safe):
             input_tensors = [m, input_tensors[1]]
@@ -595,6 +599,7 @@ class EvoformerStack(nn.Module):
         msa_dropout: float,
         pair_dropout: float,
         blocks_per_ckpt: int,
+        no_column_attention: bool,
         inf: float,
         eps: float,
         clear_cache_between_blocks: bool = False, 
@@ -658,6 +663,7 @@ class EvoformerStack(nn.Module):
                 transition_n=transition_n,
                 msa_dropout=msa_dropout,
                 pair_dropout=pair_dropout,
+                no_column_attention=no_column_attention,
                 inf=inf,
                 eps=eps,
             )