[P1] Intervention Locations more than Prefix and Suffix

[comeandcode]

Does the backbone pyvene support more than two intervention blocks on one layer? I met anaconda3/envs/reft_train/lib/python3.10/site-packages/pyvene/models/intervenable_base.py", line 1092, in _wait_for_forward_with_parallel_intervention unit_locations_base[ IndexError: list index out of range when I tried to add the third intervention block into one layer.

[frankaging]

@comeandcode Hey, what is your intervenable model config? And could you give an example of the unit_locations field you are feeding into the model forward pass? Thanks!

[comeandcode]

Thank you very much for your reply! Consider I am applying loreft to 3 layers: [3, 6, 9], and I create a list: layers = [3, 6, 9, 3, 6, 9, 3, 6, 9], and using a for loop to create representations like:reft_representations=[{ "component": component_mapping["language_model_layers"] % (int(index) - 1), "low_rank_dimension": low_rank_dimension, "intervention": pyreft.LoreftIntervention(embed_dim=model.model.config.hidden_size, low_rank_dimension=low_rank_dimension)} for index in layers].
And for the unit_locations passed in the forward, it's a 3d list like (assume batch_size =2 and the prompt length is the same for these two prompts ) [[[1, 2, 3],[1, 2, 3]], [[1, 2, 3],[1, 2, 3]], [[1, 2, 3],[1, 2, 3]], [[5, 6, 7],[5, 6, 7]], [[5, 6, 7],[5, 6, 7]], [[5, 6, 7],[5, 6, 7]], [[10, 11, 12],[10, 11, 12]], [[10, 11, 12],[10, 11, 12]], [[10, 11, 12],[10, 11, 12]]], so location 1, 2,3 for these 3 layers, location 5, 6, 7 for these 3 layers and location 10, 11, 12 for these 3 layers. (after permute(1, 0, 2).tolist()).

[frankaging]

@comeandcode Thanks. This should be a basic usecase of pyvene APIs.

Did you set up your call like:

        # run intervened forward pass
        unit_locations = None
        if "intervention_locations" in inputs:
            unit_locations={"sources->base": (
                None,
                inputs["intervention_locations"].permute(1, 0, 2).tolist()
            )}
            
        _, cf_outputs = intervenable(
            {
                "input_ids": inputs["input_ids"],
                "attention_mask": inputs["attention_mask"]
            },
            unit_locations=unit_locations,
            labels=inputs["labels"],
            subspaces=inputs["subspaces"].permute(1, 0, 2).tolist() if "subspaces" in inputs else None
        )

The unit_locations field needs to be set up as a dict.

[comeandcode]

Thank you for your reply! Yes, I followed the example you kindly provided in Pyreft/example/loreft, since using multiple intervention locations (3 or more) on the same layer is supported by pyvene, I will try to find out if there are any other mistakes and if so I will post them here! By the way, I found in the example/loreft, special tokens were skipped such as BOS or EOS right? Since I printed out the intervention locations and found they started from 1 but 0. So that only real prompts are intervened instead of affecting special tokens. I am not sure if I am correct on this. Thank you!

[nocoder657]

Hi, did you find out the reason ? Thanks!

[LvAoAo]

Hi, did you find out the reason ? Thanks!

I found that if you use get_intervention_locations function in pyreft when the num_interventions is set as an odd number (like 5 or 7 ...) and share_weights is set as False, the length of intervention_locations you get will one less than the num_interventions you set.

Take the code as below as an example:

intervention_locations = get_intervention_locations(
                    last_position=prompt_length, 
                    first_n=5, 
                    last_n=5,
                    pad_mode="last",
                    num_interventions=5,
                    share_weights=False,
                )

The function will give you

tensor([[[   0,    1,    2,    3,    4],
         [   0,    1,    2,    3,    4],
         [3007, 3008, 3009, 3010, 3011],
         [3007, 3008, 3009, 3010, 3011]]], dtype=torch.int32) # 4 layer

instead of

tensor([[[   0,    1,    2,    3,    4],
         [   0,    1,    2,    3,    4],
         [3007, 3008, 3009, 3010, 3011],
         [3007, 3008, 3009, 3010, 3011],
         [3007, 3008, 3009, 3010, 3011]]], dtype=torch.int32) # 5 layer

or

tensor([[[   0,    1,    2,    3,    4],
         [   0,    1,    2,    3,    4],
         [   0,    1,    2,    3,    4],
         [3007, 3008, 3009, 3010, 3011],
         [3007, 3008, 3009, 3010, 3011]]], dtype=torch.int32) # 5 layer

This is due to the simple logic error in
get_intervention_locations

intervention_locations = [left_intervention_locations]*(num_interventions//2) + \
            [right_intervention_locations]*(num_interventions//2)

If num_interventions is 5, then num_interventions//2 equals 2. This means [left_intervention_locations] and [right_intervention_locations] are each repeated 2 times, resulting in a list with only 4 elements instead of the expected 5.

So you can change the function get_intervention_locations directly through

intervention_locations = [left_intervention_locations]*(num_interventions//2) + \
            [right_intervention_locations]*(num_interventions//2 + 1)

or just avoid setting odd num_interventions.

I struggled with this and found this error. Hope this can help you.