Closed tt6746690 closed 11 months ago
As long as there are a sufficient number of tokens in the batch, the two objectives should be about the same (and they converge to the same value with more samples). In my tests, the two objectives seem to return almost the same average domain weights. Our implementation follows Pytorch's class weighting implementation for the CrossEntropyLoss, and makes sure that the loss scale is preserved even when there are some domains missing (could be important for situations where there are more domains).
About your code, I think that if you choose to normalize by the observed token frequency, you don't need to do this part:
# if doing non-uniform sampling, normalize by inverse sampling weight
train_domain_weights = train_domain_weights / self.sampling_weights.to(train_domain_weights.device)
train_domain_weights = train_domain_weights / train_domain_weights.sum()
Thanks for the clarification!
It seems that the loss implementation (https://github.com/sangmichaelxie/doremi/blob/main/doremi/trainer.py#L360) is not exactly the same as the loss in the paper. In the implementation, the normalizer is
Σ_i α_i Σ_{x\in Dᵢ} |x|
but should just beΣ_{x\in Dᵢ} |x|
for samples from i-th domain. Any comments on this observation?Here is the code that implements the loss in the paper. It seems you get smoother domain weights using the following implementation.