Add return_attention_scores support to CachedMultiHeadAttention #2213
+16
−4
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Call the parent class constructor | ||
| super().__init__(num_heads, key_dim, **kwargs) | ||
| # New flag to optionally return attention scores | ||
| self._return_attention_scores = return_attention_scores |
There was a problem hiding this comment.
this is a call arg in the super class here - https://github.com/keras-team/keras/blob/44a655bdb28037046ab279a49d4cd679fea7ca50/keras/src/layers/attention/multi_head_attention.py#L523
Also if flash attention is used using ops.dot_product_attention then attention scores will not be returned
There was a problem hiding this comment.
this is a call arg in the super class here
makes sense now — instead of manually setting the flag, it made sense to just pass return_attention_scores into super().init() since the base MHA layer handles it internally.
I’ve pushed the fix with that change; let me know if further changes are needed
… updating call logic
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR addresses #2055, where attention_scores was always None due to the _return_attention_scores flag not being set in the CachedMultiHeadAttention subclass.
In recent Keras versions, the base MultiHeadAttention layer uses a private flag self._return_attention_scores to decide whether or not to return attention scores from _compute_attention.
However, CachedMultiHeadAttention was not passing or setting this flag at all, which meant attention_scores were silently dropped — making them inaccessible for debugging or analysis.
In this PR we did the following-
1.Adds return_attention_scores as an optional argument to the constructor (default False, just like in base MHA).
2.Sets self._return_attention_scores appropriately.
3.Updates the call() method to return attention_scores alongside attention_output and cache when requested — fully preserving existing behavior otherwise.