Merge pull request #102 from stanfordnlp/zen/binarymask

[Minor] Add in trainable intervention based on binary mask intervention

pyvene/__init__.py

@@ -26,6 +26,7 @@
 from .models.interventions import DistributedRepresentationIntervention
 from .models.interventions import SourcelessIntervention
 from .models.interventions import NoiseIntervention
+from .models.interventions import SigmoidMaskIntervention
 
 
 # Utils

pyvene/models/intervenable_base.py

@@ -309,7 +309,8 @@ def set_temperature(self, temp: torch.Tensor):
         Set temperature if needed
         """
         for k, v in self.interventions.items():
-            if isinstance(v[0], BoundlessRotatedSpaceIntervention):
+            if isinstance(v[0], BoundlessRotatedSpaceIntervention) or \
+                isinstance(v[0], SigmoidMaskIntervention):
                 v[0].set_temperature(temp)
 
     def enable_model_gradients(self):

pyvene/models/interventions.py

@@ -362,7 +362,7 @@ def __init__(self, **kwargs):
         self.rotate_layer = torch.nn.utils.parametrizations.orthogonal(rotate_layer)
         # boundary masks are initialized to close to 1
         self.masks = torch.nn.Parameter(
-            torch.tensor([100] * self.embed_dim), requires_grad=True
+            torch.tensor([100.0] * self.embed_dim), requires_grad=True
         )
         self.temperature = torch.nn.Parameter(torch.tensor(50.0))
 
@@ -398,6 +398,39 @@ def __str__(self):
         return f"SigmoidMaskRotatedSpaceIntervention()"
 
 
+class SigmoidMaskIntervention(TrainableIntervention, LocalistRepresentationIntervention):
+
+    """Intervention in the original basis with binary mask."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.mask = torch.nn.Parameter(
+            torch.zeros(self.embed_dim), requires_grad=True)
+
+        self.temperature = torch.nn.Parameter(torch.tensor(50.0))
+
+    def get_temperature(self):
+        return self.temperature
+
+    def set_temperature(self, temp: torch.Tensor):
+        self.temperature.data = temp
+
+    def forward(self, base, source, subspaces=None):
+        batch_size = base.shape[0]
+        # get boundary mask between 0 and 1 from sigmoid
+        mask_sigmoid = torch.sigmoid(self.mask / torch.tensor(self.temperature)) 
+
+        # interchange
+        intervened_output = (
+            1.0 - mask_sigmoid
+        ) * base + mask_sigmoid * source
+
+        return intervened_output
+
+    def __str__(self):
+        return f"SigmoidMaskIntervention()"
+
+
 class LowRankRotatedSpaceIntervention(TrainableIntervention, DistributedRepresentationIntervention):
 
     """Intervention in the rotated space."""