How can we get attention weights from example sequence and structure?

[what-is-what]

How can we get attention weights from example sequence and structure?
There were no arguments to get attention weights in transformer blocks, unlike esm2.

[gelnesr]

also interested in this feature, if available!

[ebetica]

Unfortunately, pytorch flash attention doesn't let you do this. You'll have to hack it in, we'll look into support it officially. Here's where the attention is computed, you'll just have use a pytorch implementation of attention to expose the attention matrix.

esm/esm/layers/attention.py

Lines 62 to 75 in 39a3a6c

	if seq_id is not None:
	# Where True, enable participation in attention.
	mask_BLL = seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)
	mask_BHLL = mask_BLL.unsqueeze(1)
	context_BHLD = F.scaled_dot_product_attention(
	query_BHLD, key_BHLD, value_BHLD, mask_BHLL
	)
	else:
	# Shortcut, if we don't use attention biases then torch
	# will autoselect flashattention as the implementation
	context_BHLD = F.scaled_dot_product_attention(
	query_BHLD, key_BHLD, value_BHLD
	)