Add masks when calculating the CAKLD coefficient

[HsuWanTing]

I noticed that the coefficient of CAKLD is calculated without masks. (https://github.com/DD-DuDa/BitDistiller/blob/main/train/train.py#L375)

Have you tried adding the mask to calculate the mean on the real input part only?

[BrownianNotion]

Masking makes sense to me, @HsuWanTing did you try this? If so, did you observe any differences in the gamma coefficient for the cakld loss and were they large?

Update: So I tried this myself and there's a few percent difference. But I'm also confused as to why we are doing (prob * mask).mean() instead of (prob * mask) / sum(mask), doesn't it make sense to only average the probabilities over the non-pad tokens?

I ran a few simulations for each method below using the 8000 wikitext + alpaca mixed dataset, there's a few % difference in the gamma:

Method 1: (prob * mask).mean()

Gamma: unmasked 0.8069294095039368 | masked 0.778532087802887 | samples=6388
Gamma: unmasked 0.789540708065033 | masked 0.7581791877746582 | samples=6388
Gamma: unmasked 0.8421652913093567 | masked 0.8078484535217285 | samples=6388
Gamma: unmasked 0.8558153510093689 | masked 0.8342512249946594 | samples=6388
Gamma: unmasked 0.8523220419883728 | masked 0.8338661193847656 | samples=6388

Method 2: (prob * mask) / sum(mask)

Gamma: unmasked 0.8471 | masked 0.8992 | samples=6388
Gamma: unmasked 0.8615 | masked 0.8991 | samples=6388
Gamma: unmasked 0.8476 | masked 0.8991 | samples=6388
Gamma: unmasked 0.8514 | masked 0.8987 | samples=6388

In conclusion, the first method decreases gamma a few %, as we'd expect from averaging over the same number of probabilities, but with some of them now equal to 0, while the second method increases gamma, since we average over a fewer number of probabilities that should on average be higher than the pad token probability.