@sangmichaelxie It seems that the loss used for optimizing the proxy model in the code is different from the one described in the paper.

loss = (pertoken_loss * curr_domain_weights.detach()).sum() / normalizer

In the code, you directly use the proxy model's own loss here to optimize. But in the paper, the loss seems to be the minimax loss which uses the excess loss. So which one should I conform? Or there is something wrong with my understanding. Thanks.