I’m using register_forward_hook to perform a controlled intervention on a layer’s outputs of a transformer model. For instance:
hook = model.roberta.encoder.layer[args.layer].register_forward_hook(lambda model, inputs, outputs: intervention(outputs=outputs))
The intervention() function performs some operations on the hidden state, and I want to observe changes in model behavior when I perform these changes. However, I would like to make the interventions dynamic - I want to perform it only on certain tokens in the hidden state, whereas I’m doing it to all tokens currently. I have a variable listing which positions need to be intervened separately - how do I access this element within a forward hook if it doesn’t get passed through the forward pass? My current idea is to make a global or static variable that keeps track of what batch I am at currently, and to access these from inside the hook function, but this seems quite hacky. Any suggestions would be appreciated.