prep new ideas

mjc_notes.md

@@ -1871,3 +1871,16 @@ maybe I should just use the huggingface chat template? https://huggingface.co/mi
 OK intervening in minstral is likely to give nans :bug:
 
 OK it happens with some interventions. The negative one? Maybe it's not valid to reverse it?!
+
+# 2023-10-30 07:41:10
+
+ideas:
+  - [x] alredy tried positive vs -1*positive it leads to a nan
+  - [x] aready tried neutral vs positive, like 80%, with good generalization, but I wonder if it's trivial?
+- [x] tried ranking vs mse. mse is still richer
+- [x] tried big vs small model, big seems to help
+- [x] Try a negative vs positve intervention - nope
+- [x] try intervening on every layer = nope
+- [x] flip x0, x1... oh wait with ranking it doesn't know which is which anyway. If I try SL I will need to
+
+I'm out of idea? it does overfit, so maybe only giving it later layers?

notebooks/make_dataset2.py

@@ -77,8 +77,10 @@ model, tokenizer = load_model(cfg.model)
 tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True)
 
 # %%
-# # cache busting for the transformers map and ds steps
-# !rm -rf ~/.cache/huggingface/datasets/generator
+if TEST:
+    # # cache busting for the transformers map and ds steps
+    import shutil
+    shutil.rmtree('~/.cache/huggingface/datasets/generator')
 
 
 # %% [markdown]
@@ -86,54 +88,17 @@ tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=
 
 # %%
 
-            
-def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_token = -1, n_difference = 1, direction_method = 'pca'):
-    """
-    We want one set of interventions per model
-    
-    So we always load a cached version if possible. to make it approx repeatable use the same dataset etc
-    """
-    model_name = cfg.model.replace('/', '-')
-    intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}.pkl'
-    intervention_f.parent.mkdir(exist_ok=True, parents=True)
-    if not intervention_f.exists():        
-        
-        hidden_layers = list(range(cfg.layer_padding, model.config.num_hidden_layers, cfg.layer_stride))
-        
-        dataset_fit = load_preproc_dataset('imdb', tokenizer, N=N_fit_examples, seed=cfg.seed, num_shots=cfg.num_shots, max_length=cfg.max_length, prompt_format=cfg.prompt_format)
-        
-        rep_reading_pipeline = pipeline("rep-reading", model=model, tokenizer=tokenizer)
-        honesty_rep_reader = rep_reading_pipeline.get_directions(
-            dataset_fit['question'], 
-            rep_token=rep_token, 
-            hidden_layers=hidden_layers, 
-            n_difference=n_difference, 
-            train_labels=dataset_fit['label_true'], 
-            direction_method=direction_method,
-            batch_size=batch_size,
-            **tokenizer_args
-        )
-        # and save
-        with open(intervention_f, 'wb') as f:
-            pickle.dump(honesty_rep_reader, f)
-            logger.info(f'Saved interventions to {intervention_f}')
-    else:
-        with open(intervention_f, 'rb') as f:
-            honesty_rep_reader = pickle.load(f)
-        logger.info(f'Loaded interventions from {intervention_f}')
-            
-    return honesty_rep_reader
-
-
 # %%
 
 # N_fit_examples = 20
 N_fit_examples = 30
 rep_token = -1
 
-honesty_rep_reader = create_cache_interventions(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size, rep_token=rep_token)
+honesty_rep_reader1 = create_cache_interventions(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size, rep_token=rep_token)
 
-hidden_layers = sorted(honesty_rep_reader.directions.keys())
+honesty_rep_reader2 = create_cache_interventions(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size, rep_token=rep_token, get_negative=True)
+
+hidden_layers = sorted(honesty_rep_reader1.directions.keys())
 hidden_layers
 
 
@@ -154,7 +119,6 @@ hidden_layers
 
 # %%
 
