Visualising attention

[ilyalasy]

Hey, I was wondering why ParlAI has no framework for outputting attention weights from models for later visualisation?
I think this is pretty important for interpretability of large language models. Or I have missed something?

For example, Huggingface has output_attentions parameter which can help to use such libraries as bertviz, lit, etc.

[klshuster]

Both our BasicAttention and MultiheadAttention weights can (and do) output the attention weights from their forward pass (for MultiheadAttention, these are pre-softmax), though they are indeed ignored in the higher levels of the model (at least for our transformer/generator models)

[ilyalasy]

Yeah, I looked through the code and the output is non retrievable, I am quickly writing some clunky code for retrieving attention weights and tokens, maybe will post it here if interested

Are you interested in adding this feature later, maybe we should create an issue?

[stephenroller]

If you look at the distillation code, we use hooks to actually pull out the attention weights. You could use the same trick elsewhere. This code could be simplified but you still get the idea:

ParlAI/projects/anti_scaling/distillation.py

Lines 256 to 272 in f773ef7

	def _register_series_of_hooks(
	self, model: nn.Module, module_map: Dict[str, Type[nn.Module]]
	) -> Dict[str, OutputRecorder]:
	"""
	Register hooks in modules of the model, given the mapping of module types.
	`module_map` is a dict whose keys are module-type names and whose values are
	module types. For each module type, during each forward pass of `model`, all
	outputs of modules of that type will be saved to `hooks[module_type].outputs`.
	"""
	hooks = {}
	for module_name, module_type in module_map.items():
	hooks[module_name] = OutputRecorder()
	for module in model.modules():
	if isinstance(module, module_type):
	module.register_forward_hook(hooks[module_name])
	return hooks

[ilyalasy]

This is very helpful, thank you!