Upload 11 files

.gitattributes

@@ -53,3 +53,9 @@ ioi_submission/ioi_task_Gemma2ForCausalLM_output_token/AttentionHead(Layer-7,Hea
 ioi_submission/ioi_task_Gemma2ForCausalLM_output_token/AttentionHead(Layer-7,Head-6,Token-all)_inverse_featurizer filter=lfs diff=lfs merge=lfs -text
 ioi_submission/ioi_task_Gemma2ForCausalLM_output_token/AttentionHead(Layer-8,Head-1,Token-all)_featurizer filter=lfs diff=lfs merge=lfs -text
 ioi_submission/ioi_task_Gemma2ForCausalLM_output_token/AttentionHead(Layer-8,Head-1,Token-all)_inverse_featurizer filter=lfs diff=lfs merge=lfs -text
+4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol_period)_featurizer filter=lfs diff=lfs merge=lfs -text
+4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol_period)_inverse_featurizer filter=lfs diff=lfs merge=lfs -text
+4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol)_featurizer filter=lfs diff=lfs merge=lfs -text
+4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol)_inverse_featurizer filter=lfs diff=lfs merge=lfs -text
+4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-last_token)_featurizer filter=lfs diff=lfs merge=lfs -text
+4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-last_token)_inverse_featurizer filter=lfs diff=lfs merge=lfs -text

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol)_featurizer

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dab9bccd2ea775eb56ad98fa9bac02c8d6a170d41d3a58b4b300c7e97eb80af8
+size 21531300

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol)_indices

@@ -0,0 +1 @@
+null

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol)_inverse_featurizer

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a5f69e8af1b271494715d6d2cf3936a9f1897065b5cd7a1e35417c0eb19a665
+size 21531356

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol_period)_featurizer

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d3ca3b99e9badc80119a4d711f60f35caf610ae7a8bcf08689385b490a197c0
+size 21531349

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol_period)_indices

@@ -0,0 +1 @@
+null

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-correct_symbol_period)_inverse_featurizer

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:630c183b185a53d826f9e17f6932dfa7e7d1011d8fb8435bc29b42fb1ac45189
+size 21531533

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-last_token)_featurizer

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8483d8ff87d3a188f542bcf17e545d63bb2039a644249982d82ef8c45a65964e
+size 21531208

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-last_token)_indices

@@ -0,0 +1 @@
+null

4_answer_MCQA_Gemma2ForCausalLM_answer_pointer/ResidualStream(Layer-0,Token-last_token)_inverse_featurizer

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f02db1835842cbffc192c82973dc8d08dcc9b2f5667f57ac8f11c7af32c684b8
+size 21531328

featurizer.py

@@ -0,0 +1,52 @@
+"""
+Copy of the existing SubspaceFeaturizer implementation for submission.
+This file provides the same SubspaceFeaturizer functionality in a self-contained format.
+"""
+
+import torch
+import torch.nn as nn
+import pyvene as pv
+from CausalAbstraction.neural.featurizers import Featurizer
+
+
+class SubspaceFeaturizerModuleCopy(torch.nn.Module):
+    def __init__(self, rotate_layer):
+        super().__init__()
+        self.rotate = rotate_layer
+        
+    def forward(self, x):
+        r = self.rotate.weight.T
+        f = x.to(r.dtype) @ r.T
+        error = x - (f @ r).to(x.dtype)
+        return f, error
+
+
+class SubspaceInverseFeaturizerModuleCopy(torch.nn.Module):
+    def __init__(self, rotate_layer):
+        super().__init__()
+        self.rotate = rotate_layer
+        
+    def forward(self, f, error):
+        r = self.rotate.weight.T
+        return (f.to(r.dtype) @ r).to(f.dtype) + error.to(f.dtype)
+
+
+class SubspaceFeaturizerCopy(Featurizer):
+    def __init__(self, shape=None, rotation_subspace=None, trainable=True, id="subspace"):
+        assert shape is not None or rotation_subspace is not None, "Either shape or rotation_subspace must be provided."
+        if shape is not None:
+            self.rotate = pv.models.layers.LowRankRotateLayer(*shape, init_orth=True)
+        elif rotation_subspace is not None:
+            shape = rotation_subspace.shape
+            self.rotate = pv.models.layers.LowRankRotateLayer(*shape, init_orth=False)
+            self.rotate.weight.data.copy_(rotation_subspace)
+        self.rotate = torch.nn.utils.parametrizations.orthogonal(self.rotate)
+
+        if not trainable:
+            self.rotate.requires_grad_(False)
+
+        # Create module-based featurizer and inverse_featurizer
+        featurizer = SubspaceFeaturizerModuleCopy(self.rotate)
+        inverse_featurizer = SubspaceInverseFeaturizerModuleCopy(self.rotate)
+            
+        super().__init__(featurizer, inverse_featurizer, n_features=self.rotate.weight.shape[1], id=id)

token_position.py

@@ -0,0 +1,65 @@
+"""
+Token position definitions for MCQA task submission.
+This file provides token position functions that identify key tokens in MCQA prompts.
+"""
+
+import re
+from CausalAbstraction.neural.LM_units import TokenPosition, get_last_token_index
+
+
+def get_token_positions(pipeline, causal_model):
+    """
+    Get token positions for the simple MCQA task.
+    
+    Args:
+        pipeline: The language model pipeline with tokenizer
+        causal_model: The causal model for the task
+        
+    Returns:
+        list[TokenPosition]: List of TokenPosition objects for intervention experiments
+    """
+    def get_correct_symbol_index(input, pipeline, causal_model):
+        """
+        Find the index of the correct answer symbol in the prompt.
+        
+        Args:
+            input (Dict): The input dictionary to a causal model
+            pipeline: The tokenizer pipeline
+            causal_model: The causal model
+            
+        Returns:
+            list[int]: List containing the index of the correct answer symbol token
+        """
+        # Run the model to get the answer position
+        output = causal_model.run_forward(input)
+        pointer = output["answer_pointer"]
+        correct_symbol = output[f"symbol{pointer}"]
+        prompt = input["raw_input"]
+        
+        # Find all single uppercase letters in the prompt
+        matches = list(re.finditer(r"\b[A-Z]\b", prompt))
+        
+        # Find the match corresponding to our correct symbol
+        symbol_match = None
+        for match in matches:
+            if prompt[match.start():match.end()] == correct_symbol:
+                symbol_match = match
+                break
+                
+        if not symbol_match:
+            raise ValueError(f"Could not find correct symbol {correct_symbol} in prompt: {prompt}")
+        
+        # Get the substring up to the symbol match end
+        substring = prompt[:symbol_match.end()]
+        tokenized_substring = list(pipeline.load(substring)["input_ids"][0])
+        
+        # The symbol token will be at the end of the substring
+        return [len(tokenized_substring) - 1]
+
+    # Create TokenPosition objects
+    token_positions = [
+        TokenPosition(lambda x: get_correct_symbol_index(x, pipeline, causal_model), pipeline, id="correct_symbol"),
+        TokenPosition(lambda x: [get_correct_symbol_index(x, pipeline, causal_model)[0]+1], pipeline, id="correct_symbol_period"),
+        TokenPosition(lambda x: get_last_token_index(x, pipeline), pipeline, id="last_token")
+    ]
+    return token_positions