Questions about distill_loss

[haoren55555]

sorry to bother, I see the distill_loss in distill.py as :
distill_loss = F.kl_div(
F.log_softmax(distill_logits / T, dim=-1),
F.softmax(teacher_logits / T, dim=-1).detach(),
reduction='batchmean')
I wonder why the teacher part uses the softmax function rather than log_softmax one, thanks.

[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.