Hook positions

[kknakka2026-cmd]

x_l
 │
 ├─ Attention ──> attn_out ──┐
 │                           +
 └───────────────────────────┘ → residual
                              │
                              ├─ MLP ──> mlp_out ──┐
                              │                    +
                              └────────────────────┘ → x_{l+1}

HI Team,

Thank you for opensourcing the codebase. May I ask, if the hook positions inside the library are already exposed to different locations inside the transformer block.

[emiehling]

Hi @kknakka2026-cmd

Thanks for the question. The answer is (mostly) yes. The InterventionControl declare an Intervention with a site and a boundary and those cover four points in your diagram:

x_l ─────────────────────── site="decoder_layer", boundary="layer_input"  (pre-hook on the layer)
 │                           site="norm_input" also sees x_l, as the input to input_layernorm / ln_1
 ├─ Attention ──> attn_out   site="o_proj": pre-hook on self_attn.o_proj (attn.c_proj on GPT-2), i.e. the
 │                           per-head concat before the output projection; what ITI's head_additive uses
 +──> residual (mid)         site="norm_input": pre-hook on post_attention_layernorm /
 │                           pre_feedforward_layernorm / ln_2 
 ├─ MLP ──> mlp_out          not an IR site
 +──> x_{l+1} ────────────── site="decoder_layer", boundary="layer_output"  (forward hook; the default)

One note though is that norm_input hooks the input of each sublayer's pre-norm and returns a new tensor, so an edit there changes what the attention or MLP branch sees (but not the tensor carried along the skip connection). We don't currently have native functionality to edit the mid-block residual itself for the change to also travel down the skip path if that's what your looking for.

It seems reasonable to discuss a design to make this a bit easier. I'll bring it up with the team.