Questions about distill_loss about vit-pytorch

[lucidrains]

@haoren55555 yeah, you could do log softmax for teacher too by setting log_target = True https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html . i'm just rolling with what pytorch offers

[vivekh2000]

That was insightful. However, in your implementation of class DistillMixin, you are passing the cls token from the MLP lead.

The MLP head of the distillation token is implemented using LayerNorm, whereas the cls token MLP head is implemented in vit.py without LayerNorm. Why is it so? What am I missing? Please correct me if I am wrong.

[vivekh2000]

@lucidrains
I have also noticed that in the paper, the authors mentioned that at test time, they fused the two heads, i.e., MLP heads for cls token and distilled token.

for which we add the softmax output by the two classifiers to make the prediction. (p.8, paragraph 3)

This fusion of two heads is not implemented in the forward method of the DistillMixin class. Please let me know if it is implemented elsewhere or if I am missing something. Thanks for the excellent codebase.