-
 rep_control_pipeline2 = pipeline(
     "rep-control2", 
     model=model, 
@@ -165,14 +129,10 @@ rep_control_pipeline2
 
 
 # %%
+from src.datasets.intervene import get_activations_from_reader
 
-coeff=4.0
-
-activations = {}
-for layer in hidden_layers:
-    activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer]).to(model.device).half()
-
-assert torch.isfinite(torch.concat(list(activations.values()))).all()
+activations1 = get_activations_from_reader(honesty_rep_reader1, hidden_layers, dtype=model.dtype, device=model.device)
+activations2 = get_activations_from_reader(honesty_rep_reader2, hidden_layers, dtype=model.dtype, device=model.device)
 
 # %%
 
@@ -186,7 +146,7 @@ if TEST:
     input_types = {'single':dataset_train[0], 'list':[dataset_train[i] for i in range(3)], 'generator':iter(dataset_train.select(range(3))), 'dataset':dataset_train.select(range(3)).to_iterable_dataset()}
     for name, ds in input_types.items():
         print(f"==== {name} ====")
-        r = rep_control_pipeline2(ds, activations=activations, batch_size=2)
+        r = rep_control_pipeline2(ds, activations=activations1, batch_size=2)
         if isinstance(r, dict):
             r = [r]
         elif isinstance(r, list):
@@ -204,7 +164,9 @@ if TEST:
 from src.datasets.intervene import test_intervention_quality
 
 if TEST:
-    test_intervention_quality(dataset_train, activations, model, rep_control_pipeline2, batch_size=batch_size)
+    test_intervention_quality(dataset_train, activations1, model, rep_control_pipeline2, batch_size=batch_size)
+    
+    test_intervention_quality(dataset_train, activations2, model, rep_control_pipeline2, batch_size=batch_size)
 
 # %%
 
@@ -258,7 +220,7 @@ if cfg.disable_ds_cache:
     
 from src.datasets.load import ds2df, load_ds, get_ds_name, filter_ds_to_known, qc_ds
 
-    
+activations=[activations1, activations2]
 for ds_name in cfg.datasets:
     
     # load dataset

src/datasets/dm.py

@@ -30,12 +30,14 @@ class imdbHSDataModule(pl.LightningDataModule):
     def __init__(self,
                  ds: Dataset,
                  batch_size: int=32,
-                 x_cols = ['end_hidden_states']
+                 x_cols = ['end_hidden_states'],
+                 skip_layers = 0,
                 ):
         super().__init__()
         self.save_hyperparameters(ignore=["ds"])
         self.ds = ds
         self.x_cols = x_cols
+        self.skip_layers = skip_layers
 
     def setup(self, stage: str):
         h = self.hparams
@@ -60,6 +62,8 @@ class imdbHSDataModule(pl.LightningDataModule):
         # take the diff between layers. Shape batch, layers, hidden_states, inferences
         hs = torch.tensor(self.ds_hs['end_hidden_states'])
         hs = hs.diff(1, axis=1) # this makes it the residual between layers
+        if self.skip_layers:
+            hs = hs[:, self.skip_layers:] # drop the first 10 layers to prevent overfitting?
         self.hs0 = hs[..., 0]
         self.hs1 = hs[..., 1]
         

src/datasets/intervene.py

@@ -10,7 +10,7 @@ import torch
 import pickle
 from src.config import root_folder
 from src.prompts.prompt_loading import load_preproc_dataset
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer, pipeline, Pipeline
 from loguru import logger
 
 Activations = NewType("Activations", Dict[str, torch.Tensor])
