Questions about auxiliary loss values

[guanhaisu]

​I have been greatly inspired by your work and have attempted to reproduce it. While following your provided code, I noticed that the auxiliary loss remains quite high during training, seemingly because it is accumulated across every layer of the model.

​Even in the late stages of convergence, and after applying a coefficient of 0.2, the total auxiliary loss hovers around 40. In contrast, the language loss and KL divergence loss are approximately 4 and 2, respectively. Could you please advise if this magnitude of difference is expected? Is this phenomenon consistent with your observations?

[CURRENTF]

Thank you for your interest in our work! Approximately how many tokens and steps did you train for? Could you provide the hyperparameters you used? The values "the language loss and KL divergence loss are approximately 4 and 2" are not indicative of normal convergence. When the LRC-1.7B we trained converged, the language loss was at least around 2.4, the KL loss was around 1.37, and the total train/loss was approximately 13. Providing your training commands and specific GPU configuration would help me better resolve your issue.

[arianakc]

Hi @CURRENTF I think it is related to aux_loss_scale_factor, what is your aux_loss_scale_factor?

[arianakc]

Another question, how do you choose aux_loss_scale_factor and kl temperature, do you make it have similar magnitude with ntp loss?

[CURRENTF]

Thanks for your interest in our work.

Regarding the first question, for Llama, the aux_loss_scale_factor is typically set to 0.2, while for Qwen, it is generally set to 0.5.

As for the second question, based on ablation studies in small-scale training, we set the KL temperature to 40. We did not attempt to keep the aux loss and NTP loss at the same order of magnitude because they operate on different scales. For instance, in a randomly initialized LLM with a vocabulary size of approximately 1e5, the initial NTP loss is around 12, whereas the aux loss is expected to be significantly larger than this value. Currently, we primarily determine hyperparameter selection through empirical experiments. From my personal intuition and experimental results, setting the aux_loss_scale_factor to a smaller value to force the aux loss and NTP loss into the same order of magnitude tends to result in poorer performance.