Add toggle to disable perturbation encoder in state transition model

src/state/configs/model/state.yaml

@@ -21,6 +21,7 @@ kwargs:
   nb_decoder: False
   mask_attn: False
   use_effect_gating_token: False
+  zero_perturbation_encoder: False
   distributional_loss: energy
   init_from: null
   transformer_backbone_key: llama

src/state/configs/model/state_lg.yaml

@@ -21,6 +21,7 @@ kwargs:
   nb_decoder: False
   mask_attn: False
   use_effect_gating_token: False
+  zero_perturbation_encoder: False
   use_basal_projection: False
   distributional_loss: energy
   init_from: null

src/state/configs/model/state_sm.yaml

@@ -20,6 +20,7 @@ kwargs:
   nb_decoder: False
   mask_attn: False
   use_effect_gating_token: False
+  zero_perturbation_encoder: False
   use_basal_projection: False
   distributional_loss: energy
   gene_decoder_bool: False

src/state/tx/models/state_transition.py

@@ -182,6 +182,7 @@ def __init__(
             raise ValueError(f"Unknown loss function: {loss_name}")
 
         self.use_basal_projection = kwargs.get("use_basal_projection", True)
+        self.zero_perturbation_encoder = kwargs.get("zero_perturbation_encoder", False)
 
         # Build the underlying neural OT network
         self._build_networks(lora_cfg=kwargs.get("lora", None))
@@ -392,6 +393,8 @@ def forward(self, batch: dict, padded=True) -> torch.Tensor:
 
         # Shape: [B, S, input_dim]
         pert_embedding = self.encode_perturbation(pert)
+        if self.zero_perturbation_encoder:
+            pert_embedding = torch.zeros_like(pert_embedding)
         control_cells = self.encode_basal_expression(basal)
 
         # Add encodings in input_dim space, then project to hidden_dim