@@ -46,70 +46,9 @@ def intervention_meta_fn2(
         return intervene(outputs, activations[layer_name])
     else:
         raise ValueError(f"outputs must be tuple or tensor, got {type(outputs)}")
-    
-    
-
-# def get_magnitude(activations: np.ndarray, labels: np.ndarray) -> Tuple[np.ndarray,np.ndarray]:
-#     """
-#     get center of mass direction and magnitude per layer and head
-    
-#     refactored to from https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/utils.py#L698
-#     to use einops and vector ops instead of for loop
-#     """    
-#     # batch length hidden_dim
-#     # TODO: maybe I should just get COM for last token instead?
-#     true_mass_mean = reduce(activations[labels], 'b l d -> l d', 'mean')
-#     false_mass_mean = reduce(activations[~labels], 'b l d -> l d', 'mean')
-#     direction = true_mass_mean - false_mass_mean
-#     direction = direction / np.linalg.norm(direction, axis=1, keepdims=True) # sq norm per layer
-#     activations = reduce(activations, ' b l d -> l d', 'mean')
-#     proj_vals = activations * direction
-#     proj_val_std = reduce(proj_vals, 'l d -> l', np.std)
-#     return direction, proj_val_std
 
 
-# def get_interventions_dict(activations:np.ndarray, labels: np.ndarray, layer_names: List[str]) -> InterventionDict:
-#     """
-#     Make an intervention dict that works with baukit.TraceDict's edit_output.
-
-#     see https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56
-#     """
-#     direction, proj_val_std = get_magnitude(activations, labels)
-#     out = InterventionDict({l:[] for l in layer_names})
-#     for layer_i, ln in enumerate(layer_names):
-#         out[ln].append((direction[layer_i].squeeze(), proj_val_std[layer_i]))
-#     return out
-
-# def intervention_meta_fn(outputs: torch.Tensor, layer_name:str, interventions: InterventionDict, alpha = 15) -> torch.Tensor:
-#     """see 
-#     - honest_llama: https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/validation/validate_2fold.py#L114
-#     - baukit: https://github.com/davidbau/baukit/blob/main/baukit/nethook.py#L42C1-L45C56
-    
-#     Usage:
-#         intervention_fn = partial(intervention_meta_fn, interventions=interventions)
-#         with TraceDict(model, layers_to_intervene, edit_output=intervention_fn) as ret:
-#             ...
-    
-#     """
-#     if type(outputs) is tuple:
-#         # head_output
-#         output = outputs[0]
-#     elif type(outputs) is torch.Tensor:
-#         output = outputs
-#     else:
-#         raise ValueError(f"outputs must be tuple or tensor, got {type(outputs)}")
-        
-#     for direction, proj_val_std in interventions[layer_name]:
-#         # head_output: (batch_size, seq_len, layer_size)
-#         output[:, :, :] += torch.from_numpy(alpha * proj_val_std * direction).to(output.device)[None, None, :]
-#     if type(outputs) is tuple:
-#         return tuple([output, *outputs[1:]])
-#     else:
-#         return output
-
-
-
-def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_token = -1, n_difference = 1, direction_method = 'pca'):
+def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_token = -1, n_difference = 1, direction_method = 'pca', get_negative=False):
     """
     We want one set of interventions per model
     
@@ -118,7 +57,7 @@ def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_s
     tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True)
     
     model_name = cfg.model.replace('/', '-')
-    intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}.pkl'
+    intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}_{"-" if get_negative else "+"}.pkl'
     intervention_f.parent.mkdir(exist_ok=True, parents=True)
     if not intervention_f.exists():        
         
@@ -126,6 +65,9 @@ def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_s
         
         dataset_fit = load_preproc_dataset('imdb', tokenizer, N=N_fit_examples, seed=cfg.seed, num_shots=cfg.num_shots, max_length=cfg.max_length, prompt_format=cfg.prompt_format)
         
+        train_labels = np.array(dataset_fit['label_true'])
+        if get_negative:
+            train_labels = -1 * train_labels
         rep_reading_pipeline = pipeline("rep-reading", model=model, tokenizer=tokenizer)
         honesty_rep_reader = rep_reading_pipeline.get_directions(
             dataset_fit['question'], 
@@ -173,6 +115,7 @@ def intervention_metrics(control_outputs_neg, baseline_outputs, control_outputs)
         print()
 
 def test_intervention_quality(dataset_train, activations, model, rep_control_pipeline2, batch_size=2):
+    # TODO: this have bugs and is not used yet
     inputs = dataset_train[:3]
     activations_neg = {k:-v for k, v in activations.items()}
     activations_none = {k:v*0 for k, v in activations.items()}
@@ -184,3 +127,20 @@ def test_intervention_quality(dataset_train, activations, model, rep_control_pip
 
 
     intervention_metrics(control_outputs_neg, baseline_outputs, control_outputs)
+
+
+def get_activations_from_reader(honesty_rep_reader: Pipeline, hidden_layers: list, coeff=1, dtype=None, device=None) -> Dict[str, float]:
+    """Get activations from the honesty_rep_reader"""
+    
+    # FIXME: coeff is a magic number. The representation_engineering repo used 8, but it seems to vary by model?
+
+    activations = {}
+    for layer in hidden_layers:
+        activations[layer] = torch.tensor(coeff * honesty_rep_reader.directions[layer] * honesty_rep_reader.direction_signs[layer])
+        if device:
+            activations[layer] = activations[layer].to(device)
+        if dtype:
+            activations[layer] = activations[layer].to(dtype)
+
+    assert torch.isfinite(torch.concat(list(activations.values()))).all()
+    return activations

src/extraction/config.py

@@ -33,10 +33,10 @@ class ExtractConfig(Serializable):
     """The number of prompt templates to use for each example. If -1, all available
     templates are used."""
 
-    layer_stride: InitVar[int] = 2
+    layer_stride: InitVar[int] = 1
     """Shortcut for `layers = (0,) + tuple(range(1, num_layers + 1, stride))`."""
     
-    layer_padding: InitVar[int] = 4
+    layer_padding: InitVar[int] = 6
     """Skips this amount of first layers"""
 
     seed: int = 42

src/probes/pl_ranking.py

@@ -158,9 +158,12 @@ class PLConvProbeLinear(PLRankingBase):
             else: # last layer
                 layers.append(nn.Conv1d(hs*4, 1, 1))
         self.conv = nn.Sequential(*layers)
+        
+        n = c_in[0]
         self.head = nn.Sequential(
-            LinBnDrop(c_in[0], 1),
-            nn.Linear(1, 1),            
+            LinBnDrop(n, n),
+            LinBnDrop(n, n),
+            nn.Linear(n, 1),            
         )
         
     def forward(self, x):

src/repe/rep_control_pipeline_baukit.py

@@ -50,7 +50,7 @@ class RepControlPipeline2(FeatureExtractionPipeline):
     def __call__(self, model_inputs, **kwargs):
         return super().__call__(model_inputs, **kwargs)
     
-    def _sanitize_parameters(self, truncation=None, tokenize_kwargs=None, return_tensors=None, activations=None, **kwargs):
+    def _sanitize_parameters(self, activations=None, truncation=None, tokenize_kwargs=None, return_tensors=None, **kwargs):
         """This processed the init params."""
         if tokenize_kwargs is None:
             tokenize_kwargs = {}
@@ -83,16 +83,14 @@ class RepControlPipeline2(FeatureExtractionPipeline):
         inputs["attention_mask"] = torch.tensor(inputs['attention_mask'], dtype=torch.bool, device=self.model.device)
         return inputs
 
-    def _forward(self, inputs, activations) -> ModelOutput:
+    def _forward(self, inputs: dict, activations: List[Dict[str, float]]) -> ModelOutput:
         assert inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {inputs['input_ids'].shape}"
 
         # make intervention functions
-        layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()]                
-        activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()})
-        activations_neg_i = Activations({self.layer_name_tmpl.format(k):-1. * v for k,v in activations.items()})
-        activations_neut = Activations({self.layer_name_tmpl.format(k):0. * v for k,v in activations.items()})
+        layers_names = [self.layer_name_tmpl.format(i) for i in activations[0].keys()]                
+        activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations[1].items()})
+        activations_neut = Activations({self.layer_name_tmpl.format(k):0. * v for k,v in activations[0].items()})
         edit_fn_pos = partial(intervention_meta_fn2, activations=activations_pos_i)
-        edit_fn_neg = partial(intervention_meta_fn2, activations=activations_neg_i)
         edit_fn_neu = partial(intervention_meta_fn2, activations=activations_neut)
         
         self.model.eval()