Unused cls_token in PatchEmbeddingBlock

[night-gale]

Describe the bug
When I was training the ViT with torch DistributedDataParallel, during backward, torch raises error and reports that

Parameters which did not receive grad for rank 0: vit.patch_embedding.cls_token

which means that the cls_token did not participate in the backward process.

I checked the implementation of ViT and PatchEmbeddingBlock and found the unused cls_token in monai.networks.blocks.patchembedding.py: PatchEmbeddingBlock.

To Reproduce
Steps to reproduce the behavior:

1. set environment variable in shell TORCH_DISTRIBUTED_DEBUG=INFO
2. train ViT with torch DistributedDataParallel

[Nic-Ma]

Thanks for raising the issue.
Hi @ahatamiz ,

Could you please help double confirm the issue?
If we really don't need the cls_token, please remove it.

Thanks in advance.

[ahatamiz]

Hi @night-gale

Thanks for your comments. If you utilize ViT for classification application only, then the classification flag needs to be activated. Doing so will enable the use of cls_token as shown here.

Originally, ViT is used as segmentation backbone for UNETR, hence the application needs to be specificed.

Lastly, cls_token plays an important role in the ViT for classification as it assigns the class type. Hence, removing it will be against the original ViT design. I recommend reading the paper here:
https://arxiv.org/pdf/2010.11929.pdf

Thanks

[night-gale]

Hi! @ahatamiz
Thanks for your reply!

I understand that the cls_token is an essential component of ViT and can be toggled off by passing classification as False.

However, the redundant cls_token I found is in the PatchEmbeddingBlock. It is not reference in the forward method and cannot be turned off by passing argument.

I currently removed the cls_token in my local copy of Monai and everything now works fine.

It would be great if you could double check the implementation of PatchEmbeddingBlock.

Thanks!

[ahatamiz]

Hi @night-gale

Thanks for pointing out the issue. I see that there is an unused cls_token in here
I will address this in a new PR.

Thanks

[ahatamiz]

Hi @Nic-Ma

Thanks for the efforts. I would be appreciate it if this can be addressed in future PRs.

Thanks.

[Nic-Ma]

Hi @ahatamiz ,

OK, sure, I will fix it in a PR soon.

Thanks.