Title: PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts
Venue: NeurIPS 2025
Accepted: yes
Keywords: Multi-task learning

Large language models

Large-scale optimization
Abstract: Modern foundation models are trained on diverse datasets to enhance generalization across tasks and domains. A central challenge in this process is determining how to effectively mix and sample data from multiple sources. This naturally leads to a multi-task learning (MTL) perspective. While prior work in MTL has emphasized mitigating gradient conflicts, we observe that large-scale pretraining scenarios—such as multilingual or multi-domain training—often exhibit little to no gradient conflict. Motivated by this observation, we propose $\textbf{PiKE}$ ($\textbf{P}$ositive gradient $\textbf{i}$nteraction-based $\textbf{K}$-task weights $\textbf{E}$stimator), an adaptive data mixing algorithm that dynamically adjusts sampling weights during training. PiKE exploits non-conflicting gradient interactions to minimize a near-tight upper bound on the average loss decrease at each step, while incurring negligible computational overhead. We provide theoretical convergence guarantees and  show that PiKE  outperforms static and non-adaptive mixing baselines. Furthermore, we extend PiKE to promote balanced learning across tasks. Extensive experiments on large-scale language model pretraining confirm that PiKE achieves faster convergence and improved downstream performance compared to existing approaches.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the problem of creating a batch composed of multiple tasks for the training of large language models (LLMs). To this end, the authors adopt a multitask learning (MTL) perspective and pose the problem as a dynamic gradient weighting problem (studied in MTL in the context of negative transfer). Then, the authors verify existing claims on the lack of gradient conflict to larger models than before, and study the effect of weighted-mixing the batch creation under SGD and bounded gradient cosine similarities. This ideal scenario is then relax for more realistic scenarios where we may use other optimizers and expect a balanced decreases in task losses. The authors then verify the effectiveness of the proposed approach in a number of datasets and GPT-2 models, comparing with previous baselines.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n1. The paper is generally well-written, making all their points clear and easy to read.\n2. I like the idea of posing the generation of batches in the context of LLMs as a MTL problem, and their connections with gradient conflicts.\n3. The empirical results are convincing.\n4. The validation on the low-conflict of large models is very welcome.\n5. The proposed algorithm is well-motivated, although I would've liked a bit more of detail in the derivations.\n6. The problem is clearly of relevant for the community.\n\n**Weaknesses**\n1. I would have liked to see the performance of the idealized version (or the one which updates every X iterations) to see its drawbacks.\n2. While efficient, the algorithm still needs to compute and store the gradients for each task. Moreover, I could see how rich-get-richer situation, where tasks with low representation in the batch get worse estimates of their variance due to the smaller number of samples.\n3. It is not completely clear the flowchart to use the proposed method: Should we pretrain the weights? What value of $\\tau$ should we pick?\n\nQUESTIONS:\n1. How do you derive the mirror descent step in section 3.2?\n2. Is there any assumption regarding the capacity of the neural network to hold remark 1?\n3. What are the empirical values for the bound of the cosine similarity that you find in practice? Based on l165, it seems that gradients are almost orthogonal for LLMs.\n3. How is the bound in Theorem 3.4 tight (L164)? Isn't this bound some form of a truncated Taylor expansion?\n4. Do you have any intuition on how to select $\\tau$?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes PiKE, an adaptive data mixing algorithm. PiKE draws inspiration from MTL, which focuses on mitigating gradient conflicts. However, the paper finds that in pre-training, gradient conflicts are rare, so using this finding, they derive a theoretical bound on the decrease in training loss as a function of mixture weights and gradient properties. By maximizing this decrease, they propose PiKE, which updates weights using mirror descent as a function of gradient norm and gradient variance. An additional balanced PiKE is proposed to encourage good performance on all tasks. Pretraining experiments are conducted on multilingual C4 and GLAM.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\n## Quality:\n- The series of building blocks establishes why the authors arrived at this particular setup/method, providing strong motivating evidence.\n- Theoretical derivation aside, this idea of having gradient variance in the mixing algorithm makes intuitive sense to me. One unsatisfying thing about existing algorithms is that they do not encode the possibility that existing domains are poorly defined/noisy. To me, this variance term factors in this upstream noise into the mixing algorithm, which seems like the right direction.\n\n## Significance:\nThe work proposes a theoretically motivated data mixing algorithm that is also computationally efficient. It differs from several other data mixing works that simply propose an algorithm and do not describe its assumptions. I hope this work can continue to push for more principled data mixing algorithms. \n\n## Clarity: \nTheoretical results are presented at a good level of abstraction and are easy to follow.\n\n# Weaknesses\n\n## Originality:\n- It would be interesting to compare the theoretical derivation in this paper to the derivation in https://arxiv.org/abs/2310.15393. That paper uses first order approximations to arrive at an objective that is linear in $w_k$; except that it uses a gradient dot product rather than gradient norm, and also doesn't have the variance term. For clarity, it would also help to explain why there is no gradient dot product term in PiKE, since its nod to MTL is based on gradient conflict/similarity.\n\n## Quality:\n- While example 2.1 shows how an adaptive mixing procedure is optimal, the loss/data assumptions are not realistic. It would be good to have a real pretraining experiment that motivates adaptive mixing. As an example, in https://arxiv.org/abs/2411.05735 Table 10, they do a two-stage brute-force search over data mixes and show that the best two-stage approach does not result in the same $w$ at each stage, for certain domains.\n- One property of the mirror descent weight update is that in the limit, the updates move $w$ less and less (because the weight update accumulates). For instance, in your figure 8 you can see that the mixture weights converge fairly quickly. As a result, $w$ is not able to \"move dramatically\" later on in pretraining, even if that may be optimal.  \n- In Table 1, it is surprising to me that regular PiKE performs poorly in terms of validation PPL. My understanding is that PiKE aims to maximize the training loss reduction at each time-step, so it would be good to have some verification of regular PiKE's performance (i.e., comparing training loss curves for PiKE + baselines).\n- In Table 2, it is surprising to me why PiKE (regular and balanced) have high perplexity but do well on downstream tasks. I understand that perplexity and downstream are not always correlated; however, PiKE and most of these other data mixing methods are designed to minimize some function of perplexity/loss - so why aren't they able to minimize these properly in practice?  \n\n## Clarity:\n- No related work in the body; more comparison to other data mixing methods would be helpful \n- To demonstrate how balanced-PiKE promotes fiar learning on many tasks/domains, it would be helpful to also display the range/variance across tasks/domains in Table 1. After writing out the ranges of PPL in table 1 by hand, only then did I convince myself of what balanced-PiKE is doing.\n\nQUESTIONS:\n1. Can you compare derivation/results with DoGE (https://arxiv.org/abs/2310.15393)?\n2. More realistic pre-training experiment to motivate when adaptive data mixing is optimal.\n3. More validation of regular PiKE, such as arguing that the training loss curve goes down more quickly than other mixing approaches.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose the framework of PiKE, aiming to solve the multi-task training in large language models. The authors propose an adaptive data mixing algorithm that dynamically adjusts sampling weights during training, ensuring efficient and effective learning across diverse datasets. with rigorous theoretic analysis, this algorithm can near the optimality. In experiments conducted on pretraining, including the mC4 and GLaM datasets, PiKE demonstrates superior performance in multilingual and multi-domain settings, achieving improved downstream performance compared to existing approaches.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:\n1. PiKE dynamically assign sampling weights for each task based on the gradient norm and variance, leading to better applicability on  downstream taks.\n2. This paper includes formal convergence guarantees and a solid theoretical analysis.\n3. PiKE outperforms all other baselines in downstream tasks, highlighting its effectiveness in multilingual and multi-domain settings.\n\n**Weaknesses**:\n1. The experiments primarily concentrate on the pre-training phase. It would be beneficial to explore how PiKE performs when directly trained on downstream datasets jointly, to assess its applicability in end-to-end training scenarios.\n2. The hyperparameter $\\zeta_2$ significantly influences the results, as evidenced in Table 11. For a thorough understanding of the outcomes, it is crucial to clearly provide the specific value of hyperparameter utilized in the experiments.\n3. While PiKE shows promise in balancing learning across tasks, its impact on dataset imbalance, particularly for low-resource and high-resource tasks, remains unexplored.\n\nQUESTIONS:\n1. How does PiKE perform when directly trained on downstream datasets, rather than just during the pre-training phase?\n2. Could you provide the specific values of the hyperparameters used in the experiments?\n3. Can PiKE’s adaptive weighting scheme be extended to incorporate dataset size or task-specific downstream importance, beyond just gradient statistics?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: SUMMARY:
The paper proposes PiKE - a data sampling approach during foundation model pre-training and proposes "how to effectively mix and sample data from multiple sources". The paper postulates that pre-training LLMs on a big language corpus usually has low gradient conflicts and provides PiKE as a method that leverages this phenomenon. Empirical results on multiple datasets and varied model sizes confirms the viability of PiKE.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The paper is extremely well written with sound figures and colors used appropriately, improving readability and understanding. I especially appreciate the box detailing the key features of PiKE and relating to the structure of the paper .
2. Empirical results are solid, with clear improvements over baselines. The datasets used are plentiful.
3. The paper has the potential to be high impact with many key insights used in practice.

Weakness:
1. How do these insights translate into other problem domains, especially more structured tasks? Currently, the evaluations have been performed on simple QA style datasets and multilingual (en/hi/de). I wonder how the results would be on structured understanding, like tabular data, time series, code, etc. Does the gradient non-conflict hold there too?
2. How do hyperparameters affect convergence? How do hyperparameters $k$, $\psi_1$, $\psi_2$ affect the performance. I understand that due to high running costs, it is hard to perform a comprehensive empirical study, but an intuitive paragraph might be helpful.

QUESTIONS:
See Weakness.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03