[P1] Feeling confused about num_interventions and intervention_locations

[NaOH678]

Hello! I'm interested in the ReFT but have some questions.

I want to intervene the 15th and 16th layer of the model, I use the following code:

TARGET_LAYER = 15

# get reft model
reft_config = ReftConfig(representations=[{
    "layer": TARGET_LAYER, "component": "block_output",
    "intervention": SubloreftIntervention(
    embed_dim=model.config.hidden_size, low_rank_dimension=8)},
    {
    "layer": TARGET_LAYER + 1, "component": "block_output",
    "intervention": SubloreftIntervention(
    embed_dim=model.config.hidden_size, low_rank_dimension=8)
    }])
reft_model = get_reft_model(model, reft_config)
reft_model.print_trainable_parameters()

and use the code to get intervention_locations:

dataset = make_multiple_position_supervised_data_module(tokenizer, model, num_interventions=len(reft_model.interventions),
                                                    inputs=subspace_dataset['instruction'], outputs=subspace_dataset['output'], 
                                                    positions='f5+l5') 

What I want to is to intervene the f5 and l5 tokens at both layer 15 and layer 16，but confused with the following results.
Here are first five results of all data:

dataset['train_dataset']['intervention_locations'][:5]

[
[[0, 1, 2, 3, 4], [22, 23, 24, 25, 26]],
 [[0, 1, 2, 3, 4], [8, 9, 10, 11, 12]],
 [[0, 1, 2, 3, 4], [13, 14, 15, 16, 17]],
 [[0, 1, 2, 3, 4], [17, 18, 19, 20, 21]],
 [[0, 1, 2, 3, 4], [13, 14, 15, 16, 17]]

]

I am confused about the three-dimention list. for each line in the list is a two-dimention list represents the first five tokens and last five tokens, does it mean intervene the first five token at layer 15 and last five token at layer 16 or intervene both first five and last five tokes at both layer 15 and layer 16? what is the meaning of the thee-dimention list?

Thanks for answering me question!!!

[frankaging]

@NaOH678 Thanks for the question! Our unit location list is a bit complex. Here are some annotations for the 3d list.

• First of all, the 3d list from make_multiple_position_supervised_data_module is batched.
• The shape is [batch_size, number of interventions, number of intervening positions]
• In your case, the actual shape is [5, 2, 5]?
• Last two dims, 2*5=10 makes up your f5+l5, 10 intervening positions.

Let me know if these help!

[NaOH678]

@frankaging Thanks for your answer! The core problem is that I don't understand what is the num_interventions meaning. Sometimes the num_interventions , the positions='f5+l5' , the weight_share may conflict with each other resulting intervention_locations to [] or subspaces to []. (I want do some research about the orthogonal subspace.)

Let's back to the initialization I did

TARGET_LAYER = 15

# get reft model
reft_config = ReftConfig(representations=[{
    "layer": TARGET_LAYER, "component": "block_output",
    "intervention": SubloreftIntervention(
    embed_dim=model.config.hidden_size, low_rank_dimension=8)},
    {
    "layer": TARGET_LAYER + 1, "component": "block_output",
    "intervention": SubloreftIntervention(
    embed_dim=model.config.hidden_size, low_rank_dimension=8)
    }])
reft_model = get_reft_model(model, reft_config)

I want to intervene the Layer 15 and Layer 16.

when I use the reft_model.interventions，it returns the intervention layers.

ModuleDict(
  (layer_15_comp_block_output_unit_pos_nunit_1#0): SubloreftIntervention(
    (rotate_layer): ParametrizedLowRankRotateLayer(
      (parametrizations): ModuleDict(
        (weight): ParametrizationList(
          (0): _Orthogonal()
        )
      )
    )
    (learned_source): Linear(in_features=4096, out_features=8, bias=True)
    (dropout): Dropout(p=0.0, inplace=False)
    (act_fn): LinearActivation()
  )
  (layer_16_comp_block_output_unit_pos_nunit_1#0): SubloreftIntervention(
    (rotate_layer): ParametrizedLowRankRotateLayer(
      (parametrizations): ModuleDict(
        (weight): ParametrizationList(
          (0): _Orthogonal()
        )
      )
    )
    (learned_source): Linear(in_features=4096, out_features=8, bias=True)
    (dropout): Dropout(p=0.0, inplace=False)
    (act_fn): LinearActivation()
  )
)

the result of the len(reft_model.interventions) is obviously 2. I have seen the usage in the examples that num_interventions=len(reft_model.interventions)

then I use the code

