PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts
Authors: Zeman Li
Yuan Deng
Peilin Zhong
Meisam Razaviyayn
Vahab Mirrokni
Venue: NeurIPS 2025 spotlight ()
Keywords: Multi-task learning
Large language models
Large-scale optimization
Primary area: 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.

PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts
Zeman Li
Yuan Deng?
Peilin Zhong?
Meisam Razaviyayn" 1.2
Vahab Mirrokni
University of Southern California Google Research
{zemanli, razaviya edu www.omoiiioooooooeeee.o. com
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 multi-task learning (MTL) perspective. While prior work in MTL has emphasized mitigating gradient conflicts, we observe that large-scale pretraining scenarios-suchs multilingual or multi-domain training-often exhibit little no gradient conflict. Motivated by this observation, we propose PiKE (Positive gradient interaction-basel K-task weights Estimator), an adaptive data mixing algorithm that dynamically adjusts sampling weights during training. PiKE exploits non-conficcing gradient interactions to minimize 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.
1 Introduction
Foundation models, such as large language models (LLMs), owe their strong generalization and multitask abilities to pretraining on diverse datasets spanning multiple domains 62
36
10]
The
effectiveness of these models depend heavily on the composition of their training data 29
18 How-
ever, standardized practices for curating optimal pretraining data are lacking. Common approaches
involve heuristic filtering, deduplication, and categorization into heterogeneous domains (c.g., The Pile [2 has 22 domains; GLaM [18] has 6). Even after such preprocessings, determining the optimal data mixing remains key challenge-amplifeed by the scale of modern models and datasets.
A common strategy in pre-training is smaller proxy models. For example, mT5 71 weights datasets by relative size, while GLaM [18 downstream performance from proxy models. DoReMi [69] also relies on proxy models, using group distributionally robust optimization (group DRO) to set dataset weights. However, these approaches have notable limitations. First, their optimality is unclear: heuristic methods lack theoretical backing, and policies from small models may not necessarily transfer to larger ones [74). Second, proxy models introduce substantial computational overhead, often scaling linearly or worse with the number of domains. Third, static weights fixed at initialization may become suboptimal as training progresses 77 a limitation discussed further in Section 2
fixed data mixtures, typically chosen heuristically via
In this work, we frame adaptive data mixture selection problem as a multitask optimization problem, enabling a principled approach dynamically adjusting data mixture. This view is natural: each data domain typically is related to a (set of) tasks and yields a distinct gradient. Prior multitask optimization methods 75 66 [6] focus on resolving gradient conflicts that impede convergence. However,
Work done while interning Google Research.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

30
Mix
DoReMi
3.4 berler
Round-Rohhn
45
GLaM
29
Random
Accuracy (%
PiKE
Accuracy (%)
PiKE
228
10
27
35
26
30
25
0
20
40
60
80
100
120
0
20
10
60
80
100
120
# Training Steps (K)
# Training Steps (K)
Figure I: PiKE adaptively optimizes task weights in pre-training, outperforming baselines. Left (1B models, multilingual C4 en/hi): PiKE boosts average downstream accuracy by 7.1 % and reaches baseline accuracy 1.9x faster. Right (750M models, GLaM six domains): PiKE achieves 3.4% higher accuracy DoReMi (69
PiKE scales efficiently with model size and number tasks. Detailed results are Tables
and 16
most impractical for current LLMs due to their O(Kd) memory cost for storing K gradients or O(K²) computation from repeated backpropagation, A notable exception FAMO [38] which scales more efficiently. Additionally, most MTL methods are tailored for vision tasks where gradient conflicts are prevalent, a condition less common in LLM pretraining: For example, Gradnaccine [66]
observed that multilingual BERT (178M parameters), task gradients mostly positively aligned or nearly orthogonal. We extend this finding to much larger autoregressive, decoder-only models (e.g.. IB parameters), which better reflect modern LLMs. Our results show minimal gradient conflict in multilingual and multi-domain pretraining (Section 2). suggesting that modern LLMs naturally exhibit cooperative (or sometimes hearly-orthogonall gradient structures- p-t--tiaeeenll reducing the need for explicit conflict mitiigttinn. This motivates the central question of our work:
Can we design an effective adaptive data mixing strategy based on multitask optimization, that exploits non-conflicting nature of gradients and scales to large settings with minimal memory and computational overhead?
To address this, we propose PiKE, an adaptive data mixing that is empirically effective and scalable for pretraining LLMs. PiKE enjoys theoretical guarantees while remaining practical for large-scale settings- with many tasks, large models. and diverse data-at negligible memory and compute overhead. We summarize PiKE's key features below; related work discussed in Appendix
Key Features of PiKE
1. Adaptively adjusts mixture weights using per-task gradient magnitude and variance.
Theoretically, achieves near-optimal per-iteration objective decrease and enjoys conver- gence guarantees (Section 3)
3. Incorporates tilted empirical risk minimization
promoting balanced learning across tasks and preventing task under-repreesentaion (Section 3.3
35
4. Scales efficiently (linearly) with model size and the number of tasks (Section
5. Consistently outperforms existing methods across different scales (1 10M to IB parameters) and scenarios (multilingual domain mixing) (Figure T and Section
We aim train single model with parameters 0 e Rd to perform K 2 tasks simultaneously. Each task k is associated with smooth (possibly non-convex) loss & (0, x) Rd d I where is the data point. It is common minimize the total loss:
Problem Definition and Notations
K
min C(0)
Σ A=1
0RR
(0; For where Dj represents the data distribution for task We define La (0) |e notation, represents the Fuclidean norm, Tr(-) denotes the trace operator, and a function h is

---

0.3
Mix
45
books-forun books-web
Round-Robin
Raudom
Cosine Similarity
Accuracy (%)
0.2
web-wiki
40
0.1
35
0.0
30
0
20
40
60
80
100
120
20
40
60
80
100
120
# Training Steps (K)
# Training Steps (K)
Figure Left:: Average accuracy on downstream tasks (ArcE, CSQA, HellaSwag, PIQA) for 750M GPT-2 style models pre- pre-trained with Mix. Gound-Robin, and Random strategies. See Appendix [.1 for more. Right: Task gradient cosine similarity for 750M GPT-2 style model pre-trained on GlaM datasets. "datal-data?" indicates gradient similarity between tasks datal and data2. Further results in Appendix F2
L-Lipschitz if (0) h(0)) < L - 0'll for any 0, 0' in the domain of /(-). A function f(-) is L-smooth if its gradient is L-Lipschitz continuous.
Main Building Blocks for PiKE
This section presents our main observations which form the main building blocks of PiKE.
Bulding Block #1: Mixing Domains per Batch Improves LLM Generalization
Optimizing equation T with stochastic methods, such Adam SGD, requires forming batches from one more tasks at each step. The batch selection strategy strongly influences model perfor- mance 24 74
69 40 Even with fixed data proportions, key question remains: how should one form a batch from K data domains at each step?
Batch construction is critical design choice in multitask training, it directly impacts learning dynamics and final model performance 4 [2 74
69 40). We focus on three standard strategies:
Random, Round-Robin, and Mix. Let first define these methods assuming static, uniform sampling weights (i.e., 1/ K per task). Let b, = (br. bt,k) denote the number of samples from each task
Dr iteration t, with total batch size b ΣK k= bi, The strategies are defined as:
CE b ee
Random, where k* ~ Uniform((1,... K})
b,
b C(t mod K)+1 ( k)
Round-Robin
Mix
where e₁ R is the k-th standard basis vector. That is, Random selects domain per batch, Round-Robin cycles through domains, and Mix includes all tasks in each batch.
Historically, Mix has been widely used in computer vision multitask settings 14
45
while earlier language modeling efforts favored Random or Round-Rold [4
42 Recent studies
on large-scale language models 51 have revisited this question and found that Mix generally
performs best mess-esppccially in diverse-date settings 18 IT 69 22 (65). Our experiments reaffirm this
trend: across wide range of tasks and scales. Mix consistently outperforms the alternatives (Figure 2 Figure 4p. This observation underpins the design of our proposed method, PiKE.
Bulding Block #2: Large-Scale Training: Low Gradient Conflict & MTL Scalability Challenges
The Mix batching strategy offers natural lens for analyzing data mixing in multitask learning (MTL.). When batch of total size is constructed with by samples from each task k the overall gradient at

---

iteration t, denoted is the average of the individual sample gradients:
by
g:
Vla (0;3H)))
=
gr,ks
where gr.k formulation highlights how batch gradients blend contributions from different tasks and serves as the foundation for analyzing task interactions.
Ei (0 Fx.;) is the average gradient from task k. and Tk,i ~ DR- This
central challenge in prior MTL literature gradient conflict oppose the overall update direction. Formally, conflict exists if: shared update could increase the loss for task k.
75 66
where task gradients
(g
Bt.k)
0, indicating that the
While such conflicts are well-documented in vision tasks, we observe they are in large-scale lan- guage model pretraining. In particular, as shown in Figures 2 [ 6 task gradients are generally non conflicting in pretraining LLMs. This observation shifts the our goal: instead of mitigating conflict, we can leverage naturally non-confticting gradi-
10
10
Training Objective
102
ents to improve training efficiency.
10
Adapiie Movnn
Static Mixing lb, 40,20)
This insight renders many conflict-aware MTL methods-such PCGrad [75]. AdaTask (73)
Static Mixing (b, -0.40)
1
Static Mixing (b -0.60)
Static Mixing lb -0.70)
MGDA [56] and NashMTL 47 less effective
Stati Mixing (b. +0.8o)
for pre-training LLMs: PCGrad acts only when conflicts occur (rare in LLMs) and thus performs similarly to simple Mix. AdaTask treats different tasks as completely independent and thus does not utilizes potential cooperative nature of tasks, lead- ing to slower convergence. More fundamentally. these methods incur substantial memory and compute overhead, making them unsuitable for large-scale models (see Appendix [A] for discussions). A notable exception is the recent work FAMO (38), which is designed for scalability.
0
500
1000
1500
Iteration
Figure 3: Adaptive mixing consistently outper-
forms static mixing in the Example 2.1
In summary, the lack of gradient conflict in LLM pretraining, combined with the limitations of existing MTL approaches, motivates the need for scalable methods that exploit non-conflicting gradient structures rather than fixating on rare conflicts.
Bulding Block #3: Adaptively Changing Mix Sampling Weights Is Necessary
Prior work using the Mix sampling strategy typically relies on fixed (static) sampling weights, keeping (b,...., bx) constant throughout training. However, dynamically adjusting batch composition can significantly enhance efficiency. We illustrate this fact with a simple example:
Example 2.1. Consider training K tasks with losses & (0; x1) (0Tec) + and
(2 x2) (0Te)) 2 x 0, where e1 [1 TO_ e2 [0 1]_. and 0 RR.. Data for task follows Φ NO, f). while task 2 follows 22 (O, o²1). The loss for task simplifies to (0) ( )² Using by samples from task k; the gradient at iteration given by
1 g: bn + b2 (breje) + b2e2e?) O,
+ 62
where z ~ NO,
1) with = bj + bz. Updating 0, via SGD, 0,41 1 = = 77gt. we have
bj 0z07, ! 17 by 20g + 1/2 b bjof + b20g
E[C(0(+1) (1
62
where 0,, and 02,, denote the first and second component of the vector 0, The derivation details of equation can be found in Appendix G.1 Letting w1 := bi 6 W2 ==
by and relaxing them to 6
take real values, we can optimize the mixing weights W1 and u2 as
b
(0
of) + (0 02 + 02
: =1 II
(4)
8 + 03.

---

and w² = wi where [T) min(max(5 0}, 1 is the projection onto [o, 1]. This result shows that optimal batch composition should change over iterafes to maximize training efficiency.
the superiority of the adaptive mixing approach based on equation 4 over various static mixing strategies. Moreover, the adaptive mixing strategy in this example does not require any hyperparameter tuning, while finding the best static mixing requires tuning.
Despite its simplicity, this example mirrors key aspects of MTL in large models: 1) The optimal solution 0" 0 minimizes all task losses simultaneously, reflecting the high expressive power of large models. 2) Task gradients are non-conflicting, resembling real-world gradient interac- tions observed in Building Block #2. Moreover, equation (4) further reveals that optimal data mixing depends on (1) the gradient norm squared per task [ = Bf. [2) 83 and (2) gradient variance (o o2). As will see in the next section, these factors play a crucial role in defining optimal mixing strategies for more general settings.
3 Method
3.1 PiKE: Conceptual Version
To develop our method, we first start by quantifying gradient conflicts:
Definition 3.1. For a given point 0, we say gradients are g-conflicted (with g 2 0) if, for all task pairs j,j # k,
v( [VLx [V <
(O),
The above definition is implied by lower bound on the gradients cosine similarity. In particular, if (0)) (.) (O))
[1; (e) 2 - then the gradients are c-conflicted for c = 2/2. Therefore, experiments in F.2 show that e is typically small for LLM training.
Section and Figures 5 and
G in Appendix
While task gradient conflict are rare in LLM training, we observed in Section that task gradients are not fully aligned either. To quantify the level of alignment, we define the following concept:
Definition 3.2. For given point 0, we say that the gradients are ž-aligned (with 2 0) if, for all task pairs j * k,
(V), (L) < <1
Notice that Definition 3.1 and 3.2 always hold for č and 1/2. However, smaller values allow for more refined analysis and more desirable properties. Notably, when both é and e are small (as we observed in our experiments), the value of V is small if and only if (x))))) is small for all (see Lemma G. in Appendix G). To proceed further, we make the following standard assumption:
Assumption 3.3. Per-task gradients are L-Lipschitz, unbiased, and have bounded variance, i.c., Vk
[ - (O2)) < L|[0; - 0211. 02 VL,
(5)
En.DD [f: VC+ (0), A0
(6)
EE EE VO, V., (0) < of 4A
Using mix batch with by samples per task k; the estimated gradient follows equation ( The next theorem characterizes the descent amount in one step of SGD under low conflict conditions:
Theorem 3.4. Assume Assumption 3.3 holds and the gradients at θ, are c-conflicted and e-aligned,
with g K K-2-0/06 for all k. If gradients are computed using the Mix strategy equation 22, then:
K
LT² 212 ok
K
Ln
Σ bn ( k=1
+
(O).) (8)
212
k=1
for SGD update O.+1 = 0, 1g- Here, 3 2 min, (1 + gl-g(-K- + 2- by b) 1) 2 E(K 1), and the expectation is over batch sampling randomness under the mix strategy (b.....,bk)
5

---

A formal proof is provided in Theorem G.3 the decrease in the objective depending on the mix batch composition (b). (......, bx) As we show in Theorem G.4 in the appendix, this upperbound is tight in the non-conflicting, non-aligned regime of C = - (which is good approximation for pre-training L.LMs according to our experiments).
(Appendix G). Theorem 3.4
provides an upper- poper-bound on
According to Theorem 3.4, to maximize descent in Mix sampling, we need to minimize the RHS of equation (8. Thus, relaxing wa bu/b to continuous values, we need solve:
K
min W0.....KK' 20
Σ
wx + W s.t.
wk
1=1
k=1
where Aạ A [ +
Laj²
2b of and Kik &
Using KKT conditions, the
optimal solution is given by w = max {0, # where is chosen such that Σk= "x (see Lemma
G.2 Appendix G This leads to the conceptual version of PiKE (Positive gradient Interactions baaset K-task weight Estimator), summarized in Algorithm 2 in Appendix
The conceptual version of PiKE (Algorithm 2 adaptively adjusts sampling weights. This adaptive adjustment makes the stochastic gradients biased, i.e., E ge] * V ((2)). Due to this introduced bias, the classical convergence results of SGD the convergence of conceptual PiKE:
longer be applied. The following theorem establishes
Theorem 3.5. Suppose the assumptions in Theorem 3.4 hold and the Conceptual PiKE Algorithm (Algorithm 2 in the Appendix) initialized at 0g with the SGD optimizer in Step 10 of the algorithm. Let Ar [())) ming [(0) and O max = max; σg- Suppose ô >> 0 is a given constant and the
stepsize 1) S Lo /4465 Then, after T
2 iterations, Algorithm 2 finds a point O such that nB6
B LA (0)1 8
Vk: =k= K.
(10)
Particularly, we choose 1l
36 Lo Lo Lo Logo then the Conceptual PiKE algorithm requires at most
T = 2o maz (b41)) iterations to find a point satisfying equation 10
82
This theorem (restated as Theorem G.5 is proved in Appendix G
Remark 1. Theorem 3.5 guarantees that, given enough iterations, all tasks are learned jointly, i.e., the gradients of losses become small. Moreover, the convergence rate = 0(1/8), which matches the optimal iteration complexity for smooth, nonconvex stochastic optimization.
Remark The Conceptual PiKE algorithm maximizes the expected decrease each iteration by finding the optimal mix batching strategy. This property leads to improved iteration complexity bound than (static) uniform mix batching, as discussed in Appendix G.3 when e and ē are small.
3.2 PiKE: Simplified Computationally Efficient Version
k Solving equation 99) requires estimating ( and for 1i= however, per-iteration estimation of these terms with sufficient accuracy often necessitates large batch computations, thereby impeding convergence. To speed the algorithm. we can update these estimates every To iterations. However, this can cause abrupt changes (wi..... ... wx). leading to instability, especially with optimizers like Adam, where sudden shifts may disrupt momentum estimates. To mitigate this, we update (Wi,...,wy. using single mirror descent step on equation ensuring gradual adjustments:
gradiens noem
abo²
gradiest variance
Wh +- wk exp
con(B Loyywe)
(0)
25
followed by normalization:
w/|lw/11 where is the mirror descent step size.
Even after the above simplifications, fine-tuning L, 7, Or, and can be challenging in practice. We simplify this finetuning by noting two observations: 1) The coefficient in front of of is constant, independent uig. 2) For small y and w& 1, the coefficient in front of | Lx (0) remains nearly constant: any(33 - Lopyws) ss on33. Therefore, in our practical implementation, PiKE employs tunable constant coefficients for terms related to task-specific gradient variance and gradient norms, which simplifies its application. The resulting algorithm is detailed in Algorithm

---

Algorithm PiKE: Positive gradient Interaction-based K-task weights Estimator
Input: 0, To. total batch size b, task dataset Dx, hyperparameters Çi and (2. prior weights w'
2: Initialize: wi & 1/K
f w/
3: for = 0, ... do
4: if't mod To = 0 then
5: Estimate (0:)ll and of for every k
6: Wk f- wk exp <i (
26 o
7:
w/ ///lw/
8:
(b)..... bx:) + round (b(u), (b.,,...., wx))
9:
end if
10: Sample by data points from each task &
11:
Compute the gradient g using the estimates samples
12:
Update: oz+1 < Optimizer(1,, Or. g)
end for
Remark 3. The inclusion the variance-related term, of. is crucial in PiKE; its omission would cause PiKE to disproportionately prioritize tasks with currently large losses even if their gradients are very noisy (high variance). Our ablation study, presented Appendix underscores the importance of this term for achieving robust and balanced multi-task performance.
Remark 4. PiKE updates sampling weights only every To steps. our experiments, this additional step incurs negligible memory and computational overhead (see Appendix F.4
3.3 Balanced-PiKE: Balanced Learning Across Different Tasks
While Algorithm optimizes the average loss across tasks equation II. this objective lead to unbalanced learning. In practice, some tasks may dominate due to larger datasets or easier loss reduction, causing others to be under-optimized. This imbalance is problematic when certain domains (e.g., code, math) are more important to user, or when balanced performance across all tasks is desired. To address this shortcoming, we consider balance-promen promoting objective based on filted empirical risk minimization
35 also known as the o-fairness utility [46]:
K
min O(T; = logg Σ
Cc((0)
(1I)
0
This objective interpolates between average loss (r - 0) and worst-case loss (r o0), allowing
users to control the trade-off between efficiency and fairesss. Moderate values of (7 R 1 ~ 3 in our experiments) encourage balanced learning without sacrificing overall performance.
To leverage our PiKE developments in solving equation (1I,, we need to connect it to our intital objective in equation []]. The following lemma achieves this through the use of Fenchel duality (53
Lemma 3.6. Assume (+ 0, vk and let 0 < E R. Then equation
11 is equivalent to solving
min
max YERK
Σ Ca
2/
log
1/1
(12)
Ek V6
Moreover, for any fixed 0 the inner maximization problem is maximized at vi
Σk coltoss Vk.
This lemma, which proved in Appendix G.4 provides natural alternative minimization algorithm for solving equation (1I: Fixing y. equation 12)
reduces to a weighted minimization over tasks, where regular PiKE sampling with proper weights yi front of each loss can be applied to determine the optimal mixing strategy. On the other hand, fixing 0, the optimal solution y can be computed algorithm, described in Appendix [D] which balances overall loss minimization and balanced/fair learning of all tasks. Although prior MTL 3 such methods often grapple with scalability limitations, critical aspect where PiKE is designed to excel through efficient large-scale performance.
according to Lemma 3.6) This leads to Balanced-
literature explored fair learning
7

---

Table 1: We report the perplexities (lower the better) the validation split of multilingual C4 datasets. We also report the accuracies (%, higher the better) of different models on HellaSwag and its corresponding translated version. Bolding indicates the best model in the task; Metrics means the average across different tasks, Additional results can be found in Table 15
C4 (ca)
(ba)
C4 (dc)
HellaSwag (cn)
HellaSwag (hi)
HellaSwag (dc)
Perplexity
Perplexity
Perplexity
Perplexity
Avcuracy
O-shot
O-shot
O-sbot
C4 (hd), C4
GPT-2 large
params, Layers default, 120K training steps
Mix
8.29
11.13
4.45
9.29
27.5
28.
27.1
27.6
Round-Robae
8.41
11.31
4.97
9.46
26.5
27.6
26.7
26.3
FAMO
8.25
11.04
4.48
9.23
27.2
27.3
26.9
27.3
ADO R0
8,30)
11.12
4.45
9.31
27.5
27.7
27.5
27.2
9.56
28.7
26.2
Balancel-PiKE
8.29
11.12
4.46
9.31
27.9
28.3
27.4
28.0
Ralancad-PIKI-
Balanced-PiKE 8.42
10.02
6.30
8.94
28.9
31.2
26.9
28.6
Table 2: We report perplexity (lower is better) the validation split of the GLaM datasets, averaging perplexities across six domains. We also compare the accuracies (%, higher the better) of different models four different Q/A tasks. PiKE (Uniform) means PiKE using initial uniform sampling weights and PiKE (GLaM) means PiKE using GLaM tuned weights as initial weights. Bolding indicates the best model in the task, underlining indicates PiKE beating Mix, Round-Robin, Random methods. Additional detailed (per-domain) results with different model sizes can be found in Table T6
GLaM
ArcE
CSQA
HellaSwag
PIQA
Perplexity
Accuracy
7-shat
7-shot
7-shot
7-shot
Six domains of GLaM dataset, GPT-2
large style, 750M paramns, 36 layers default
Mix
12.77
46.4
47.2
39.6
37.9
60.9
Round-Robo Robin
12:98
44.3
43.5
36.7
36.8
60.3
Random
12.99
42.7
41.7
34.2
36.6
58.2
FAMO 38
13.25
45.0
43.7
40.0
36.4
59.8
ADO
12.77
45.9
45.5
38.7
38.1
61.1
GLaM
18
13.20
45.3
46.9
39.8
38.0
564
DoRcMi
13.25
46.5
48.6
PiKE (Uniform)
13.22
47.6
49.6
43.2
37.2
60.4
PIKE (GLaM)
13.35
48.1
49,8
43.5
38.0
61.2
Balanced-PiRG
13.21
47.5
48.8
42.5
37.6
61.2
Balanced PIKE (r
13.26
47.2
48.8
41.5
37.2
61.3
Balanced (7
13.19
48.2
49.3
42.6
38.5
62.4
4 Experiments
We evaluate PiKE two multitask pretraining scenarios: 1) Pretraining language models on multilingual mC4 dataset [71), dataset covering diverse languages from Common Crawl corpus. 2) Pretraining language models on the GLaM dataser [18), an English dataset spanning six domains. As we will show, across multiple model sizes (110M, 270M, 750M, and 1B parameters), PiKE consistently outperforms all existing methods, including heuristic or static data mixing, multitask optimization, and adaptive data mixture approaches. We first start by explaining our setup.
4.1 Experiment Setup
Baselines: We evaluate range of sampling strategies: (1) (Uniform) Mix, (2) Round-Robin (3) Random, (4) FAMO (38) , (5) ADO (30), (6) GLaM (18), (7) DoReMi 69 (8) PiKE, and (9) Balanced- PiKE. Among these, DoReMi, GLaM. and ADO are specifically designed for LLM pretraining. We also include FAMO. a recent multitask learning (MTL) method, because it is scalable to large model sizes and has demonstrated competitive crrformancce-sevieee work. DoReMi estimates task weights by training small proxy model, while GLaM assigns static weights based on downstream performance from smaller models. However, since both methods report weights only for the GLaM dataset and do not provide configurations for multilingual C4, we exclude them from our multilingual experiments. dynamically updates task sampling weights during training using gradient information, enabling adaptive optimization throughout training. PiKE scalable to both large models and many tasks, with minimal overhead, as detailed in Appendix F4
as useful comparison point for our
contrast to these static methods, PiKE
8

---

Datasets: For multilingual experiments, we use mC4 77]. focusing on English (en), Hindi (hi), and German (de). An overview of these datasets provided in Table For GLaM-based experiments, we use the six-domain GLaM dataset [8] Additional details are presented in Table 4
Perplexity is measured on held-out validation data. Downstream evaluation follows the OLMES suite [26]. For multilingual downstream tasks, we use multilingual HellaSwag [33, covering 26 languages. For models trained on GLaM, we evaluate on downstream tasks ARC-Easy [12). CommonsenseQA [59], PIQA [S], and HellaSwag [76]. HellaSwag and ArcE tasks have 4 choices, CSQA has 5 choices, and PIQA has 2 choices.
Evaluation:
Further details on our experimental setup and evaluation are in Appendix E
4.2 Main Observations on Multilingual Pretraining Experiments
49 on English, Hindi, and German. Mix batching outperforms Round-Robi and Random strategies, supporting our analysis in Section and justifying our focus on Mix as the foundation for PiKE. Additional results with different model sizes (270M and 1B) and language settings are reported in Table [5] Our main observation is that PiKE and its Balanced variant achieve the highest average downstream accuracy across all language settings and model sizes, demonstrating their effectiveness for multilingual pretraining,
Table I presents results for pretraining a IB multilingual GPT-2 model
We also observe that Balanced- PiKKE promotes more fair learning across tasks, In particular, we {1, 5}. As
pre-trained IB models using Balanced- with different values of parameter T increases, task losses become more uniform, reflecting improved balanced learning. T 5, perplexity becomes more balanced across languages, while 3 offers the best trade-off both the lowest perplexity and highest downstream accuracy. These results highlight the importance of incorporating fairness/halanced learning into data mixing strategies during pretraining.
4.3 Main Observations on Pretraining Experiments with GLaM Datasets
Table results for pretraining 750M GPT-22 model the GLaM dataset. Additional results with both 110M and 750M models across six domains are provided in Table Across both 110M and 750M model sizes, PiKE consistently outperforms DoReMi, GLaM, and Mix in downstream accuracy, For the 750M model, PiKE improves average accuracy by 3.4% over DoReMi, 6.2% % over GLaM, 7.1% over FAMO, and 4.8% over ADO. In the 110M setting, PiKE achieves 37.8% accuracy, exceeding DoReMi (36.0% )), GLaM (35.3%) )), and FAMO (35.9 %). Unlike DoReMi, which requires a separate proxy model, or GLaM, which depends on tuning with smaller models, PiKE delivers these gains with negligible additional overhead.
PiKE benefus from apriori downstream-tured weights. We evaluate PiKE with two initializations: (1) uniform weights by b/ K and (2) GLaM-tuned weights. In both small and large GPT-2 configurations, PiKE benefits from utilizing already fine tuned weights initialization, achieving 48.1% % accuracy with GLaM-tuned weights vs. 47.6% with uniform initialization. This shows that PiKE can effectively leverage pre-existing fine-tuned weights while still outperforming other methods with uniform initialization.
Mixing datasets improves language model generalization. We compare models trained on individual domains to those trained on -oxxed-ommin datasets. Table [6 shows that single-domain training underperforms compared
training, even with simple Mix sampling. This reinforces the importance of diverse data for pretraining and aligns with prior work [40]
Perplexity versus downstream performance. Table 16 reveals that validation perplexity does not always align with downstream performance. For instance, while Mix sampling yields lower perplexity in 750M models, PiKE achieves better downstream accuracy. This aligns with prior findings 23 60 67], suggesting that perplexity alone is not a reliable performance metric.
5 Related Work
Data Curation and Selection. The effectiveness of language models heavily depends on the quality of the pre-training corpus. Consequently, significant efforts have been made to enhance pre-training data.
These efforts include heuristic-based filtering [51 50 32 48 57) and deduplication [] 33 10 19].

---

Recently, 664 proposed an automated method for constructing large, diverse, and balanced datasets
for superviee learning by applying hierarchical k-means clustering. [55] introduced techniques that leverage instruction-ture -unned models to assess and select high-quality training examples, along with density sampling to ensure diverse data coverage by modeling the data distribution. Additionally, [27] simulated training runs to model the non-additive effects of individual training examples, enabling the analysis of their influence on a model's predictions.
Multitask Learning Optimization Multitask learning (MTL) optimization is closely related to our
framework. Among different MTL approaches, one idea is modify gradient updates to mitigate gradient conflicts situatiins where task gradients point in opposing directions, slowing down optimization [63 [75]. The Multiple Gradient Descent Algorithm (MGDA) [16] [56] updates the model by optimizing the worst improvement across all tasks, aiming for equal descent in task losses. Projected Conflicting Gradient Descent (PCGrad) 75 modifies task gradients by iteratively removing conflicting components in randomized order, ensuring that updates do not interfere destructively across tasks. Conflict-Averse Gradient Descent (CAGRAD) (37] optimizes for the worst task improvement while ensuring a decrease in the average loss. NASHMTL 47 determines gradient directions by solving a bargaining game that maximizes the sum of log utility functions. While these methods improve performance, they introduce computational and memory overhead, making them impractical for large-scale models with numerous tasks [70]. Similar challenges exist in AdaTask (73), which improves multitask learning by balancing parameter updates using tusk-wise adaptive learning rates, mitigating task dominance, and enhancing overall performance. Unlike previous approaches that requires (K) storage for task gradients (e.g. PCGrad) or optimizer states (e.g. AdaTask),, FAMO [38] balances task loss reductions efficiently using (1)) space and time. However, these methods fail exploit the non-conflicting interactions among tasks, focusing instead on resolving conflicts that seldom arise. This motivates our work that actively leverages lack of gradient conflicts to enhance training efficiency.
Data Mixture Problem Another closely related line of work to our paper is data mixture problem. Early works relied on manual heuristics for forming data compositions to construct datasets [18] 71 22
for training LLMs. To create more principled and less labor- intensive methods, subsequent research has focused optimizing these mixtures. One line of work uses offline optimization with proxy models, where a smaller, less expensive model is trained to find optimal data weights that are then used for the final large model training. Methods in this area include DoReMi (6), which optimizes domain weights to accelerate training convergence, and DOGE (20), which re- weights data by relying on the inner product between task gradients. A third category involves adaptive online data mixing, where the data distribution is adjusted dynamically during the training process. This includes methods that efficiently adjust sampling rates based on training signals ADO 30) and skill-based curriculum learning like Skill-it! 81 which dynamically samples data corresponding to specific skills to improve model proficiency. Our work PiKE is also in this category. In particular, PiKE dynamically leverages the near-orthee ortoonnlty of task gradients, making it far more computational and memory -fffcienn than offline approaches like DoGE by only needing two scalars per task: per-task gradient norm and variance. PiKE also uniquely incorporates gradient variance or noise, fundamental factor impacting the performance of the training, which was largely overlooked in recent data mixing literature such as DoGE. Overall, PiKE offers a reliable way to enhance final model performance.
6 Conclusion
In this work, we introduced PiKE, an adaptive data mixing algorithm for multitask learning that dynamically adjusts task sampling based observed gradient interactions. Unlike prior methods that aim to resolve gradient conflicts, PiKE exploits the predominantly non-confficting gradients seen in large-scale language model pretraining. We provided theoretical analysis and showed, through extensive experiments, that PiKE improves both convergence speed and downstream performance across multilingual and multi-domain settings. To promote balanced task learning, we extended PiKE with a fairness-aware objective, resulting in Balanced-PiKAe which reduces task-level performance gaps without sacrificing overall accuracy.
One limitation of PiKE is its lack of sensitivity to dataset size when assigning sampling weights. Future work could incorporate data abundance or downstream performance feedback into the sampling strategy. Additionally, extending PiKE to other domains beyond language modeling remains a promising direction for further research.
10

---

References
[1] A. Abbas, K. Tirumala, D. Simig, S. Ganguli, and A. S. Morcos. Semdedup: Data-efficient learning at web-scale through semantic deduplication. arXiv preprint arXiv:3303.09540, 2023.
[2] J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization, arXiv preprint arXiv:1607.06453 2016.
[3] H. Ban and K. Fair resource allocation in multi-task learning. arXiv:2402. 15638, 2024.
arXiv preprint
[4] Y. Bengio. J. Louradour, R. Collobert, and J. Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pages 41-48, 2009.
[5] Y. Bisk, R. Zellers, R. Le Bras, Gao, and Y. Choi. Reasoning about physical commonsense in natural language, 2019.
[6] J. Bradbury, R. Frostig, P. Hawkins, M. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman- and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL pettp://githnn com/ com/ google/jax
[7] Brown, B. Mann, Ryder, M. Subbiah, J. Kaplan, Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[8] M. F. Chen, N. Roberts, K. Bhatia, J. WANG, C. Zhang, F. Sala, and C. Re. Skill-it! a data- driven skills framework for understanding and training language models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openrrviies
[9] Z. Chen, V. Badrinarayanan, C.-Y. Lee, and A. Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing deep multitask networks. International conference on machine learning, pages 794-803. PMLR, 2018.
[10] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[11] A. Chowdhery, S. Narang. J. Devlin, M. Bosma, G. Mishra, A. Roberts, Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
[12] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457 2018.
[13] V. Dac Lai, C. Van Nguyen, Ngo, Nguyen, Dernoncourt, R. A. Rossi, and H. Nguyen. Okapi: Instruction-tuned large language models in multiple languages with reinforce- ment learning from human feedback. arXiv e-prints, pages arXiv-2307, 2023.
[14] J. Dai. K. He, and J. Sun. Instance-aware semantic segmentation via multi-task network cascades. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 3150-3158, 2016.
[15] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron, R. Geirhos, I, Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. arXiv preprint arXiv:2302.00544, 2023.
[16] J.-A. Désidéri. Multiple-gradien descent algorithm (mgda) for multiobjective optimization. Comptes Rendus Mathematique, 350(5-6):313-3) 2012.
[17] Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1:Xiv:1110.04805 2018

---

[18] N. Du, Y. Huang, A. M. Dai, S. Tong, D. Lepikhin, Y. Xu, M. Krikun, Y. Zhou, A. W. Yu, O. Firat, al. Glam: Efficient scaling of language models with mixture-of-ene -eeeetess.e..s....e.eeeee In International Conference on Machine Learning, pages 5547-5569. PMLR, 2022.
[19] A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Yang, A. Fan, et al. The llama 3 herd of models. arXiv preprini arXiv:2407.21788 2024.
[20] S. Fan, M. Pagliardini, and M. Jaggi Doge: Domain reweighting with generalization estimation. arXiv preprint arXiv:2310.15309 2023.
[21] C. Gafffny, D. Li, R. Sang, A. Jain, and H. Hu. Orbax, 2023. URL http::/github. com/
google/orbax
[22] L. Gaa, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, al. The pile: An 800gb dataset of diverse text for language modeling, arXiv preprint arXiv:2101.00027, 2020.
[23] T. Gao, A. Wettig, L.. He, Y. Dong, S. Malladi, and D. Chen. Metadata conditioning accelerates language model pre-training. arXiv preprint arXiv:2501.01956 2025.
[24] C. Ge, Z. Ma, D. Chen, Y. Li, and B. Ding. Data mixing made efficient: A bivariate scaling law for language model pretraining. arXiv preprint wxiv:2445..49008 2024.
[25] Google. Grain - feeding jax models, 2023. URL http::/giituubb... www.auubb..ooooooole//.oooo /grain
[26] Y. Gu, O. Tafjord, B. Kuehl, D. Haddad, J. Dodge, and H. Hajishirzi. Olmes: A standard for language model evaluations. arXiv preprint 240644640, 2024.
[27] K. Guu, A. Webson, E. Pavlick, L. Dixon, I. Tenney, and T. Bolukbasi. Simfluence: Modeling the influence of individual training examples by simulating training runs. arXiv preprint arXiv:2303.08114, 2023.
[28] Heek, A. Levskaya, A. Oliver, M. Ritter, B. Rondepierre, A. Steiner, and M. van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL http:://///t com/ comp/oogle///// flax
[29] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, Welbl, A. Clark, et al. An empirical analysis of compute-optimal large language model training. Advances in Neural Information Processing Systems, 35:30016-30030, 2022.
[30] Y. Jiang, A. Zhou, Z. Feng, S. Malladi, and Z. Kolter. Adaptive data optimization: Dynamic sample selection with scaling laws. arXiv preprint arXiv:2410.1 /820, 2024.
[31] N. P. Jouppi, C. Young, N. Patil, D. Patterson, G. Agrawal, R. Bajwa, S. Bates, S. Bhatia, N. Boden, A. Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
[32] H. Laurençon, Saulnier, T. Wang, C. Akiki, A. Villanova del Moral, T. Le Scao, Von Werra, C. Mou, E. González Ponferrada, H. Nguyen, et al. The bigscience roots corpus: A 1.6 tb composite multilingual dataset. Advances Neural Information Processing Systems, 35: 31 31809-31822 2022.
[33] K. Lee, D. Ippolito, Nystrom, C. Zhang, D. Eck, C. Callison-Burchh and N. Carlini. Dedupli- cating training data makes language models better. arXiv preprint arXiv:2107.06499 2021.
[34] C. Li, H. Farkhoor, R. Liu, and J. Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations (ICLR), 2018. https //arxiv. org/abs//////oooooeoggggoteaeere 1804. 08838
[35] Li, A. Beirami, M. Sanjabi, and V. Smith. Tilted empirical risk minimization. arXiv preprint arXiv:2007.01162 2020.
12

---

[36] A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024.
[37] B. Liu, X. Liu, X. Jin, P. Stone, and Q. Liu, Conflict-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems, 18888-8888800 2021.
[38] B. Liu, Y. Feng, Stone, and Q. Liu. Famo: Fast adaptive multitask optimization. Advances in Neural Information Processing Systems, 36, 2024.
[39] H. Liu, M. Xie, Z. Li, and T. Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning, pages 22188-22214. PMLR, 2023.
[40] Q. Liu, X. Zheng. N. Muennighoff, G. Zeng, L. Dou, T. Pang, J. Jiang, and M. Lin. Regmix: Data mixture regression for language model pre-training. arXiv preprint arXiv:2407.01492 2024.
[41] X. Liu, J. Gao, X. He, L. Deng, K. Duh, and Y.-Y. Wang. Representation learning using multi- task deep neural networks for semantic classification and information retrieval. In Association for Computational Linguistics, 2015.
[42] X. Liu, P. He, W. Chen, and J. Gao. Multi-task deep neural networks for natural language understanding, arXiv preprint arXiv:1901. 10::::9911111000 11504, 2019.
[43] Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Con- ference on Learning Representations (ICLR), 2019. http:////ppeereeeeeeeeeee/ooo/oe Bkg6RicqY7
[44] M.-T. Luong, Q. V. Le, I. Sutskever, O. Vinyals, and L. Kaiser. Multi-task sequence to sequence learning, arXiv preprint arXiv:1511.066 14, 2015.
[45] I. Misra, A. Shrivastava, A. Gupta, and M. Hebert. Cross-stitch networks for multi-task learning. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 3994 4003, 2016.
[46] J. Mo and J. Walrand. Fair end-to-end window- congestion control. /EEE/ACM Transac- tions on networking, 8(5):556-566 567, 2000.
[47] A. Navon, A. Shamsian. I. Achituve, H. Maron, K. Kawaguchi, G. Chechik, and E. Fetaya. Multi-task learning as bargaining game. arXiv preprint arXiv:2202.01017. 2022.
[48] G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, H. Alobeidli, A. Cappelli, B. Pannier, E. Almazrouei, and Launay. The refinedweb dataset for falcon Ilm: Outperforming curated corpora with web data only. Advances in Neural Information Processing Systems, 36:79155- 79172, 2023.
[49] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and Sutskever. Language Models are Unsupervised Multitask Learners, 2019. Sstps://openate wwweenai.combbrie com/blog/ Matter-language-monn
[50] J. W. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoffmann, F. Song, J. Aslanides, S. Henderson, R. Ring. S. Young, al. Scaling language models: Methods, analysis insights from training gopher. arXiv preprint arXiv:21 12. /446, 2021.
[51] C. Raffel. N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Zhou, W. Li, and J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
[52] J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He. (ZeRO-Offload): Democratizing (Billion-Seale) model training. In 202/ USENIX Annual Technical Conference (USENIX ATC 2/), pages 551-564, 2021.
[53] R. T. Rockafellar. Convex 28)))). Princeton university press, 2015.
13

---

[54] S. RRdder, J. Bingel, I. Augenstein, and A. Søgaard, Latent multi-task architecture learning. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 4822 4829, 2019.
[55] N. Sachdeva, B. Coleman, W.-C. Kang, J. Ni, L. Hong, E. H. Chi, J. Caverlee, J. McAuley, and D. Z. Cheng. How train data-efficieed llms. arXiv preprint arXiv:2402.00666 2024.
[56] O. Sener and V. Koltun. Multi-task learning as multi-objective optimization. Advances in neural information processing systems, 31, 2018.
[57] L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, J. Dumas, Y. Elazar, al. Dolma: An open corpus of three trillion tokens for language model pretraining research. arXiv preprint wwXiv:2402.00/52 2024.
[58] J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu, Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864 2021.
[59] A. Talmor, J. Herzig, N. Lourie, and J. Berant. Commonsensequ: A question answering challenge targeting commonsense knowledge. arXiv preprint arXiv:1811.00937 2018.
[60] Y. Tay, M. Dehghani, J. Rao, W. Fedus, S. Abnar, H. W. Chung, S. Narang, D. Yogatama, A. Vaswani, and D. Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers, arXiv preprint (idi::2009.06688 2021.
[61] G. Team, R. Anil, S. Borgeaud, J.-B. Alayrac, Yu, R. Soricut, Schalkwyk, A. M. Dai, A. Hauth, K. Millican, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. /1800, 2023.
[62] G. Team, Georgiev, I. Lei, R. Burnell, L. Bai, A. Gulati, G. Tanzer, D. Vincent, Z. Pan, S. Wang, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint vrXiv:2403.055300 2024.
[63] S. Vandenhende, Georgoulis, W. Van Gansbeke, M. Proesmans, D. Dai, and L. Van Gool. Multi-task learning for dense prediction tasks: A survey. IEEE transactions on pattern analysis and machine intelligence, 44(7):3614 4444443333220 2021.
[64] H. V. Vo. V. Khalidov, T. Darcet, T. Moutakanni, N. Smetanin. M. Szafraniec, H. Touvron, C. Couprie, M. Oquab, A. Joulin, et al. Automatic data curation for self self-supervised learning: A clustering- based approach. arXiv preprint 2rXiv:2405.15611 2024.
[65] A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, Michael, F. Hill, O. Levy, and S. Bow- man. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
[66] Z. Wang, Y. Tsvetkov, O. Firat, and Y. Cao. Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models. arXiv preprint arXiv:2010.05874, 2020.
[67] A. Wettig, A. Gupta, S. Malik, and D. Chen. Qurating: Selecting high-quality data for training language models. arXiv preprint arXiv:2402.097.3) 2024.
[68] M. Wortsman, P. J. Liu, L. Xiao, K. Everett, A. Alemi, B. Adlam, J. D. Co-Reyes, I. Gur, A. Kumar, R. Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprini arXiv:2309. /4322, 2023.
[69] S. M. Xie, H. Pham, X. Dong, N. Du, H. Liu, Y. Lu, S. Liang, Q. Le, T. Ma, and A. W. Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2024.
[70] D. Xin, B. Ghorbani, J. Gilmer, A. Garg, and O. Firat. Do current multi-task optimization methods in deep learning even help? Advances in neural information processing systems, 35: 13597-13609, 2022.
[71] I. Xue. mt5: A massively multilingual pre-trained text-to-text transformer. arXiv preprint arX1v:2010.11934, 2020.
14

---

[72] L. Xue, A. Barua, N. Constant, R. Al-Rfou, S. Narang, M. Kale. A. Roberts, and C. Raffel. Byt5: Towards a token-free free future with pre-trained byte-to-bye bott models. corr, abs/2105.1362) arXiv preprini arXiv: 105. 1 13626, 2021.
[73] E. Yang, J. Pan, X. Wang, H. Yu, L. Shen, X. Chen, L. Xiao, J. Jiang, and G. Guo. Adatask: A lask-aware adaptive learning rate approach to multi-task learning. Proceedings of the AAAI conference artificial intelligence, volume 37, pages 10745-10753 2023.
[74] J. Ye, P. Liu, T. Sun, Y. Zhou, J. Zhan, and X. Qiu. Data mixing laws: Optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403 Ttt::2400..695555 16952, 2024.
[75] T. Yu, S. Kumar, A. Gupta, S. Levine, K. Hausman, and C. Finn. Gradient surgery for multi-task learning, Advances in Neural Information Processing Systems, 33:5824-5836, 2020.
[76] R. Zellers, Holtzman, Y. Bisk, Farhadi, and Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: /90. 905300 2019.
[77] Y. Zhang, C. Chen, T. Ding, Z. Li, R. Sun, and Z.-Q. Luo. Why transformers need adam: A hessian perspective. arXiv preprint 4rXiv:2402.16778 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the claims made in the abstract and introduction accurately reflect the paper's contributions and scope. Each claim is clearly supported in the main text, with appropriate pointers provided for readers to easily verify and locate the corresponding results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Yes discuss the limitations of the work in the conclusion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification:
Yes provide the full set of assumptions and complete and correct proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
Yes disclose all the information needed reproduce the main results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
17

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification: We do not include the full code for pre-training language models because tailored to the specific framework and system we use. However, our method is general, and a professional programmer should able to reproduce it within one to two days using the detailed implementation description provided in the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niis. cc/
sublic/guidess
CoodeSubmissionooll1 hataaigggiiies/oeeeuubm.msee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
cilips.cc/publle httee//ullccgguuddeeeeeeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we have specify all the training and test details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due to the high computational cost of each experiment, which takes approxi- mately 1-2 days to run, we did not perform multiple runs and therefore do not report error bars other statistical significance measures.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Yes we report the information on the computer resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttp::///euriiii wttpuurriisccppppupeeeeee
Answer: [Yes]
Justification:
Yes we strictly follow the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes our work should only have positive societal impacts and no negative societal impacts. There are many potential positive societal consequences of our work, none of which we feel must be specifically highlighted here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
19

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper doess not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy pettps://neurr CcCC/oofeeeeecccs///2////// for what should or should be described.
22

---

A The Scalability Problem of Existing MTL Methods
73 require computing and storing all task-specific gradients at each training iteration, where K is the number tasks. While exceptions like FAMO [38] are designed for efficient scaling, the aforementioned general approach typically leads to K d) space complexity for storing gradients and O(Kd) time complexity per iteration for their computation, where d denotes the number of model parameters. This cost is further exacerbated when these methods also involve solving auxiliary optimization problems to combine or re-weight gradients [70]. In stark contrast, when the overall loss simple average of task sse-sspccfiic K losses (e.g., C = ズ Ei=1 Lx), its gradient V can be computed via single backward pass. This results in ((d) space and time complexity for the gradient computation, with the aask-aggregation step being effectively independent of K.
Many existing multitask learning (MTL) methods 16 56 66 75
While early MTL research often focused on relatively small-scale vision models or language models with fewer than 100 million parameters, modern language models frequently exceed 100 billion parameters (100B). Storing the gradients or weights for such a model (assuming 32-bit Hloating-poin precision) requires approximately 400 GB of memory. This capacity already surpasses that of single high-performance GPU like the NVIDIA H100 (80 GB), necessitating at least five such GPUs for storing just one set of gradients or weights.
Applying typical prior MTL methods in this large-scale context becomes impractical. Storing K full gradients for 100B parameter model would demand approximately 400K GB of memory, equivalent to 5K H100 GPUs. This scaling of memory and computational requirements-4 least linear with the number of tasks K for gradient operations and often compounded by super-linear costs from auxiliary optimization steps-render naive applications of many existing MTL methods infeasible for current large language models. Indeed, even if storing per-task gradients were feasible, Xin al. [70] reported significant reduction in training throughput (steps per second) for such MTL techniques even on small-scale vision models, let alone for modern LLMs (e.g.. those with 1B+ parameters). Collectively, these challenges highlight the pressing need for novel MTL methodologies engineered for large-scale pre-training demanding both exceptional scalability and minimal additional computational and memory overhead.
B Discussion of Gradient Conflicts
GradVaccine (66 made a similar observation regarding low gradient conflicts among task gradients in multilingual BERT (178M parameters). We extend this finding to substantially larger autoregressive, decoder-only models (up to IB parameters), which are more characteristic of current large-scale language modeling paradigms.
Our experiments demonstrate that task gradients in large-scale models indeed exhibit minimal conflicts. To illustrate, we conduct two pre-training experiments: (i) a 1B parameter GPT-2-style model 49 on the multilingual mC4 dataset [7]] (covering six languages: English, Hindi, German, Chinese, French, and Arabic), and (ii) 750M parameter model on The GLaM dataset [8] (English text from six diverse domains). Experimental details are provided in Appendix E Figures and depict cosine similarity trends for task gradients, revealing several key observations: 1) Gradient similarity is initially high but generally decreases as training progresses. 2) In multilingual settings, gradient similarity correlates with linguistic proximity (e.g., English- closely), whereas gradients from GLaM's 's diverse domains exhibit more uniform positive alignment. 3) Task gradients rarely monfict-multiliggeee cosine similarities seldom drop below 10.1, and GLaM domain gradients remain predominantly positive.
gradients align more
C PiKE: Conceptual Version
Here, we present the conceptual (basic) version of PiKE. As discussed in the main text, this approach lacks computational efficiency due to the frequent estimation of the norm and the variance of the per-task gradient.
23

---

Algorithm 2 Conceptual version of PiKE: Positive gradient Interaction-based K-task weights Estima-
tor
1: Input: 0, total batch size b, stepsize y, task & dataset Dg. constants B3, I, and prior weights w/
2: Initialize: we 1/K / wk +- w vk
3: for = 0, 1, 1.... do
4: Estimate (0.)ll and of for every
5: Compute A - Vz
Lnd of and K3 L 2b
6: set W = max (0, - where μr is found (by bisection) such that Ek=1 W = Set (bi,..., bx) +- (owi... wk))
7:
8: Sample by data points from each task k
9: Compute the gradient g using the estimates samples
Update: 0c+1 < Optimizer(1), 0c g)
11: end for
As discussed in section 3.2 this algorithm is computationally inefficient it requires estimating (l,( and dn at each iteration. To improve efficiency, we introduced modifications that led to the development of the PiKE algorithm (Algorithm) I in the main body).
D Balanced-PiKE: Fairness Considerations Across Tasks
Here, we present the Balanced-Pite algorithm in more detail. As discussed in the main body, the main difference with PiKE is that the fair version requires the computation of the coefficients
TE (1)
k== e Vk
Then updating the sampling weights by
Wk wy exp
(y
Vk
2b
The overall algorithm summarized in Algorithm For our experiments, we evaluate three different values of r: 1, 3, and 5. A larger T results in a stronger balancing effect between different tasks.
Algorithm Balanced- Pikkn: Balanced considerations across tasks
1: Input: 0, To. total batch size b, task & dataset DR. hyperparameters (2, prior weights w/
2: Initialize: Wi f 1/1 or wk f- w/
3: for t. = 0, 1, ... do
4: if mod To = 0 then
5: Estimate L. of, and (0:) for every
6:
Vi Σn 12
7:
<- wk exp (
(vt) 2
8:
w + w w//w
9:
(bi..... .b) + round (bwnddb(ww.......... ..,.... wx))
10:
end if
Sample by data points from each task k
12: Compute the gradient using the estimates samples
Update: 0.+1 <- Optimizer(1), A g)
14: end for
E Experiments Setup
E. Dataset Details
Our experiments construct two primary scenarios for multitask learning: multilingual tasks and diverse task mixtures spanning multiple domains. We consider two widely-used datasets for our study: mC4 [71] and GLaM [18].
24

---

mC Dataset The mC4 dataset [71] is multilingual text corpus derived from the Common Crawl web archive, covering diverse range of languages. It has been widely used for pretraining multilingual models, such as m'T5 71 and ByT5 72 'he dataset is curated by applying language-speciln filtering to extract high-quality text, ensuring a balanced representation across languages. Mixture weights for training models on mC4 are often chosen based on token counts. In our cases, we mainly focus on English (en), Hindi (hi), and German (de). We report their details in Table 3
Table 3: Partial statisties of the mC4 corpus, totaling 6.3T tokens.
ISO code Language
Tokens (B)
English Hindi
2,733
hi
24
de
German
347
GLaM Dataset The GLaM dataset [8] comprises English text from six distinct sources and has been used to train the GLaM series models and PaLM IT. Mixture weights for GLaM training were determined based on small model performance 18 while [69] employed group distributionally robust optimization (Group DRO) to compute domain-specific weights. Table summarizes the six domains in the GLaM dataset and the mixture weights selected by GLaM and DoReMi. We use these weights as oracle baselines for comparison with PiKE, which dynamically adjusts task weights over time using gradient information, unlike the fixed weights employed by GLaM and DoReMi.
Table 4: GLaM dataset
g
and fixed mixture weights used in GLaM 18 and DoReMi 66
Dataset
Tokens (B)
Weight chosen by GLaM [8]
Weight chosen by DoReMi (69
Filtered Webpages
143
0.42
0.51
Wikipedia
3
0.06
0,05
Conversations
174
0.28
0.22
Forums
247
0,02
0.04
Books
390
0,20
0.20
News
650
0.02
0.02
E.2 Training Details
Our experiments explore two distinct scenarios for multitask learning: multilingual training and diverse task mixtures spanning multiple domains. To achieve optimal results, customize the training setups for each scenario and present them separately in this section. All training is performed from seratch.
Multilingual Training address complexities of tokenizing multilingual data, we utilize mT5 tokenizer [7I], which features vocabulary size of 250K. Both GPT-2 small and GPT-2 large models are trained with context length of 1024 and batch size 256. The AdamW optimizer [43] is employed with consistent hyperparameters and learning rate scheduler. Additional details hyperparameter configurations are provided in Appendix E.5
GLaM Training For GLaM training, we use the T5 tokenizer 51 implemented as SentencePiece
tokenizer trained on the C4 dataset with vocabulary size of 32,000. Both GPT-2 small and GPT. 2 large models are trained with a context length of 1024 and batch size of 256. The AdamW optimizer 43 is used, and additional details on hyperparameters is in Appendix E.5
E.3 Model Architecture
The detailed architecture is summarized Table 5
Our implementation utilizes pre-
normalization 49 Transformers with qk-layernorm [15]. Consistent with [10), we omit biases, and the layernorm 2 value remains set the Flax 28 default of 1e-6. Additionally, we incorporate rotary positional embeddings (58).
25

---

Table Architecture hyperparameters for different model scales used in the paper. All models are GPT-2-like decoder-only architectures. The multilingual models employ a vocabulary size of 250K. whereas GLaMM training uses vocabulary size of 32K. Differences the total number parameters arise due to the variation in vocabulary sizes.
Size
# Params
Layers
Attention heads
Attention head dim
Hidden dim
GPT-2 small
110M/270M
12
12
64
768
GPT-2 large
750M/IB
36
20
64
1280
Table Hyperparameter settings for our experiments.
Hyperparameters
Values
Optimizer
AdamW (ßi 0.95, B2 = 0.98)
Initial and final learning rate
7e 6
Peak learning rate
76 4
Weight decay
0.1
Batch size
256
Context length
1024
Gradient clipping norm
1.0
Training step
120, 000
Warm-up step
10, 000
Schedule
Lincar decay to final learning rate
Table 7: Hyperparameter settings for running PiKE (Algorithm)
Hyperparameters
Values
PiKE hyperparameter (S1,Ca)
{(le
2, le - 31. le 3),(5e (5e 2,5e 3),
(7.5c
2, 5e 3), (1e 1, 2), (1.5c 1, 2))
Check interval To Batch size for estimation
1,000
256
E.4
Experimental Resource
All experiments are conducted on 8 Google TPUv5. The training time for GPT-2 small and GPT-2 large models for 120K steps approximately day and 2 days per run, respectively.
E.5
Hyper-parameters
Table showws the detailed hyperparameters that we used in all our experiments. We also report our hyperparameters grid for tuning PiKE in Table 7
E.6
Implementation Details
Our implementation builds upon the Nanodo training infrastructure 68.
incorporating enhancements
for efficiency. This framework relies on Flax (28). JAX (6), and TPUs 31.
To enable training of larger models, we shard both model and optimizer states, following the methodology of FSDP [52], and define these shardings during JIT compilation. Checkpointing is
handled using Orbax (21. while deterministic data loading is facilitated by Grain [25].
For data loading, sequences are packed to avoid padding. When sequence contains fewer tokens than the context length hyperparameter, an end-of-sequence token is appended. This differs from Nanodo (68) where both Degin-off-squence and end-of-sequence tokens are added.
E.7 Evaluation
For multilingual downstream performance, utilize the multilingual version HellaSwag 13 which supports evaluations across 26 languages. English downstream tasks are assessed using ARC-Easy [2], Commonsense [59]. PIQA (5),
Our evaluation adheres to the OLMES suite 26
26

---

and HellaSwag [76]. Unless specified otherwise, multilingual evaluations are performed in a O-shot setting, while GLaM pretraining evaluations employ 7-shot in-context learning, with demonstration candidates separated by two line breaks. For HellaSwag and its translated variants, we evaluate the first 3,000 examples. For all other downstream tasks, evaluations are conducted on their respective validation sets. In the case of multiple-choie tasks, different candidates are included in the prompt, and the average log-likelihood for each candidate is computed. The candidate with the highest score is then selected as the predicted answer.
F
Additional Experiment Results
F11
Comparison of Performance Using Mix, Random, and Round-Robin Sampling Strategies
Figure presents the average downstream accuracies language models pre-trained using Mix, Random, and Round-Robin sampling strategies. In both multilingual pre-training and GLaM pre- training, the Mix sampling strategy consistently outperforms the other two. This motivates us its use in pre-training large language models.
Mis
Mis
45
27
Barsias
Rundics
Accurscy (%)
Avcuracy (%)
40
26
35
25
30
0
20
40
60
80
100
120
0
20
40
60
80
100
120
#
Training Steps (K)
Training Steps (K)
(a): 1B models on multilingual C4 (en), C4 (hi), (b): 750M models on GLaM datasets with six domains
and C4 (de) datasets
Figure 4: Average downstream task accuracy of pretraining language models using Mix, Round- Robin, and Random sampling strategies. Mix and Random use equal batch size for each task (b& 6/K, vk K).
F.2 Cosine Similarity and c-Confficted Gradients
and show the cosine similarity, defined as (C) and the "ratio," defined as
Figures
(0) I12,. ( In particular, if (2,01 (L)
č then the gradients are g-conflicted for
c 8/2, which aligns with the observations in Figures and 6
0.3
0.3
Genan--irgla
Englab-tited
Cosine Similarity
0.2
0.2
Chirmse Peeneh
Ratio
0.0
0.0
-0.1
-0.1
100 120
100 120
Training Steps (K)
Training Steps (K)
Figure 5: IB models trained multilingual mC4 datasets. Left: Cosine similarity between task gradients during language model pre-training over time. Right: The "ratio," which defined as (C,(0).C,()))
[c,(0) between task gradients during language model pre-training over time. "datal- data2" denotes the cosine similarity ratio between the gradient of data/ and the gradient of data2.
27

---

0.3
0.3
booke- foruue- booke-wweb
foruns-boo books
Cosine Similarity
web books
0.2
web
0.2
wilki-wee
Ratio
0.1
0.1
0.0
0.0
100 120
20
100 120
# Training Steps (K)
# Training Steps (K)
Figure 750M models on GLaM datasets with six domains. Left: Cosine similarity between task gradients during language model pre-training over time. Right: The "ratio," which defined as (z()))
time. "datal- data2" denotes the cosine similarity or ratio between the gradient datal and the gradient of data2.
[,) [[,,,,+[1000 (o)ls between task gradients during language model pre-training
F.3
Comparison of Performance Using PCGrad, AdaTask, and Mix
Figure 7 presents the average downstream task performance on HellaSwag (en) and HellaSwag (hi) for 270M multilingual language models pre-trained using PCGrad, AdaTask, and Mix. As shown in Figure 7 1) PCGrad performs similarly to Mix, as only adjusts gradients when conflicts occur which is rare. 2) AdaTask converges more slowly due to noisy gradients and suboptimal optimizer state updates. Additionally, both methods are memory-intensive requiring O(K) storage for task gradients (PCGrad) or optimizer states (AdaTask), making them impractical for large-scale models such as the 540B PaLM [00
Prrra
AdoTiak
Perplexity
20
40
60
80
100
120
Training Steps (K)
Figure 7: Eval perplexity of pretraining 270M GPT-2 style multilingual language models on mC4 datasets (English and Hindi) using Mix, PCGrad, and AdaTask.
F.4
PiKE: Minimal Training and Memory Overhead
Table & reports the total training time and the computational overhead incurred by PiKE. our experiments, per-task gradient squared norms and variances are estimated only once every To 1, 000 training steps. Profiling results indicate that this infrequent estimation introduces minimal overhead during the pre-training of LLMs. Specifically, for a IB parameter multilingual model pre-trained for approximately 71 hours, this estimation process accounted for only about 2.3% of the total training time (roughly 1.6 hours). This overhead percentage can be further reduced by using larger estimation interval, To (e.g.. 5,000 steps), or by adopting a more efficient parallelization implementation for handling multiple tasks.
Regarding memory requirements, PiKE only needs store the periodically estimated scalar values for each task's gradient variance and squared norm. This amounts an additional memory footprint of (()). which is negligible compared to the model parameters (Odd) and especially to the O K d) memory demanded by prior MTL methods that store K full per-task gradients. This efficiency starkly
28

---

contrasts with those earlier approaches, which also often involve solving auxiliary optimization problems that further increase computational and memory burdens [70]. Unlike approaches such
as DoReMi or GLaM, PiKE simplifies deployment due to its automated learning of task sampling weights during the training procedure. This dynamic weight adaptation facilitates straightforward implementation and seamless integration into existing workflows, eliminating the need for manual intervention or training disruptions.
Table 8: We report the total training time (hrs) and how much overhead time that running PiKE (hrs). Compared with the total training, the overhead for running PiKE is minimal only taking from 1.2% 2.4% training large models.
Training thrs)
Overhend running (hrs)
Overhead running (%)
GPT-2 small (I 10M, GlaM)
10
0.12
1.2
GPT-2 small (270M, Multllnguul)
12
0.24
2
GPT-2 large (750ML GLaM)
41
1.0
2.4
largo (IB, Multilingual)
71
1.6
2.3
F55 Additional Zero-shot and Seven-shot Results
In the main text, our evaluation focused on 0-shot performance for multilingual pre-training and 7-shot performance for GLaM pre-training. This appendix provides more complete picture by presenting both O-shot and 7-shot evaluation results for the checkpoints from these respective experimental setups. The results are in the Table and [0
F.6 Details on Estimating Gradient Variance and Magnitude
We estimate the per-task gradient variance and magnitude (measured using the L2 norm) every To 1000 training steps. At each estimation step, the number of samples used for each task corresponds to its assigned portion (b,) of the standard mixed training batch. Our JAX- implementation enables efficient computation of these statistics with negligible overhead relative to total training time.
To evaluate the impact of estimation accuracy on PiKE's performance, we conduct a sensitivity analysis while keeping all other experimental settings and hyperparameters fixed. Specifically, we compare two strategies. The first, used in our main experiments, computes statistics for each task once using samples from single training batch. The second estimates statisties multiple times using different subsets of samples (e.g., distinct micro-batcher and averages them to produce the final values used by PiKE. Results are shown in Table We find that PiKE is robust to the estimation
Table 9: We compare the accuracies (%, higher the better) of different models HellaSwag and its corresponding translated version using 0- and 7-Shot settings. Bolding indicates the best model in the task; Metrics means the average across different tasks.
HellaSmag.cen
HellaSwag (hi) HellaSwag
HelloSaag
HellaSwag (hi)
HellaSwag
Accuracy
0-shot
O-shor
D-shot
Accumic((7)
7-shot
7-shox 1
7-shot
C4 (en), C4 (hi), and C4 (de) datasets, GPT-2 large style, 1B parnms, 36 Layers default, 120K training steps
27.6
32.5
Round-Rohin 26.5
27.6
26.7
26.3
32.0
34.00
29.7
32.4
Random
26.6
27.0
26.9
26
31.9
34.0
30.0
31.7
30.9
Table 10: We compare the accuracies (%, higher the better) of different models on four different Q/A tasks using 0- and 7-Shot settings. Bolding indicates the best model in the task, Metrics means the average across different tasks.
Arch
CSQA
HollaSwag
PIQA
Arct:
CSQA
HellaSwag
PIQA
Accumay
0-shot
0-shot
O-shott 1
D-shat ?
Accuracy(
7-shot 1
7-shot
7-shot
7-shoe
Six domains GLaM dataset, GPT-2 style, 750M
36 layers default
Mix
20.8
29.5
$3.8
46,4
47.2
39.6
37.9
Round-Rae Robin
32.7
30.4
20.3
26.1
53.9
14.3
43.5
36.7
36.8
60.3
52.3
42.7
GLaM
31.7
28.8
19.9
26.3
51,9
45.3
46.9
39.8
38.0
DoReMi
35.6
33.0
23.8
30.0
55.7
46.5
48.6
40.1
37.5
$9.6
PIKE (Uniform)
37.4
PIKE (GLMM)
35.5
33.5
20.4
31.2
36.8
48.1
49.8
43.5
38.0
61.2
29

---

Table 11: Mean accuracies (%, higher is better) of GLaM 740M models pretrained with PiKE on four question answering (Q/A) tasks, evaluated under different PiKE hyperparameter settings. The table compares the effect estimating gradient statistics for each task either once (using single batch) or multiple times using different subsets of samples.
Number of Estimation
2
4
Accuracy
46.8
47.4
47.1
46.8
Table 12: Mean Accuracies (%. higher values indicate better performance) of GLaM 740M models pre-trained with PiKE on four settings. PiKE's SI hyperparameter was fixed constant at 0.1, while (2 was varied. value of C2 means an ablation where the per-task gradient variance term in PiKE is omitted.
(Q/A) tasks, under different PiKE hyperparameter
C2
0
0.001
0.005
0.01
0.05
0.1
Accuracy
45.3
46.2
47.0
46.4
46.9
45.3
Table 13: Pre-training 110M GLaM for 20K steps with different Çi and fix <2 0.01
0.08
0.09
0.1
0.11
0.12
Accuracy 37.18
37.75
37.92
37.81
37.46
noise: computing gradient statistics from a single batch is sufficient to achieve strong performance, making the method both effective and efficient.
E.7 Impact of K, (1, and (2 on Performance
K (Number of Domains) Our experiments show that PiKE achieves robust performance across a varying number of domains (K 2 to 6), consistently outperforming baselines. A key advantage of PiKE is its linear computational complexity with respect to K, in contrast to some existing methods with quadratic complexity. This makes PiKE highly scalable for large K. However, PiKE's formulation is based on the observation that per-task gradients are nearly orthogonal. This assumption typically holds in high-dimensional settings where the number of parameters d is much larger than the number of tasks K (i.c., d >>> K). Therefore, we anticipate that PiKE's 's advantages may diminish when training smaller models on a very large number of tasks.
Ç1, S2 Hypeerparameters of the PiKE algorithm) As derived from our theoretical analysis, PiKE's data sampling rate is designed prioritize tasks with high gradient norm (indicating difficulty) and low gradient variance (indicating training stability). The hyperparameters Çi and {2 control the model's s sensitivity to these two terms, respectively.
To underscore the importance of the gradient variance component, we conducted an ablation study by fixing Çi = 0.1 and varying (2 (Table [2). The case (2 0 represents the complete omission of this term. Optimal performance (47.0% mean accuracy) was achieved at (2 0.005. In contrast, removing the variance component (C2 = caused performance drop 45.3% This degradation highlights the critical role of the per-task gradient variance in regulating PiKE's sampling weights and achieving optimal results.
We also conducted a sensitivity analysis for Ç1. which weights tasks by their gradient norm to prioritize difficulty. Due computational constraints, these experiments were for limited number of training steps. The results (Table [3) demonstrate that PiKE is robust across a range of Çi values.
Based on our ablation studies, we recommend setting Çı in the range 0.05-0.15 and Ç2 in the range 0.005-0.015 for practical use.
E88 Advantages of Adaptive over Static Data Mixtures
To demonstrate the advantages of an adaptive data mixture, we conducted an exhaustive search over various static (fixed-weight) data mixtures of English and Hindi. We then compared the performance
30

---

of these static baselines against our adaptive method, PiKE. The results, presented in Table 14 show that PiKE consistently outperforms all static configurations. This finding empirically highlights the limitations of fixed mixing strategies and validates the necessity of an adaptive approach to achieve optimal pre-training performance.
Table 14:
Pre-training 270M GPT-2 style LMs with context length 2048 for 20K steps using C4 (English and Hindi) datasets.
Method
Wen
Whi
Mean Accuracy
Mix
0.2
0.8
27.43
Mix
0.4
0.6
27.38
Mix
0.5
0.5
27.65
Mix
0.6
0.6
27.90
Mix
0.8
0.2
28.67
PiKE
Adaptive
Adaptive
29.27
F.9 Full Pre-training Results
Tables 15 and 16 present the complete results of pre-training language models across various scales (110M, 270M, 750M, and IB) and scenarios (Multilingual GLaM datasets). PiKE consistently outperforms all baselines across all scales and scenarios.
Table 15:
We report the perplexities (lower the better) on the validation split of multilingual C4 datasets. We also compare the accuracies (%, higher the better) of different models on HellaSwar wag and its corresponding translated version. HellaSwag and its translated versions have 4 choices. Bolding indicates the best model the task, Metries means the average across different tasks.
C4
(bil
C4 (de)
HellaSwag
HelloSwag (hi)
HellaSwag
Perplexity
Perplexity
Perplexity
Perplexity
Accumcy(X)
0-shot
0-shot
O-shot
Single datnset, GPT-2 small style. 270M params, 12 layers default, 120K training stops
(cn)
13.25
13.25
(hi)
4.97
4.97
26.4
26.4
C4 (de)
11.27
11.27
26.1
26.1
(hi)
params, layers default. 120K training steps
Mix
10.50
15.46
5.55
25.5
24.4
26.5
Round-RR Round-Robin
10.57
15.57
5.57
25.6
25.2
10.57
26.3
FAMO 8X
10.38
15.18
5.57
25.7
24.8
26.5
R0
10.45
PIKE
10.15
14.31
5.99
26.5
26.0
27.0
C4 (en), C4 (hä), and C4 (de) datasets, GPT-2 small style, 270M params, 12 layers default, 120K training steps
12,00
25.3
25.5
Round Robin
12.10
16.44
5.91
13.95
25.1
24.3
26.0
21.9
Randon
12.16
16.49
5.95
14.03
25.1
24.7
26.6
23.9
ADO 30
12.01
16.31
$.88
13.54
24.9
24.1
25.4
21.8
PIKE
12.01
15.48
5.92
14.64
25,6
25.4
24.8
Single
large style, parums, Layers
truining steps
C4 (cn)
9.30
9.30
33.6
33.6
(hi)
3,87
3.87
27.5
27.5
(de)
7.72
7.72
28.1
28.1
C4 (en) and C4 (hi) datasets, GPTT22 large sty le. 1B params, 36 Layers default, 120K truining steps
Mix
7.41
10.60
4.22
27.3
28.2
26.5
Round- Robin
7.49
10.72
4.25
27.5
28.0
27.0
Random
7.52
10.76
4.28
28.0
28.9
27.0
ADO
H0
7.41
10.59
4.23
26,
26.0
26.9
PIKE
7,21
9,63
4,30
30,0
32.7
27.3
(en),
large sty
params, Layers 120K training sleps
Mix
8.29
11.13
4.45
9.29
27.5
28.
27.1
27.6
Round- Robin
8.11
11.31
1.97
9.16
26.5
27.6
26.7
26.3
Random
8.48
11.38
9.55
26.6
27.0
26.9
26.1
FAMO 38
8.25
11.04
4.48
9.23
27.2
27.3
26.9
27.3
30
PIKE
9.56
9,49
5.32
13.87
28.7
33,0
27.2
2622
Balanced-PiKE
8.29
11.12
4.46
9.31
27.9
28.3
27.1
28.0
Balanced PIKE
Balanced-FFi (r
8.12
10.02
6.30
8.4
25,9
31.2
26.9
25.6
31

---

Table 16: We report perplexity (lower is better) the validation split of the GLaM datasets, averaging perplexities across six domains when applicable or reporting a single perplexity when only training with single domain. We also compare the accuracies (%, higher the better) of different models four different Q/A tasks. HellaSwag and Arcß tasks have 4 choices, CSQA has 5 choices, and PIQA has 2 choices. PiKE (Uniform) means PiKE using initial sampling weights of 1/6 for each task and PiKE (GLaM) means PiKE using GLaM tuned weights as initial task weights. Bolding indicates the best model in the task, Metrics means the average across different tasks, underlining indicates PiKE beating Mix, Cound-Robin, Random methods
GLaM
ArcF
CSQA
HellaSwag
PIQA
Perplexity
Accuracy (P6)
7-shot f
7-shot f
7-shot +
7-shot 1
Single domain GL.M dataset, GPT-2 small style, 110M params, layers default
Wikipedia
9.96
33.5
32.5
20.9
27.3
53.3
Filtered Webpage
16.05
37.2
38.4
26.8
27.6
55.8
News
9.33
33.8
31.1
22.7
27.0
54.5
Forums
22.87
35.5
23.4
28.7
57.6
Books
16.81
34.7
34.3
22.1
27.8
54.7
Conversations
18.27
36.1
32,6
25.6
28.6
57.6
Six domains of GLaM dataset, GPT-2 small style, 110M params, 12 layers default
Mix
18.27
36.2
35.6
24.1
28.5
56.7
Round-Robin
18.45
35.9
35.8
24.2
27.5
56.0
Random
18.48
35.5
34.3
22.4
28.4
56.8
FAMO 13
18.19
35.9
35.3
24.2
27.8
56.4
ADO 30
18.27
36.2
35.7
24.8
27.7
56.4
GLaM 18
18.91
35.8
35.3
24.1
28.5
55.1
DoReMi
18.98
37.0
36.0
28.3
28.2
55.3
PiKE (Uniform)
18.44
37.4
36,8
27.5
28.5
57.0
PiKE (GLaM)
19.34
37.8
39.0
27.0
28.0
57.0
Single domain of GLaM dataset, GPTT2 large style, 750M params, 36 layers default
Wikipedia
7.24
35.9
35.1
24.0
30.5
53.9
Filtered Webpage
11.12
40.9
36.7
33.2
34.2
56.5
News
6.62
37.4
33.6
24.7
34.1
57.3
Forums
16.29
43.6
38.0
35.8
39.7
60.7
Books
11.83
41.3
40.0
33.0
34.5
57.8
Conversations
13.50
42.2
36.9
33.2
39.2
59.6
Six domains of GLaM dataset, GPT-2 large style, 750M params, 36 layers default
Mix
12.77
46.4
47.2
39.6
37.9
60.9
Round-Robin
12.98
44.3
43.5
36.7
36.8
60.3
Random
12.99
42.7
41.7
34.2
36.6
58.2
FAMO 38
13.25
45.0
43.7
40.0
36.4
59.8
ADO 30
12.77
45.9
45.5
38.7
38.1
61.1
GLaM 18
13,20
45.3
46.9
39.8
38.0
56.4
DoReMi
69
13.25
46.5
48.6
40.1
37.5
59.6
PiKE (Uniform)
13.22
47.6
49.6
43.2
37.2
60.4
PiKE (GLaM)
13.35
48.1
49.8
43.5
38.0
61.2
Balanced-Pinn PPKKE (T 1)
13.21
47.5
48.8
42.5
37.6
61.2
Balanced-PiKC (T=3) 3)
13.26
47.2
48.8
41.5
37.2
61.3
Balanced-PiKEE 55 13.19
48.2
49.3
42.6
38.5
62.4
E11
Adaptive Sampling Weights of PiKE During Pre-training
Figure 8 illustrates how the adaptive sampling weights of PiKE evolve during language model pre-training. Compared to the Mix sampling strategy, which assigns equal sampling weights each task, PiKE adaptively adjusts the sampling weights among English, German, and Hindi by leveraging the positive interaction of task gradients. This adaptive data selection allows PiKE to achieve superior performance compared to fixed or beuristic-based baselines.
32

---

0.40
0.3
0.6
0.35
1022
0.2
8 0.5
FINE
0.4
0:30
0.1
0.0
0
20
40
60
80
100
120
20
40
60
100
120
Traiiing Steps (K)
Training Steps (K)
# Training Steps (K)
Figure 8: The sampling weights for each dataset during pre-training IB GPT-2-style multilin- gual language models mC4 (English), mC4 (Hindi), and mC4 (German). Here, Wen represents the sampling weight for the English dataset, Whi for the Hindi dataset, and Wht for German dataset.
G
Derivations and Proofs
G.1
Detailed Derivation of equation
3
Recall that
(heie + b2e262) 0, + 2,
g: bi + b2
Then
Oc41
0,
" bi + bi
(blage) T + b2e2e) O, 1/2
17 b
0:
0
b2
0:
172
Now consider the loss functions for task 1, L1 (fe+1), and task 2, (2(0.))), separately, taking the expectation over the randomness of
E
0c+1)²
2
E
T
1
nb
0
0.
el 1/2
27
2
n2b
1
([1
nb
0]
+ 5" e
Qei
2
1
((1
ndu
2
01
2 ej Qe1
2
Similarly, for task 2, have
1 E[C2(0++))))= 2
((1
1722
02,1
+
Qe2
b
2
where 01,e and 02 denote the first and second component of the vector 0,, Combining the losses for both tasks, the total expected loss becomes
E(((0++1) (0
E(Ce
((1
11b
1
(11
182
2
b20 + b20z
01,1
+
02,1
+ 17
2
b
b
62
(1 nb1
1
nbb
6107
+ byoz
2
1
0009
+
2
(1
20
2,t
2
52
which completes the derivations.
33

---

G.2
PiKE: Main Theoretical Results
Lemma G.I. Assume 2(K-1
S e, have
K
(L)(1
S 2c(K
k=1
Conversely, (f)] z dg, Yk, then
V < (1
62
Σ v 6g
k=
Proof: We first prove the first direction. Notice that
K (0)12
k=
K
(O)I
+
ΣΣ
Sc
1=1
1=
where we use the definition of VC()) and expand the term. Then we have
[C) +
(L)
[V e ΣE (1)2 (V2,00))))++
k=1s#k
-
2e(K 1))
where (a) uses the Definition 3.1 (b) uses symmetric identity. Thus we get
(0)1 - - 2c(K- 1)
This completes the proof of the first inequality. We now prove the second inequality. Notice that
|V Σ (0)
(L,(0),LL.(00)
K [V +
Σ
j
K
(1 2) VL+ + V
E Σ Σ
(2,1
jyth
K
2
$ (1 8 + ē
Σ
6z
where (a) use the Definition 3.2 and (b) combines the second and third terms and use the condition
that L (0) < 6R. This completes the proof of the second inequality.
Lemma G.2. For the optimization problem
1 2
min yo MC
wk u + u
(13)
S.f.
Wk 1,
wk 2 0,
4k
1=1
34

---

the optimal solution is
p + 2k
" = max
0,
(14)
Kik
where yt is chosen such that Σx= w
Proof: Consider the Lagrangian function
K
K Σ
K
[(w),....
Wk: On .ak)
w + FuR +
Wk
Og
2
where j Lagrange multiplier for the equality constraint for the constraint Σk=1 WA and Og 2 0 are Lagrange multipliers for the nonnegativity constraints wk. Take the partial derivative of C. with respect to wi and set O:
OC Own
AE + WaKK
Otk. =0
From the Karush-Kuhn-Tuue Tucker (KKT) conditions, also have wo 2 0, are 2 and on 0. If w > O, then og = 0, which implies
# +
0 = > + WINK +
wh
Kik
- (n + As) /nk is negative, then w = O must hold. Combining these, we get
0,
1
Au
3 = max
Finally, the Lagrange multiplier y is determined by enforcing the equality constraint:
with y chosen so that the w sum to 1. This completes the proof.
Theorem G.3. (Formal Statement of Theorem 3.4) Suppose Assumption 3.3 is satisfied. Assume that at the given point Og the gradients are g-conflicted and e-aligned with g K-24b/be Vk. Moreover assume batching is performed based on the mix strategy equation i.e.
= Σ by gr,ks
(la(Or;
ge
Then, we have
B[C(O. - (CC(O.--g.))10 [(O.) +
Σ by 77 (--3
+ Lip
212 ok
bi
Lai > (0)))
26
where 0 < B e ming + c(-K +2 b) and + o(K - 1), the expectation is taken over batch sampling randomness under the mix strategy (bi...., bk), and Fe denotes the natural filtration of our process.
Proof: We begin by revisiting the multi-task optimization problem under consideration. The objective defined as:
min
(()) = Σ E. BANDe x)),
(15)
Oe
where L(O) is the aggregate loss over all tasks. Assume we mix the gradients with taking by i.i.d. samples from task & for k = 1..... K. Then under, Assumption 3.3 and based on equation 2 the estimated gradient direction is given
35

---

bk Vl+
81
(16)
=k=1 bk
Dx
Let Ou+1 be the updated point after gradient descent with 0,+1 the following incquality holds for the updated parameter 0+:
0: -7--- By the descent lemma,
Lij²
((0 S ((O.) vi8. VL..))
(17)
Conditioned on the natural filtration process Fr, we take the expectation over the randomness the samples draw and obtain:
Tap² +
E (c( |F & ((0.) - lB Fi]
VL(0.)
((g())
21
2
1 K
K
(a)
((0,) -
by (.
(0,)
Ln
K
K
+ 262
by
(L)
+ (bbooit
1m
km
K
K +
((O,)
b [
by (V,, VCL((.)) ((.(0))))
b
k=1
K
2
K
+
by (0.)
+ +
(b)o
26
k=1
where (a) substitutes the definition of g. and uses the Assumption 3.3 and (b) expands the terms. The notation Fr denotes the natural filtration of the random process (loosely speaking, conditioned on the current time / and the past history). Continuing our simplifications, we have
(C((0+1) Fi
(a) I
C(O,)
n
by (0)1
K
1#
Lnj²
K
I LA (0.)1 + Σ Σ bybb VC,(1))) Σ buo
26
1=1
K
(
((0.)
1
bk
cba - 1k) c Σ bi
(0.2)
b
2n
K L (0)12
K Σ book
bybn (O.) (L)()))) +
21
k= k=1jz
((0.)
(be s(K 1)bb c(b - [2)
b
Ln
K V LA (0.)1
Σ
Σ
obyb [=b66 Σ book
26
k=1
K
(
((0.)
1
Σ
bu(1 + cccKKk
K + b/bb)) VZ
36

---

Lnj
2
K
Σ bk
V
+ (0.)ll
Σ
b
La (0:)
Σ book
21
K
((0,)
h&
(1++ 2 (0:)1
Ln
K
K
K
+
čk
Σ
b ( -1 (11-a) Σ
V
Σ
boo
212
k=1
1=1
1=1
L(11)
Σ (1 c + 2/bbbbb)]
Laj²
K
b (0,)11 (1 (1- eK)
Σ
Σ book
(18)
26
where (a) applies Definition 3.1 the second term and expands the third term, (b) expands the summation in the second term, (c) uses the identity Σ Ejje by = > - bx) in the second
term and applies Definition 3.2 to the third term, (d) combines terms in the third term, and (e) uses
the inequality | u N > N
i=1 u.. where is column vector. We define / and y such that
b
3 min + s K by #
1 + a(K 1)
Then using the definition of 3 and y. substituting back we have
E [[((0++1) [F]] C(O,) nB
VC.(0.)1P)) bkll
+
Ln²
7
K
b
V
(0:))1 +
bRo
26
k=1
1/3 (-
Laj²
K
2
-((,)) +
Σ bk
[V
o
b
21 7
VL
(.))
212
k=1
which completes the proof.
Theorem G.4. There exist loss functions {Ex(-, 1)K= satisfying Assumption 3.3 whose gradients k=1
are e-conflicted and e-aligned with == 0. For these losses, the upper bound in equation ( is fight when gradients are computed using the mix strategy equation that is, the inequality in equation (8) holds with equality.
Proof. The proof is by simply generalizing Example 2.1 and showing the steps in the proof of Theorem G.3 are tight for this example. Consider the multi-task learning problem where the loss of task is given by
L (e (e 0)² x x
l (0, mT)
2
with x NO, I is the data for task is easy verify that this loss is smooth with smoothness parameter L. Moreover, we can show that Assumption 3.3 is satisfied. Notice that
(x( E (fx
(cT
and
K
L 11012
((0) = Σ (0)
k=1
First notice that it is easy to check that c= 7= 0. Lett g: be the direction obtained by equation 2
and assume o,+1 0, 79.- Using the form of L(0), can easily show that (by expanding the foslos)):
Lij²
C(0,,) L(0.) ng. VL(..)
g 1
(19)
37

---

Conditioned on the natural filtration process Fi. we obtain:
E [ Fi] = [(0.) nE(g Fil VL((0) + Lnp²
E (lBBil) Fi)
2
(a) [(0.) 11
by (O)
FL (0.)
Lnj
+ 20
by (0))
Σ (bRo
(
((O,)
k=1
Ln²
K
K 66 (0,)l + E(bnoR)
+
Σ
26
A=1
where (a) substitutes the definition of ge and uses the Assumption 3.3 and (b) is because (O,) (O.)) = 0, vj * k. Since 3 = y = in this example, we can write
n [.)(
Lay²
Tap² bk 2b [l.
E (C(0+41) Fi] (()))
bk
20
which shows our upper-bound is tight in this example and completes the proof.
Theorem G.5. (Theorem 3.5 in the main body) Suppose the assumptions Theorem G.3 is satisfied
and we run the Conceptual PiKE Algorithm (Algorithm 2 initialized at Gg with the SGD optimizer in
Step 10 of the algorithm. Let AL = [(Ou) ming L(0) and maax max; σg- Suppose > 0 is
given constant and the stepsize 1) S Lo / Lng Then, after T
95
23AL iterations, Algorithm
n0
finds a point 0 such that
E V < 5,
=k= K.
(20)
Moreover, if we choose
b+ Lag' then the Conceptual PiKE algorithm requires at most
Lo moax
(0 m + 76)
b
T =
82
iterations to find a point satisfying equation (20
Proof: We prove this by contradiction. Assume that max; El > 8 fort O,...,T 1. First notice that Theorem G.3 implies that for all t, we have
K
T max
E[C(0++)) |R] < C(0,) +
w
(-n3)V22200))11--
2b
k=
w
V
(21)
where [wi]k_ is the minimizer of the RHS of the equation 21 on the constrained set ((w).... wx) >
E wk 1, we 2 vk Kj. Since w is the minimizer of the RHS of
equation (2T). we have
wi (-n3]
Lnj
Wx LTi² V [C.()))) +
2
2h
2
Ln² o² mmx
Ln >1
<nn3) (0.)
+
VL; (0.)
2b
2
(22)
38

---

where K E arg max; (0)11 Moreover, we have
B6
L (0)11
7 1
mas b
+
L
I
oan b
+ I Lxy (0.)
Therefore,
Laj²
Laj
B17
-n (Ox))l
0
max
V VLx; (0.)
VV (0:)
(23)
2h
2
Combining equation
21
(22
and
we obtain
By | Luz (0.)1l
B[((0++1) Fil - ((0,)
Or equivalently,
3n
[[C(O,+)) S (0.))
max [x
Taking the expectation over all sources of randomness in the algorithm, summing the above inequality from /. = 0 to = T 1, and simplifying the resulting telescoping sum, we obtain:
3n T-1
B[[(Or)] S C(0,)
max B (
2
Recalling the contradiction assumption that d < 1B[1| L+ [1) we get
E[C(Or)] (()o) 2
Using the definition A. & C(0)) ming [(0), we get
2 L
T Σ
Finally notice that by setting 17
888
get
tmn + 1,76"
2LAr
(o
T TST-
78
B1
352
b
which completes the proof.
G.3 Conceptual PiKE versus Static Uniform Mix Batching
now perform standard analysis Stochastic Gradient Descent (SGD), adapted to our specific setup where uniform (static) mini-batching is used (i.e., when by = b/ K, Vk: I,.... K). As is common in the analysis of gradient- algorithms in smooth, nonconvex settings, we begin by quantifying the expected decrease in the objective function at each iteration. This is known the descent per iteration. Once this is established, we use a telescoping sum argument to derive the iteration complexity of the algorithm, which tells us how many steps needed to achieve desired level of accuracy, given a properly chosen learning rate.
Let's begin by analyzing the descent that occurs in single iteration. Recall that the objective is defined as:
K min (()) = Σ ELND. [f((0:))] BER
(24)
Under Assumption 3.3 and using uniform static mini-batch sampling-we wwhre each source con- tributes equally to the batch with size by b/ K the estimated gradient direction at iteration t is given by:
K K
g:
Σ
(25)
1=1
Dr
39

---

This expression aggregates gradients from all sources using equally sized sub-batches, providing mini-bated estimate the full gradient. Let 0,+1 denote the updated parameter after performing
gradient descent step, given by 0,+1 0r 78t. By the descent lemma, the following inequality holds for the updated parameter: 0.+1:
Laj
((0(41) S ((0,) Vit VL(
(26)
Conditioned on the natural filtration Fr, we take the expectation with respect to the randomness in the sampled data and obtain:
Ln 2
[c(00++) ((O.) - nlEigr Fi]
VL((,)
+
E
(I|g|l
Fi)
K
T
K
-(0,) - 11
Σ
Σ
V Li (0:)
K
k=1
Laf
K
2
K
+ 21
(O,))
+
o
K
K
k=1
Observing the fact that L(0) Σx Lx((), we can write
Lid
b²
K
E [[C(O++)) < C(O,)
!
|C(01 +
+
oh
(27)
K
26
K2
K
k=
Define o² == ===i ohe By summing Rearranging the terms, we obtain:
[[C(0.+1) [i] < ((0.) + (- 7 Laj
Lab²
|VC()))
(28)
21
2bK
Taking the expectation all sources of randomness in the algorithm, summing the above inequality from / = 0 to 1 = T 1, and simplifying the resulting telescoping sum, we obtain:
[[ ((O))
(-x
7
Lnp
T-1
[[VC(()) Σ
LT
0².
(29)
2K
26K
Let A, == L(O6) ming L(O) denote the initial suboptimality gap. As is standard in iteration complexity analysis, we are interested determining the number of iterations required to reach
point 0% where E ||V )l2] 6. Suppose, for the moment, that this condition has not yet been satisfied. Then, we can rewrite the above inequality as
Lnp
La 2K
LIn T8+
AL +
0²
(30)
26K
Equivalently, under a properly chosen step size (to be defined later), we obtain the following upper bound on the number of iterations
A
T
Lxp 2K²
Lnd 26K
o²
K
By optimizing the choice of 1) minimize the right-hand side, derive the following bound the iteration complexity
2T L
(5 o K
T S
82
In summary, for the uniform mix batching strategy, the standard analysis yields the following upper bound on the iteration complexity:
2LA
8 +
o² K
I Uniform
(31)
82
40

---

In contrast, when the task gradients are nearly orthogonal i.e., when the alignment/conflic parame- ters satisfy č, c = 0 O-Thhoremm 3.5 gives the following iteration complexity bound for the Conceptual PiKE algorithm:
2 (5 max/
TAKE
(32)
6²
which clearly illustrates the advantage of PiKE for this regime in terms of iteration complexity upper-bound.
G.4 Balanced-PiKE: Theoretical Developments and Derivations
Here we prove Lemma 3.6 which is the idea behind the Balanced-PiKE algorithm. Before proving let us, for the sake of completeness, re-state the lemma
Lemma 6.6 (Restatement Lemma 3.6 tion problem
Let < <EE R. Then, the tilted empirical risk minimiza-
K
min log
Σ El((
0
equivalent to
min
max YER Σk- WL=T
yk (0)
yk log
V/k
(33)
2x(00
Moreover, for any fixed 0, the inner maximization problem maximized at yī
Vk.
The above lemma is the direct consequence of applying Lemma G.8 and Lemma G.7
Lemma G.7. For the problem
K
K
)
U/: log
max
Uk
yeat Ex
k=
the optimal y is given by
T2x -1
v
Proof: We start by forming and maximizing the Lagrangian function
K
K
K
max
Σ
log
yk
+
Σ yh
y YERK
k=1
where is a free variable. Taking the partial derivative of the objective with respect to ys and setting zero gives
v T=e
where the coefficient is independent of the index l and should be chosen such that Σx vi implying
Te
y
j==
eTa
Lemma G.8. Letx E R and T > 0. Then,
K
K
K
log
Σ
Tak
max YERK
MoTT
Uk
log
Uk
km
Proof. The proof by simply plugging in the optimal value y* obtained by Lemma G.7 in the objective function.
41


Reviews:
[{"review_id": "XPmvBl7yKs", "reviewer": "Reviewer_hDVe", "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.", "questions": "See Weakness.", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. 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 .\n2. Empirical results are solid, with clear improvements over baselines. The datasets used are plentiful.\n3. The paper has the potential to be high impact with many key insights used in practice.\n\nWeakness:\n1. 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?\n2. 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.", "was_revised": true, "final_justification": "Good work - I maintain my Accept rating", "rebuttal": "Thank you for your detailed constructive feedback. We are glad that you found our contributions solid and the paper impactful. We also appreciate your recognition of the care we put into presenting the algorithm and results. Please find our detailed responses to your questions below\n\n$\\quad$\n\n> 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?\n\nThank you for raising this important question about the applicability of our findings to structured problem domains. To explore this, we conducted an additional experiment using code, which is a highly structured data type.\n\nIntuitively, as model size increases, gradients from different tasks are more likely to become orthogonal. This is because in high-dimensional spaces, two (noisy) vectors tend to be nearly orthogonal. To test this idea, we used the final GLaM 740M checkpoint and computed task gradients on the CodeSearchNet dataset. We then measured the cosine similarity between these code-derived gradients and those from the original six GLaM domains.\n\n\n\nThe results, presented in Table 1, show that the cosine similarities are consistently near zero. This outcome supports our hypothesis, suggesting that the gradient non-conflict property does extend to structured domains such as code.\n\n\n\nTable 1: Cosine Similarity of Code vs. six different GLaM Task Gradients at the Final 740M LM Checkpoint.\n| Datasets | glam_wiki | glam_web | glam_conv | glam_forums | glam_books | glam_news |\n|---|:---:|:---:|:---:|:---:|:---:|:---:|\n| Cosine Similarity (Between 1 and -1) | 0.00148 | 0.00097 | -0.01256 | 0.00409 | -0.00213 | -0.01023 |\n\n\n\n$\\quad$\n\n\n> 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.\n\n\n\nUnderstanding the role of these hyperparameters is key to the easy implementation of our method in practice. Below we provide an intuitive explanation for the role of these hyper-parameters.\n\n\n**$K$ (Number of Domains):**  Our experiments shows robust performance of PiKE across different numbers of domains (K=2 to 6), consistently outperforming the baselines. Unlike some other existing methods with quadratic computational complexity in $K$, PiKE has linear computational complexity in $K$, which makes it favorable for large values of $K$. However, PiKE is based on the observation that per-task gradients are nearly orthogonal. This observation intuitively holds in large dimentional setting where the number of training parameters is much larger than the number of tasks, i.e., $d >> K$. Hence, we  expect that PiKE offers no benefit in training small/medium size models with a huge number of tasks.\n\n\n\n**$\\\\zeta_1, \\\\zeta_2$ (Hyperparameters of the PiKE algorithm):** We assume you are refering to $\\\\zeta_1, \\\\zeta_2$ in the PiKE algorithm here. As derived from our theoretical analysis, PiKE's data sampling rate is designed to prioritize tasks with a high gradient norm (indicating difficulty) and low gradient variance (indicating training stability). The hyperparameters $\\\\zeta_1$ and $\\\\zeta_2$ control the model's sensitivity to these two terms, respectively.\n\n\nThe hyperparams $\\\\zeta_1$ scales the influence of the gradient norm, effectively prioritizing tasks that the model currently finds more difficult.\nWe conducted a sensitivity analysis for the hyperparameter $\\\\zeta_1$, which weights tasks by their gradient norm to prioritize those the model currently finds more difficult. Due to computational constraints, these experiments were performed over a limited number of training steps. The results in Table 2 show that our method is robust across different values of $\\\\zeta_1$.\n\nTable 2. Pre-training 110M GLaM for 20K steps with different $\\\\zeta_1$ and fix $\\\\zeta_2=0.01$. \n| $\\\\zeta_1$ | 0.08 | 0.09 | 0.1 | 0.11 | 0.12 |\n|:---:|:---:|:---:|:---:|:---:|:---:|\n| Mean Accuracy (%) | 37.18 | 37.75 | 37.92 | 37.81 | 37.46 |\n\n\nThe hyperparams $\\zeta_2$ penalizes tasks with high gradient variance, encouraging the model to favor tasks with more stable and reliable gradients. When $\\zeta_2=0$, the variance consideration is removed, and the sampling strategy simplifies to a greedy approach focused only on the most difficult tasks. As shown in  Table 11 in the Appendix, this ablation causes mean accuracy to drop from $47.0 \\\\%$ (with $\\zeta_2=0.005$ ) to $45.3 \\\\%$ (with $\\zeta_2=0$). This performance degradation highlights the importance of considering gradient variance in achieving optimal sampling strategy.\n\n\nAll in all, we recommend using the value $\\zeta_1=0.05 \\sim 0.15$ and $\\zeta_2 = 0.005 \\sim 0.015$ for practical use based on our ablation. We will add these discussions on the role of hyperparameters to the paper if the paper gets accepted.\n\nWe once again thank you for your time and valuable feedback. We hope that our responses and planned revisions have fully addressed their concerns. Please let us know if you have further questions."}, {"review_id": "VNsrjzGVKT", "reviewer": "Reviewer_vude", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "**Weakness 1 (weight  50%)**. Authors added a concise fine-tuning experiment showing PiKE matches or exceeds manual mixing without hyper-parameter search. This directly answers my primary objection.\n\n**Weakness 2 (weight  25%)**. Authors now provide exact $\\zeta$ values per setting. Please briefly report the total number of hyper-parameter combinations evaluated and the approximate GPU-hours required for the $\\zeta_1$ and $\\zeta_2$ grid search. This will help readers gauge reproducibility and computational overhead.\n\n**Weakness 3 (weight  25%)**. Added empirical tracking and theoretical explanation demonstrate PiKE implicitly adapts to dataset size. Issue resolved.", "rebuttal": "Thank you for your time and for providing such a detailed constructive feedback. We are encouraged that you found our theoretical analysis solid, our experiments comprehensive, and the effectiveness of PiKE convincing. We address your specific questions below:\n\n$\\quad$\n\n> **Weakness 1:** 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.\n\n> **Question 1:** How does PiKE perform when directly trained on downstream datasets, rather than just during the pre-training phase?\n\nWhile PiKE was developed for pre-training data mixing, we agree that evaluating it in a direct multi-task fine-tuning scenario would help assess its broader applicability. Our experiments in the paper focused on the computationally intensive pre-training phase, but PiKE is designed as a general framework (assuming no gradient conflicts), not limited to any specific training stage.\n\nTo address your question, we conducted a new fine-tuning experiment during the rebuttal period:\n\n* **Base Model:** We used a Llama-1B style architecture pre-trained on 70B tokens from the Fineweb-Edu dataset.\n* **Task:** We fine-tuned the model jointly on a mixture of HellaSwag and PIQA for 2,000 steps, using a learning rate of $1 \\times 10^{-5}$.\n\nThe results are shown in Table 1. While simple static data mixing can be effective during fine-tuning—especially since the lower computational cost allows manual tuning of mixing ratios—our experiment shows that PiKE achieves competitive performance without the need for manual tuning. This experiment demonstrates that PiKE is a viable and effective method for fine-tuning on downstream datasets.\n\nDue to time constraints, we only conducted this experiment during the rebuttal phase. However, this result confirms PiKE's potential use in post-training and fine-tuning stages. It is worth noting that fine-tuning is typically easier and less computationally intensive than pre-training. As a result, reasonable mixing coefficients can sometimes be found by trial and error. Nonetheless, the need for a reliable, adaptive algorithm is more critical during pre-training, which is why our work primarily focuses on that phase.\n\n\n\nTable 1. Downstream Performance of different methods for fine-tuning a Llama-1B base model on a mixture of the HellaSwag and PIQA datasets.\n| Method | Mean Accuracy  | HellaSwag | PIQA |\n|:---:|:---:|:---:|:---:|\n| Base Model (No Fine-tuning) | 63.46 | 55.70 | 71.21 |\n| Mix (Uniform, Static)  | 65.75 | 58.43 | 73.07 |\n| PiKE (Adaptive) | 65.94 | 58.37 | 73.50 |\n\n\n\n\n$\\quad$\n\n> **Weakness 2:** 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.\n\n> **Question 2:** Could you provide the specific values of the hyperparameters used in the experiments?\n\nThank you for pointing this out. We reported our hyper-parameter grid  in Table 7 in the Appendix. To address your question, we will explicitly mention the result of the hyper-parameter search in our grid for each experiment. In particular, we will include  the table below into the experiment section of the revised manuscript:\n\nTable 1. Specific values of $\\zeta_1$ and $\\zeta_2$ for PiKE across different experiment settings. \n\n\n\n| Experiments Setting | $\\zeta_1$ | $\\zeta_2$ |\n| -------- | -------- | -------- |\n| en+hi, 270M, PiKE     | 0.1     | 0.01     |\n| en+hi+de, 270M, PiKE     | 0.075  | 0.005    |\n| en+hi, 1B, PiKE     | 0.05     |  0.005     |\n| en+hi+de, 1B, PiKE     | 0.1     | 0.01     |\n| GLaM, 110M, PiKE (Uniform) |   0.1   |   0.01   |\n| GLaM, 110M, PiKE (GLaM) |   0.075   | 0.005     |\n| GLaM, 750M, PiKE (Uniform) | 0.15     | 0.01     |\n| GLaM, 750M, PiKE (GLaM) | 0.075     | 0.005     |\n\n\n\n$\\quad$\n> **Weakness 3:** 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\n> **Question 3:** Can PiKE’s adaptive weighting scheme be extended to incorporate dataset size or task-specific downstream importance, beyond just gradient statistics?\n\nAs you mentioned, dataset size plays an important role in determining the optimal mixing strategy. While PiKE does not explicitly use dataset size as an input, it implicitly accounts for it through its gradient-based sampling mechanism.\n\nThe intuition is as follows:\n\n\n1. For a low-resource task (e.g., Wikipedia with 3B tokens in GLaM), the model quickly learns the underlying data distribution and begins to converge or overfit. As this happens, the gradient norm for that task decreases, indicating that there is less left to learn.\n2. Conversely, a high-resource task (e.g., Conv with 174B tokens in GLaM) offers a richer learning signal for a longer period, resulting in a sustained, larger gradient norm.\n\nSince PiKE increases the sampling rate for tasks with larger gradient norms, it naturally reduces focus on low-resource tasks as they converge, and shifts focus to high-resource tasks where more can still be learned. In this way, PiKE implicitly adapts based on how much information remains in each dataset.\n\nTo demonstrate this empirically, we tracked the sampling rates PiKE assigned during GLaM pre-training. As shown in Table 2, PiKE initially assigns a reasonable weight to the Wiki dataset. Over time, this weight decreases, while the weight for the larger Conv dataset increases. This dynamic behavior shows that PiKE inherently adjusts for dataset imbalance. We will include this discussion in the paper if it is accepted.\n\n\n\nTable 2. Sampling Rate of Wikipedia (3B tokens) and Conversation (174B tokens) during PiKE-based pre-training of a 740M Models with GLaM dataset. \n\n\n\n| Iteration Steps (K)  |  30 | 60 | 100 | 120 |\n| -------- | -------- |-------- |-------- |-------- |\n| $w_{\\text{wiki}}$   | 0.026     | 0.023     | 0.021 | 0.019 |\n| $w_{\\text{conv}}$   |  0.278   | 0.346 | 0.344 | 0.35 |\n\n\nRegarding your question on incorporating downstream task importance: yes, PiKE's framework can be easily extended to account for task-specific priorities. While we derived fair-PiKE under a specific utility function ($\\alpha$-fairness), the framework is modular and general. For example, one can easily introduce a set of user-defined importance weights, $\\\\left\\\\{\\\\beta\\_1, \\\\beta\\_2, \\\\ldots, \\\\beta\\_K\\\\right\\\\}$, where each $\\beta_k$ represents the downstream importance of task $k$. These weights can then scale the loss for each task, effectively guiding PiKE to allocate more resources to higher-priority tasks, even if their gradient signals are temporarily weak (before scaling).\n\nThis illustrates the flexibility of our approach. We will add this important discussion to the paper if it is accepted. \n\nThank you again for your valuable time and feedback. We hope our response has resolved your concerns and welcome any further questions you may have."}, {"review_id": "sa5ArWqY5m", "reviewer": "Reviewer_w6wi", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# 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.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and insightful feedback. We’re glad you saw our work as a meaningful step toward more principled data mixing algorithms. We appreciate that you recognized its strong theoretical grounding, efficiency, and clarity. We also agree that many existing methods lack a clear or principled foundation. In contrast, we designed PiKE with well-stated assumptions and a focus on principled development. This approach has helped us better understand both the strengths and limitations of PiKE. We are very happy that you noticed and appreciated our efforts in this direction. \n\nWe address your questions below.\n\n\n$\\quad$\n > Weaknesses: It would be interesting to compare the theoretical derivation in this paper to the derivation in DoGE. 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 > Question: Can you compare derivation/results with DoGE?\n\nThanks for bringing up this related work. We will discuss it in our revision. As you correctly pointed out, PiKE relies on both the gradient norm and the variance of the gradients, while DoGE relies on the inner product between task gradients. There are two important points to highlight here:\n\n- In large models, the gradients are almost orthogonal (under the assumptions of low conflict and low alignment, as motivated in the paper based on our experiments). This means the inner product between task gradients is either negligible or equal to the norm of the per-task gradients. Specifically, the inner product between gradients of two different tasks is close to zero, while the inner product of a task's gradient with itself equals the squared norm of that gradient. Therefore, from the expected value viewpoint, DoGE and PiKE have a similar flavor (since DoGE rely on inner-product, which is related to the squared gradient norm under our assumptions). However, PiKE uses the near orthogonality to simplify calculations and reduce memory requirements. In particular, PiKE does not need to maintain or store the overall gradient, whereas DoGE does, which adds memory overhead. PiKE only requires to store two scalar values per-task: gradient norm and variance.\n\n- As you noted, PiKE also relies on gradient variance. We believe gradient variance is an important signal. Ignoring this signal reduces performance, as shown in Table 11 in the Appendix. The importance of this signal is also intuitive: if two tasks have the same expected gradient but one has much higher variance due to noisier data, it makes sense to focus more on the task with less variance.  This point has been mostly overlooked in recent data mixing literature, though classical statistics has considered this fact. For example,  noise is considered in data weighting in the Feasible Generalized Least Squares (FGLS) method in classical statistics.\n\n\nIn summary, PiKE uses gradient variance in addition to the mean, unlike DoGE. Also, leveraging near orthogonality simplifies derivations and reduces memory overhead. We will clarify these points and the connection to DoGE in our revision.\n\n\n$\\quad$\n\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 Aioli 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 at each stage, for certain domains.\n\n> More realistic pre-training experiment to motivate when adaptive data mixing is optimal.\n\nThank you for your suggestion. To provide a more realistic motivation for adaptive mixing, we conducted a brute-force search over various static (**fixed-weight**) data mixtures of English and Hindi. We then compared these baselines against our **adaptive** method, PiKE. As shown in Table 1 below, PiKE consistently outperforms all static mixture configurations we tested. This result empirically is in line with our intuition, illustrating the necessity of an adaptive strategy for achieving the best performance in pre-training. \n\nTable 1. Pre-training 270M GPT-2 style LMs with context length 2048 for 20K steps using C4 (English and Hindi) datasets. \n| Method | $W_{en}$ | $W_{hi}$ | Mean Accuracy |\n|:---:|:---:|:---:|:---:|\n| Mix  | 0.2 | 0.8 | 27.43 |\n| Mix | 0.4 | 0.6 | 27.38 |\n| Mix | 0.5 | 0.5 | 27.65 |\n| Mix | 0.6 | 0.6 | 27.90 |\n| Mix | 0.8 | 0.2 | 28.67 |\n| PiKE  | Adaptive | Adaptive | 29.27 |\n\nIn our revision, we will include this experiment in the paper, as it addresses the important question of the relevance of adaptive mixing in complex LLM pre-training. We will also keep our current simple regression example because it is easy to replicate and can be verified very quickly on any computer.\n\n\n\n$\\quad$\n> One property of the mirror descent weight update is that in the limit, the updates move 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\nCould you please clarify what you mean by “accumulation”? In the update rule, we weight the step by both the gradient and its variance: if that combined term grows large, it can produce abrupt changes to the previous update (since our update is multiplicative). \n\n\n\n\n$\\quad$\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\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> More validation of regular PiKE, such as arguing that the training loss curve goes down more quickly than other mixing approaches.\n\nThank you for raising this important point. The confusion arises from a reporting error in our manuscript. We mistakenly reported the arithmetic mean of perplexities instead of the correct aggregate metric, which is the average loss. When using average loss (which is related to the geometric mean of ppls), PiKE's performance is highly competitive. For example, in the GLaM 740M pre-training results (Table 2), PiKE’s average loss of 2.55 is very close to the 2.50 achieved by the strong Mix baseline. We will correct this issue if our paper is accepted.\n\nWe also acknowledge that perfect alignment between our method's objective and the final validation loss is not always guaranteed. This may be due to several factors: first, our theory assumes an SGD optimizer, while we use AdamW in experiments, which is standard for LLM pre-training. Second, PiKE optimizes the training loss directly, whereas reported metrics are on the validation set, which may have a distribution shift. Third, our hyperparameters might not be perfectly tuned to maximize the descent in the objective during training (our theory is for perfectly tuned hyper-parameters).\n\nThat said, the loss values remain comparable, as noted above. Moreover, as discussed in earlier works (such as the ADO paper), simple mixing strategies can be very competitive in final perplexity but may not perform as well in model generalization. Understanding this gap requires further study and likely several follow-up works. We will highlight this finding in the paper and invite the community to explore the reasons behind this phenomenon.\n\n\n\n\n\n\n$\\quad$\n> No related work in the body; more comparison to other data mixing methods would be helpful\n\n In the camera-ready version, we will move our related work section from the appendix into the main body and expand it to include a more detailed comparison with other data mixing methods, as requested. The additional page provided for the camera-ready version will allow us to make these important improvements. We will also include more detailed discussions on related works (such as connection to DoGE, which we discussed above).\n\n\n\n$\\quad$\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\nThis is a great suggestion! In the revised version, we will add new columns to Table 1 to include this information and to quantitatively demonstrate the reduced performance disparity."}, {"review_id": "PiVD6Nrscm", "reviewer": "Reviewer_SmGd", "summary": "This 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.", "questions": "1. 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$?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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?", "was_revised": true, "final_justification": "I can confirm my support on accepting this submission. This is good and solid work. After the rebuttal, my biggest concerns were resolved and the few remaining are rather unimportant or context-dependent, let me explain:\n\n**1. Lack of empirical results for the idealized algorithm (weight 25%):** I was somewhat concerned that the authors derived an algorithm that has never been tested, but rather its approximations. The authors added experiments to compare a version of the algorithm that it is closer to the ideal one than before (still far, as it is naturally expensive), that I think solves this issue.\n\n**2. Batch-dependent variance estimators (weight 50%):** My biggest concern was that tasks with low-representation in the batch could get worse statistical estimates that lower their representation even further. After the rebuttal, authors clarified that every $T$ steps the variance of each task is estimated individually with an equal number of samples, clarifying this concern.\n\n**3. Assumptions on large-models and over-parametrization (weight 15%):** The premise of the paper heavily relies on having large models that can learn every task perfectly and with almost no-interaction between tasks. While this is said upfront, some results assumed it implicitly as acked by the authors (e.g. remark 1 and, I assume, also theorem 3.4). While this is resolved, I'd like it to be even more explicit in the final version.\n\n**4. Intuition about the hyperparameter $\\tau$ (weight 10%):** The algorithm provides a hyperparameter to balance the learning across tasks. While the authors recommend a value of 3 or 5, it is still a bit too ad hoc to my taste and the lack of interpretability for the values of $\\tau$ make it a bit hard to tune without some model selection. This is, however, a minor point.", "rebuttal": "Thank you for your time and for the thoughtful and constructive feedback. We are encouraged that you found the paper well-written, the problem important, and the empirical results convincing. We especially appreciate your recognition of our contribution on the new view of data mixing as a multi-task learning problem, and our validation of the low-gradient conflict hypothesis in large models. Below, we respond to your specific questions and suggestions in detail.\n\n\n\n$\\quad$\n> I would have liked to see the performance of the idealized version (or the one which updates every X iterations) to see its drawbacks.\n\nThis is a great question. During the rebuttal period, we ran additional experiments on the GLaM dataset to study the effect of the PiKE update interval, $T_0$. As shown in Table 1 below, more frequent updates tend to improve performance.  However, this benefit comes at an additional computational cost. While our original setting of $T_0=1000$ incurs a negligible ( $\\\\sim 1 \\\\%$ ) training overhead, shorter intervals of $T_0=500$ and $T_0=100$ increase this overhead to $\\\\sim 2 \\\\%$ and $\\\\sim 10 \\\\%$, respectively. We think that a reasonable  update interval like $T_0=1000$ balances an effective trade-off between model performance and computational efficiency. If the paper is accepted, we will include this additional table in the final version.\n\nTable 1. Pre-training 110M GPT-2 style LMs with context length 2048 for 20K steps using GLaM datasets. \n| Method | Mean Accuracy | ArcE | CSQA | HellaSwag | PIQA |\n|:---:|:---:|:---:|:---:|:---:|:---:|\n| Mix | 37.69 | 38.35 | 25.80 | 28.50 | 58.11 |\n| Round_Robin | 36.47 | 37.17 | 22.77 | 27.57 | 58.38 |\n| Random | 36.07 | 37.55 | 21.79 | 27.23 | 57.73 |\n| GLaM | 37.10 | 37.89 | 24.90 | 28.17 | 57.45 |\n| DoReMi | 38.17 | 39.49 | 26.21 | 28.50 | 58.49 |\n| PiKE (Update Interval - 100)  | 38.32 | 40.85 | 27.60 | 28.47 | 56.37 |\n| PKE (Update Interval - 500) | 38.67 | 40.68 | 26.45 | 28.80 | 58.76 |\n| PiKE (Update Interval - 1000)  | 38.23 | 39.37 | 27.35 | 28.30 | 57.89 |\n\n\n\n$\\quad$\n> 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.\n\nWe agree that computational overhead is an important consideration. PiKE is designed to be computationally efficient. In particular, PiKE does not require storing full per-task gradients. Instead, for each task, we only store two scalar values: the gradient norm and its variance. We also reported the overhead from estimating these statistics during pre-training: In our experiments, this step adds at most $2.4\\\\%$ to total training time, which we find to be relatively negligible.\n\nYour question about the impact on tasks with low representation is very important. To answer it, we would like to clarify one detail about PiKE: the variance estimation is done once every $T_0$ steps, and we perform it over tasks one at a time. In particular, for each task $k$, we select a relatively large batch (we used 256 for all experiments) and use it to estimate the gradient norm and variance. So even if a task has low representation during regular training, it still gets a large batch for variance estimation every $T_0$ steps. Moreover, all tasks use the same batch size for the variance estimation step. This is an important detail, and we will clarify it further in the revised version of the paper.\n\n\n\n$\\quad$\n> 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\nOur method doesn't require pretraining the weights. The batch mixture is computed dynamically on the fly during training. For the hyperparameter $\\tau$, we recommend setting $\\tau=3$ or $5$. Our experiments consistently show these values deliver strong and balanced performance.\n\n\n\n$\\quad$\n> How do you derive the mirror descent step in section 3.2?\n\nHere is more detailed description on how we derive the mirror descent step. Starting from equation (9), we solve the following constrained optimization:\n\n$$\n    \\min_{w_1,\\ldots,w_K\\geq0}  \\sum^K_{k=1}  w_k \\lambda_k + \\frac{1}{2} w_k^2\\kappa_k \n     \\textrm{\\ \\ \\   s.t.} \n    \\sum_{k=1}^Kw_k = 1\n$$\nwhere $\\lambda_k\\triangleq -\\eta\\beta\\\\|\\nabla\\mathcal{L}_k(\\theta)\\\\|^2 + \\frac{L\\eta^2}{2b} \\sigma_k^2$ and $\\kappa_k\\triangleq L\\eta^2\\gamma\\\\|\\nabla\\mathcal{L}_k(\\theta)\\\\|^2$. \n\nThis can be written more explicitly as:\n\n$$\n    \\\\min_{w_1, \\\\ldots ,w_K \\\\geq0}  \\\\sum^K_{k=1}  \\\\left(  - \\\\eta \\\\beta   |\\\\nabla\\mathcal{L}\\_k(\\\\theta)|^2 + \\\\frac{L \\\\eta^2}{2b} \\\\sigma_k^2 \\\\right)  w_k  + \\frac{1}{2} \\\\left(L\\\\eta^2\\\\gamma\\\\|\\\\nabla\\\\mathcal{L}\\_k(\\\\theta)\\\\|^2\\\\right) w\\_k^2 \n     \\\\textrm{\\\\ \\\\ \\\\ s.t.} \n    \\\\sum_{k=1}^Kw_k = 1\n$$\nTo apply mirror descent over the simplex, we need to use a multiplicative update followed by normalization. Taking gradient with respect to w_k, we have  \n\\begin{equation}\ng_{w_k} = -\\\\eta(\\\\beta - L \\\\eta \\\\gamma w_k \\\\| \\\\nabla \\\\mathcal{L}_k(\\\\theta) \\\\| ^2) + \\\\frac{L \\\\eta^2 }{2b} \\sigma_k^2 \\\\\n\\end{equation}\n\nLet $\\alpha$ be our mirror descent learning rate. The update rule becomes\n\n\\begin{equation}\n\\begin{split}\n    w_k &\\leftarrow w_k \\exp \\left(\\alpha \\eta\\left(\\beta-L \\eta \\gamma w_k\\right)\\left\\|\\nabla \\mathcal{L}_k(\\theta)\\right\\|^2-\\frac{\\alpha L \\eta^2}{2 b}  \\sigma_k^2 \\right) \\\\ \n\\end{split}\n\\end{equation}\nfollowed by the normalization $\\mathbf{w} \\leftarrow \\mathbf{w} /\\|\\mathbf{w}\\|_1$. \n\n\n\n\n\n\n$\\quad$\n> Is there any assumption regarding the capacity of the neural network to hold remark 1?\n\n\n\nYou are correct that Remark 1 implicitly relies on an assumption about the model’s capacity. More precisely, we assume the model is in the over-parameterized regime, with enough capacity to learn all tasks jointly without representational bottlenecks. In particular, we assume that small gradients imply small loss values. This matches our experimental setting, which uses large-scale modern networks that are, by design, highly over-parameterized\n\n\nWe will revise Remark 1 to make this assumption explicit. Thank you for highlighting the need for clarity here.\n\n\n\n\n$\\quad$\n> 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.\n\nCorrect! The per-task gradients are indeed nearly orthogonal in practice, which one of the core assumption behind our method's design. As illustrated in Figures 5 and 6, we measured the cosine similarity across different task gradients for both our multilingual and GLaM pre-training experiments. We found that the most values are consistently low, with a empirical upper bound of 0.05.\n\n\n\n$\\quad$\n> How is the bound in Theorem 3.4 tight (L164)? Isn't this bound some form of a truncated Taylor expansion?\n\nTo demonstrate that this bound is tight (in the minimax sense), we must show that there exists a case where the inequality is met with equality. As we state on L163-164, we constructed such an example in Theorem H.4 (Appendix H). In that section, we define a specific loss function for which the inequality in our bound becomes an exact equality. Therefore, by definition, the bound is tight for this example. \n\n\n$\\quad$\n> Do you have any intuition on how to select $\\tau$ ?\n\nIntuitively the hyperparameter $\\tau$ allows users to navigate the trade-off between optimizing for average performance versus worst-case performance (i.e., fair learning across tasks). For the hyperparameter $\\tau$, we recommend setting $\\tau=3$ or $5$. Our experiments consistently show these values deliver strong and balanced performance."}]

Author rebuttal:

Meta-review: The paper proposes a multi task learning algorithm based on adaptive data mixing, motivated by the observation that gradient conflicts are rare in pretraining. The method connects gradient norm and variance to mixture weight updates via mirror descent, with a balanced variant to ensure fairness across tasks. It is supported by solid theoretical derivations and convergence guarantees, and experimental results demonstrate clear benefits in multilingual and multi domain settings. Reviewers were mostly positive, and while there were concerns about applicability to fine tuning and sensitivity to hyperparameters, they were sufficiently addressed through additional experiments during the rebuttal period. The applicability of this method to LLM training, together with the motivation that it leverages nonconflicting gradients, represents a meaningful advancement in multi task learning, and thus the paper is a clear accept.
Decision: Accept (spotlight)