Commute Your Domains: Trajectory Optimality Criterion for Multi-Domain Learning

Dear Reviewer vKQ3,

Thank you for your time and effort in this thorough review!

Concerning the practical use, here we propose potential applications. Based on the theory, we may examine the optimality of the training trajectory, thus, making an analysis of different mixing algorithms and provide recommendation on adjusting the mixing weights near a given checkpoint. Moreover, we can compute the P-value of the potential weight intervention during the training and locally change the mixing proportions. Based on the P-value, we may decide for which domains it's better to increase their proportion at the end and decrease at the start.

Validation on the other multi-domain scenarios is an interesting direction for the further research. In our paper, we were concerned with presenting new method for the LLMs training analysis and show some examples on it. We mainly target the new theoretical framework to analyze the ordering of domains during the training. We took a two-languages as minimal example. We believe that results would generalize to other languages and amount of domains.

Some additional experimental results are obtained for the same bi-lingual LLM setup but with domain (i.e. language) imbalance (90% of first language and 10% of second). In this case results show that the effect of order on the loss in high-resource language is much less than the effect on the loss in low-resource language. Therefore, we can benefit the low-resource one without significant draw-down in the high-resource one. The results will be added to the paper. We were unable to perform similar experiments with 7b and larger models, since the pre-training is computationally expensive, but we believe that the predictive power of the proposed theory persists.

The introduction section indeed does not cover some of relevant works. We added some of them in the next version. In particular, we provide references to some studies ignoring the training order and dealing only with the amount of data. The paper on the order of the in-context examples seems quite interesting. Despite that the authors of this paper do not deal with any kind of training, the in-context leaning is still learning, and also benefits from the correct order (of course, by means of different mechanisms).

Thanks again for your insightful feedback. We will improve the text accordingly.

Dear Reviewer kMLr,

Thank you for your time and effort in this thorough review!

We train our bilingual models with Adam optimizer and average domain gradients over several batches to remove the noise due to the stochasticity. According to our empirical results, the theoretical predictions give reasonable estimations as generally we are only interested in the sign of the examined quantity (loss became better or not). Introducing the stochasticity requires averaging over multiple examples during the gradients and Hessian-vector products computation. The remaining discrepancy coming from the commutation of the stochastic differential operators (Limitations section of our paper) is hard to estimate, it is the topic for the following research. Initial guess is that theoretically, this additional term would not be significant. To make actual estimates, more experiments with simplified models of data should be performed.

As you mention, the method, as it is stated, suggests a direction for altering the existing training sequence. Potentially, we hope that our method can be used online during the training by making look-ahead predictions for different training schedules. Currently we have no experimental validation of this approach.

Concerning the empirical limits and practical implementation of Hessian calculation, as we mention in the end on section 4, for a given batch of data, the Hessian-vector product (HVP) can be caclulated with twice as much compute and twice as much memory compared to gradient computation. Our formula for mutual influence of two domains requires two HVPs. Therefore, no limitations on model size are imposed. To reduce noise we used averaging over 100 batches. For multiple domains -- calculation of all pairwise influences would require n(n-1)/2 such computations. However, for determining the best direction to shift an existing weight schedule, we would not need all pairwise mutual influences, only n computations would be needed.

Thanks again for your thorough and constructive comments.

Dear Reviewer 61Cp,

Thank you for your time and effort in this thorough review!

You note that the paper lacks the description of the difference between our approach and classical works on multi-task learning. Actually, the main difference is the addressed problem itself: we study the influence of the training order given the fixed amount of data in every domain. To the best of our knowledge, we give the first theoretical study of this influence. Since it is indeed not clear from the text, we will add the discussion of relevant topics in multi-domain and multi-task learning to the introduction section.

Concerning the noise level (stochasticity) and cost of the additional computations, indeed Hessian-vector products introduce additional complexity. As we state in the paper, this calculation requires two times more memory and four times more compute than ordinary gradient computation. However, we do it infrequently as in practice the changes to the mixing weights would be applied for >1K steps ahead. In this case, 600 (as in our experiment, 100 batches of size 6) Hessian-vector products for 512K (1K steps of batch size 512) gradient computation would give practically negligible computation overhead. We would add these calculation to the paper.

We note that the presented theory can be used to examine existing methods for data mixing on their optimality. We can say which parts of the trajectory aren't optimal. Based on the predictions, we can make changes to the mixing preserving the amount of the data for each domain to improve the model's performance. For simplicity, we've used bi-lingual LLM task and validated on the same distribution, but, it is common to have validation domain that differs from the training. For example, LLMs are trained on some mixture of data and then evaluated on the downstream tasks. In this case, we can re-order data to improve the downstream performance.

