Fixing gradient in sincos positional encoding in monai/networks/blocks/patchembedding.py

[Lucas-rbnt]

Description

When you choose to put the argument pos_embed_type='sincos' in the PatchEmbeddingBlock class, it still return a learnable positional encoding

To reproduce:

from monai.networks.blocks import PatchEmbeddingBlock
patcher = PatchEmbeddingBlock(
            in_channels=1,
            img_size=(32, 32, 32),
            patch_size=(8, 8, 8),
            hidden_size=96,
            num_heads=8,
            pos_embed_type="sincos",
            dropout_rate=0.5,
        )
print(patcher.position_embeddings.requires_grad) 
>>> True

In the literature, we sometimes use either positional encoding in sincos which are fixed and non-trainable as in the original Attention Is All You Need paper or a learnable positional embedding as in the ViT paper.
If you choose to use a sincos, then it seems that is must be fixed which is not the case here.
I'm not completely sure of the desired result in MONAI since there's already a learnable possibility, so if we choose sincos we'd like gradient-free parameters. However the documentation of build_sincos_position_embeddingin the pos_embed_utils.pyfiles stipulate: "The sin-cos position embedding as a learnable parameter" which seems a bit confusing. Especially as the encoding construction function seems to aim to set the require gradient to False (see below)

pos_embed = nn.Parameter(pos_emb)
pos_embed.requires_grad = False

return pos_embed

But these changes are not maintained by torch's copy_ function, which does not copy gradient parameters (see the cpp code https://github.com/pytorch/pytorch/blob/148a8de6397be6e4b4ca1508b03b82d117bfb03c/torch/csrc/lazy/ts_backend/tensor_aten_ops.cpp#L51).
This copy_is used in the PatchEmbeddingBlock class to instantiate the positional embedding.

I propose a small fix to overcome this problem as well as test cases to ensure that positional embedding behaves correctly.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[marksgraham]

LGTM, i agree the sincos embeddings should be fixed, it doesn't make sense to have them learnable.

[KumoLiu]

/build

[KumoLiu]

/build

[KumoLiu]

/build

[KumoLiu]

/build

[KumoLiu]

Merged now, feel free to create another PR for #7564 (comment).