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,6 +337,8 @@ class EvoformerBlock(nn.Module):
             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,
@@ -398,6 +401,7 @@ class EvoformerBlock(nn.Module):
             inplace=inplace_safe,
         )
 
+        if self.no_column_attention == False:
             m = add(m,
                 self.msa_att_col(
                     m,
@@ -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,
             )