dataset = make_multiple_position_supervised_data_module(tokenizer, model, num_interventions=len(reft_model.interventions),
                                                    inputs=subspace_dataset['instruction'], outputs=subspace_dataset['output'], 
                                                    positions='f5+l5') 

(weight_share is False in default)
In my case, I got the intervening positions as following:

A = [
[[0, 1, 2, 3, 4], [22, 23, 24, 25, 26]],
 [[0, 1, 2, 3, 4], [8, 9, 10, 11, 12]],
 [[0, 1, 2, 3, 4], [13, 14, 15, 16, 17]],
 [[0, 1, 2, 3, 4], [17, 18, 19, 20, 21]],
 [[0, 1, 2, 3, 4], [13, 14, 15, 16, 17]]

]

In dpo demo dpo.ipynb, related code is

reft_config = pyreft.ReftConfig(representations=[
    {
        "layer": 18,
        "component": "block_output",
        "low_rank_dimension": 2,
        "intervention": pyreft.LoreftIntervention(
            embed_dim=model.config.hidden_size,
            low_rank_dimension=4
        )
    },
    {
        "layer": 28,
        "component": "block_output",
        "low_rank_dimension": 2,
        "intervention": pyreft.LoreftIntervention(
            embed_dim=model.config.hidden_size,
            low_rank_dimension=4
        )
    }
])
reft_model = pyreft.get_reft_model(model, reft_config)

data_module = pyreft.make_multiple_position_supervised_data_module(
    tokenizer, model, prompts, correct_answers,
    positions="f1+l1", share_weights=True, num_interventions=2
)

the invervening location of the first sample is

[[0, 18], [0, 18]]

Let's talk about the first line of A, [[0, 1, 2, 3, 4], [22, 23, 24, 25, 26]]，does it mean intervene the first five token at layer 15, and last five token at layer 16? or It means invervening at both layer 15 and 16 for the 10 positions? What I want to do is latter.

In dpo demo, what does the [[0, 18], [0, 18]] meaning? It seems that intervening at both layer 18 and 28 for the first and last token?
So what is the num_interventions meaning? does it means the number of layers I want to intervene or others?

I don't know how to correctly set the num_interventions , the positions and the weight_share.

For example: If I want to intervene at layer 18, 19, 20. And for each hidden_state of sample at the layer 18, 19 and 20, the intervening positions is 'f5+l5' How can I set it correctly?

Thanks a lot for your support and the amazing work on this project!

[frankaging]

@NaOH678 Sorry for the delay! I saw your email and your pings!

Here are a more detailed description about intervention_locations and num_interventions and others.

• First of all, this library is built on our pyvene library as a wrapper. I would suggest to go through the pyvene tutorial here to get the overall picture.
• Now, interventions are the basic training primitives of ReFT. In ReFT, all we are doing is training these mounted interventions.
• Interventions can be mounted to different layers. In the current pyreft library, we actually assume, interventions cannot share weights across layers.
• For a single layer, we can mount multiple interventions, having them intervening on separate tokens. This is what we mean by f1+l1 and share_weights=False. If it is specified in this way, for one intervening layer, we are intervening on the first and the last token, and we use separate interventions. If share_weights=True, we only have one intervention acting on both the first and the last tokens.
• Now, lets say we have the following configuration:
  1. intervening layers = [10, 20]
  2. position = f5+l5
  3. share_weights=False
     In total, we would have 2 (number of intervening layers) * 2 (because we are not sharing weights among prefix and postfix tokens) = 4 interventions. Thus, num_interventions is set to 4. Note that, among the first 5 tokens, we are using the same intervention.
• Using the same setup above ^, lets now figure out how to set intervention_locations.
  1. First of all, you need to specify intervention_locations for all interventions, which is 4 interventions.
  2. For each intervention, you need to supply batched intervention locations with a shape of [batch size, number of intervening positions]. Note that this potentially needs padding, similar to other batched sequence inputs to the transformer model.
  3. Given this, at the end, you would have a 3d matrix which has a shape of [number of interventions, batch size, number of intervening positions]. Now, depending on what stage you are in in our pipeline, you might want the batch size to be the first dimension. As long as you pad this correctly, you can simply to a torch.permute().

Let me know if these pointers help! Thanks.

[NaOH678]

==@frankaging== Thank you very much! Now I have a deeper understanding of the intervention_locations, but got another issue when training the model. I focused on the code compreft.ipynb，when use the following configuration:

1. intervening layers = [15, 16, 17]
2. position = f3+l3
3. share_weights=False
4. num_interventions=3*2
5. batch_size = 5
   In a batch, the intervention_locations is:

