Allow an arbitrary mask to be used in the self attention #8235
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
Signed-off-by: Lucas Robinet <[email protected]>
Signed-off-by: Lucas Robinet <[email protected]>
Signed-off-by: Lucas Robinet <[email protected]>
ericspod
reviewed
Nov 22, 2024
|
I think this is fine with the minor proposed change. |
KumoLiu
approved these changes
Nov 22, 2024
Co-authored-by: Eric Kerfoot <[email protected]> Signed-off-by: Lucas Robinet <[email protected]>
Signed-off-by: Lucas Robinet <[email protected]>
|
/build |
|
It seems there is a TorchScript conversion issue caused by this addition. |
It seems to be due to a typing error on the
This typing method is reserved for python versions >3.10, but it seems that python 3.9 is being used in the test environment.
I used this notation because I thought I'd already seen it in MONAI. from typing import Optional
...
attn_mask: Optional[torch.tensor] = NoneI can change, as you prefer! |
Yes, could you help convert this to the older typing syntax, as TorchScript does not support the |
Signed-off-by: Lucas Robinet <[email protected]>
auto-merge was automatically disabled
November 25, 2024 13:10
Head branch was pushed to by a user without write access
Signed-off-by: Lucas Robinet <[email protected]>
|
/build |
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.
Description
The aim of this PR is to enable the use of an arbitrary mask in the self attention module, which is very useful in the case of missing data or masked modeling.
Official torch implementations allow the use of an arbitrary mask, and in MONAI the use of a mask is also made possible with the
causalargument. Here, it's just a generalization directly in the forward pass.In the
SABlockandTransformerBlock, it is now possible to input a boolean mask of size(BS, Seq_length).Only the columns of the masked token are set to
-infand not the rows, as is rarely the case in common implementations. Masked tokens don't contribute to the gradient anyway.In cases where causal attention is required, inputting a mask is not supported to avoid masks overlapping.
I haven't implemented the addition mask to the attention matrix, which allows you to use values other than
-infin certain cases, as may be the case here: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.htmlIf you think it's relevant, it could be added.
Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.