Feature/acces single heads

[llinauer]

• Add accessing single-head activations in source_forward_pass and target_forward_pass methods of Patchscope
  Warning This currently only works for models with the GPT2Attention implementation (https://github.com/huggingface/transformers/blob/v4.39.2/src/transformers/models/gpt2/modeling_gpt2.py#L123)
  We probably need one implementation per attention architecture

• Modify activation_patching_ioi.py to create a plot by layer & head

[tvhong]

Does this work with other GPT models (e.g., GPTJ)

[llinauer]

Unfortunately not.
GPT-J, despite being similar, uses a different attention implementation (GPTJAttention: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gptj/modeling_gptj.py#L100)

[llinauer]

We need to implement a mechanism that works for a range of model architectures

[tvhong]

Yeah, good news is that we're starting to see a pattern emerge.

I'm thinking we want to have a base ModelAccessor class that looks like:

class ModelAccessor(ABC):
  def get_block_output(position: list[int], layer: int) -> Tensor:
    raise NotImplementedError(...)

  def set_block_output(position: list[int], layer: int) -> None:
    raise NotImplementedError(...)

  def get_head_attn(position: list[int], layer: int, head: list[int]) -> Tensor:
    raise NotImplementedError(...)

  def set_head_attn(position: list[int], layer: int, head: list[int]) -> None:
    raise NotImplementedError(...)

and each model can implement this class.

[tvhong]

Why are we using input instead of output?

My understanding is that patchscope always uses output, and if a researcher needs an input from layer i, they can access the output from layer i-1.

[llinauer]

The problem is, that the output of the c_attn layer in GPT2Attention is not the same as the input of the c_proj.
c_attn.output gets us the Q,K & Values concatenated together into one tensor. We want the attention layer outputs (sometimes referred to as z-values), which are calculated inbetween the c_attn and c_proj forward calls in the GPT2Attention object. So they are input of c_proj, but not output of c_attn
See GPT2Attention.forward for reference (https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L306)

[tvhong]

I.. see. So, c_proj is the equivalence of W^O in the original transformer paper?

If so, I agree that the concatenated head would be at .attn.c_proj.input.

https://arxiv.org/pdf/1706.03762.pdf

[tvhong]

I think this PR is logically sound.

But do you want to merge #41 first to fix CI before merging in this PR?

[tvhong]

You know what, let's merge this. I'll rebase the other one.