{'sources->base': (None, 
[
    [[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]], 
    [[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]], 
    [[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]], 
    [[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]], 
    [[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]], 
    [[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]]
]
)}

just as you said, the shape of it is [6, 5, 3]. However, I print some internal shape during the training:

class SubloreftIntervention(LoreftIntervention):
    """
    This is a LoReFT that supports subspace interventions!
    """
    def forward(
        self, base, source=None, subspaces=None
    ):
        assert subspaces is not None
        output = []

        # print("============== Subloreft ===============")
        # print("=== Debug SubloreftIntervention.forward ===")
        # print("base shape:", base.shape)
        # print("subspaces:", subspaces)
        # print("len(subspaces):", len(subspaces))
        # print("============== Subloreft ===============")
        
        rotated_base = self.rotate_layer(base)
        print("rotated_base shape:", rotated_base.shape)

        diff = self.act_fn(self.learned_source(base)) - rotated_base
        print("diff shape:", diff.shape)

        batched_subspace = []
        batched_weights = []
        
        for example_i in range(len(subspaces)):
            LHS = (diff[example_i, :, subspaces[example_i]])
            RHS = self.rotate_layer.weight[..., subspaces[example_i]].T
            # print(diff.shape, LHS.shape, RHS.shape, base.shape, subspaces)
            # 
            # torch.Size([5, 2, 8]) torch.Size([2, 4]) torch.Size([4, 4096]) torch.Size([5, 2, 4096]) [[4, 5, 6, 7], [4, 5, 6, 7], [4, 5, 6, 7], [4, 5, 6, 7], [4, 5, 6, 7]]
            # print(f"example {example_i}:")
            # print("  LHS shape:", LHS.shape)
            # print("  RHS shape:", RHS.shape)
            batched_subspace += [LHS]
            batched_weights += [RHS]
        

        batched_subspace = torch.stack(batched_subspace, dim=0)
        batched_weights = torch.stack(batched_weights, dim=0)
        output = base + torch.bmm(batched_subspace, batched_weights)

        return self.dropout(output.to(base.dtype))

I found the both the shape of the rotated_base and diff are [5, 3, 8], I think it should be [5,6,8]. (Note: here the 6 should be f3+l3, not the 2 (number of intervening layers) * 2 (because we are not sharing weights among prefix and postfix tokens) )

Then I dive into the source code of the pyvene, and print some variables (unit_locations_base, self.sorted_keys.index(key), unit_locations_base[self.sorted_keys.index(key)]) in the _wait_for_forward_with_parallel_intervention，

# about line 1730
print("unit_locations_base:",unit_locations_base)
for group_id, keys in self._intervention_group.items():
    for key in keys:
        # skip in case smart jump
        if key in self.activations or \
            isinstance(self.interventions[key], LambdaIntervention) or \
            self.interventions[key].is_source_constant:
            
            print(unit_locations_base[
                        self.sorted_keys.index(key)
                    ])
            print("self.sorted_keys.index(key):",self.sorted_keys.index(key))

            set_handlers = self._intervention_setter(
                [key],
                [
                    unit_locations_base[
                        self.sorted_keys.index(key)
                    ]
                ],
                # assume same group targeting the same subspace
                [
                    subspaces[
                        self.sorted_keys.index(key)
                    ]
                ]

the result of the print is

unit_locations_base: [
[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]], 
[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]], 
[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]], 
[[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]], 
[[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]], 
[[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]]
]

# unit_locations_base[self.sorted_keys.index(key)]
[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]

self.sorted_keys.index(key): 0 (1 or 2 )

but I can only got [[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]] during the training(I observed the print during one training epoch but not found [[61, 62, 63], [154, 155, 156], [44, 45, 46], [45, 46, 47], [176, 177, 178]]), that's why the shape of rotated_base shape and diff are [5,3,8], it just slices the first 3 tokens.

I think it is because the keys is odict_items([(0, ['layer_15_comp_block_output_unit_pos_nunit_1#0']), (1, ['layer_16_comp_block_output_unit_pos_nunit_1#0']), (2, ['layer_17_comp_block_output_unit_pos_nunit_1#0'])])
, so the self.sorted_keys.index(key) is from 0,1,2, it seems we can't get other locations(in this case, we can't get last 3 tokens).
How can I solve the problem?

I also tried the following configuration:
3. share_weights=True
4. num_interventions=3
I kept other config the same as before.

and commented out one line of code in ReftDataset, if I don't do this, list out of range will be raised.

I got the shape of the diff and rotated_base are [5,6,8], it seems I intervened the first 3 and last 3 tokens.

Am I doing this right?

Looking forward to your response!!

Thanks!!!!!!!