Your questions on gradual domain adaptation are insightful. Indeed, since our method is general enough to be applied to any setup with two or more domains, GDA, among others, could be a good application. Optimizing the data mixing from intermediate pseudo-domains using our method is an interesting research topic.

Thanks again for highlighting your concerns about the parameters. If you consider the explanations helpful, we shall improve the manuscript accordingly.

Dear Reviewer zRLt,

Thank you for your time and effort in this thorough review!

First of all, with the presented approach, we show that the data ordering can't be neglected during the training and making small interference into the pre-specified mixing weights, we can change the model's performance forming Pareto front. Also, we showed that if we interfere early in the training, it can severely degrade the performance. There are some works like [1] that fix the amount of examples per-domain and use some kind of a sampling method to form a schedule, but they do not account for the ordering of examples. With our approach, we can analyse such methods on their optimality.

Concerning the potential benefits of changing the weight schedule: for the mentioned setup with bi-lingual model training with equal data budget in each language, we agree that constant schedule works well if the target is also equal-weight scalarization of two language losses. However, for non-equal target loss, other points of our Pareto front would be better. More importantly, we will add the results of the similar experiment with bi-lingual LLM pre-training, when we have significant data imbalance 90% of one language and 10% of another. In this case results show that the effect of order on the loss in high-resource language is much less than the effect on the loss in low-resource language. Therefore, we can benefit the low-resource one without significant drawdown in the high-resource one. This is an example of benefit from the proposed approach.

Another practical application of the study is the training regime were the order of the domains is changed based on the predictions from the theory. As experiments with bi-lingual model suggest we can do that infrequently and the estimation is still quite reasonable due to the averaging effects.

Thank you for the provided list of papers relevant to multi-task learning. We will cite them in the work since they would help to create a general picture of the field. However, in our work we have a specific setting of analyzing the influence of data mixing given the fixed total data budget in all domains. The gradient manipulation (e.g. PCGrad) and even rescaling (e.g. "Multi-Task Learning as a Bargaining Game") are not strictly falling to our setup.

[1] UniMax: Fairer and more Effective Language Sampling for Large-Scale Multilingual Pretraining

Thank you for the answers. I am not still convinced of the practical applications of the proposed method. Without any numerical experiments to show its use-case, it is hard for me as an empirical researcher to validate the practical applications.

Dear Reviewer Xgd4,

Thank you for your time and effort in this thorough review!

Presentation (citation) issues are fixed.

Just as you stated, the theory is formulated for the gradient descent, despite that in practical deep learning other optimizers are used. We discuss the discrepancy in Remark in section 3 and also in the limitation subsection. However our experiment with training bi-lingual models was conducted using Adam optimizer. Though it was not proven theoretically, the update direction in Adam is derived from gradients in the points of training trajectory, and therefore we believe that gradient descent could in our case be a valid approximation.

Concerning the cost of the additional computations, as you mentioned, Hessian-vector products introduce additional complexity. As we state in the paper, this calculation required two times more memory and four times more compute than ordinary gradient computation. However, we do it infrequently as in practice the changes to the mixing weights would be applied for >1K steps ahead. In this case, 600 (as in our experiment, 100 batches of size 6) Hessian-vector products for 512K (1K steps of batch size 512) gradient computation would give practically negligible computation overhead. We will add these calculation to the paper.

Concerning the question on two-domain scenarios, we should say that the method and its theoretical explanation works well for any number of domains. We chosen two-domain setups for toy example and bilingual LLM for simplicity and clarity. The Corollary 3.2 is stated for any number of domains and states that for an optimal trajectory at each step any two domain gradients either commute or the weights of some domains are zero. We will add some words to our theory section, since it was indeed not clear from the text.

We note that the presented theory can be used to examine existing methods for data mixing on their optimality. We can say which parts of the trajectory aren't optimal. Based on the predictions, we can make changes to the mixing preserving the amount of the data for each domain to improve the model's performance. For simplicity, we've used bi-lingual LLM task and validated on the same distribution, but, it is common to have validation domain that differs from the training. For example, LLMs are trained on some mixture of data and then evaluated on the downstream tasks. In this case, we can re-order data to improve the downstream performance.

The curriculum learning is an interesting direction of the further research. Currently, we fix the amount of data for each domain and make only local changes to the schedule. But, in curriculum learning it's more interesting to make global changes and stop learning on some of the domains if they are easy to learn. We will add the topic of curriculum learning to the list of relevant topics in the introductory section.

Thanks again for the valuable ideas and comments.