Pith. sign in

REVIEW 3 major objections 4 minor 77 references

PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts

T0 review · 3 major / 4 minor · reviewed 2026-08-08 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Adaptive data mixing that reweights tasks by gradient size and variance speeds up LLM pretraining when task gradients barely conflict.

desk verdict PiKE has a genuinely new adaptive data-mixing rule with consistently positive empirical results, but the convergence guarantee is proved for the conceptual version, not the simplified algorithm that produced the numbers. read the letter →

arxiv 2502.06244 v2 pith:VA2EYTMY submitted 2025-02-10 cs.LG

Zeman Li , Yuan Deng , Peilin Zhong , Meisam Razaviyayn , Vahab Mirrokni This is my paper · ORCID
classification cs.LG
keywords datamixingmulti-tasklearninggradientconflictsLLMpretrainingadaptivesamplingvariancetiltedempiricalriskminimizationconvergenceanalysis
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Large language models are pretrained on mixtures of domains, but standard practice fixes the mixture by heuristics or proxy models. This paper makes the case that in large-scale pretraining, domains' gradients rarely conflict, and that this fact can be exploited directly: an online algorithm that reweights domains by gradient magnitude and variance can maximize the per-step loss decrease. The proposed PiKE does exactly that, with a theoretical guarantee that all task gradients shrink at the optimal O(1/$δ^{2}$) rate for smooth nonconvex optimization. Empirically, PiKE reaches higher downstream accuracy and faster convergence than static and proxy-based mixtures on multilingual and multi-domain pretraining at 110M–1B scale, and Balanced-PiKE trades a small amount of average loss for more even progress across tasks.

What carries the argument

The load-bearing object is the one-step descent bound of Theorem 3.4, which controls the expected decrease of $L(\theta)=\sum_k L_k(\theta)$ under gradients that are $c$-conflicted, meaning pairwise inner products satisfy $\langle \nabla L_j,\nabla L_k\rangle \ge -c(\|\nabla L_j\|^2+\|\nabla L_k\|^2)$, and $\bar c$-aligned. Optimizing the bound over sampling weights $(w_1,\dots,w_K)$ gives the conceptual solution $w_k^*=\max(0,-(\mu+\lambda_k)/\kappa_k)$ via KKT conditions. The practical algorithm replaces exact optimization by a single mirror-descent step $w_k\leftarrow w_k\exp(\zeta_1\|\nabla L_k(\theta)\|^2-\tfrac{\zeta_2}{2b}\sigma_k^2)$, followed by renormalization. Balanced-PiKE relies on the tilted objective $\tfrac1\tau\log\sum_k e^{\tau L_k(\theta)}$ and its Fenchel-dual reformulation, whose inner maximizer $y_k^*\propto e^{\tau L_k(\theta)}$ supplies fairness-adjusted weights.

What would settle it

Re-run the 750M GLaM pretraining while logging the c-conflict constant c(t) from Definition 3.1 at every estimation step; if c(t) ever reaches 1/(K-2+b/b_k) for some task k, Theorem 3.4's bound ceases to apply, and PiKE should lose its edge over a static Mix in exactly those intervals.

Watch

Extended reading notes

Core claim

PiKE's central claim is that in the low-conflict regime typical of large-scale pretraining, the per-step expected decrease of $L(\theta)$ has a near-tight upper bound given by $\sum_k b_k(-\tfrac{\eta\beta}{b}\|\nabla L_k(\theta)\|^2+\tfrac{L\eta^2}{2b^2}\sigma_k^2)+\sum_k b_k^2\tfrac{L\eta^2}{2b^2}\gamma\|\nabla L_k(\theta)\|^2$, with $\beta$ and $\gamma$ determined by the conflict and alignment constants. Minimizing this bound at every step reweights tasks toward those with large gradient norms and small noise, and the resulting adaptive mixture drives all per-task gradients to small norm at rate $O(1/\delta^2)$. The paper verifies the premise by measuring mostly positive gradient cosine similarities on mC4 and GLaM runs up to 1B parameters, and reports consistently higher downstream accuracy than static mixtures, GLaM and DoReMi proxy weights, FAMO, and ADO, with 1.2–2.4% training-time overhead.

Load-bearing premise

The theory assumes that per-task gradients stay only mildly conflicting throughout training, so the descent coefficient β remains positive; the paper measures this on its own runs but does not prove it holds for other tasks or later training stages.

Editorial extensions

If this is right

  • Static or proxy-tuned mixtures can be replaced by online reweighting from gradient statistics, removing the need for separate small proxy models in data mixing.
  • Given low conflict, PiKE's iteration complexity is O(1/δ²) to drive every task's gradient small, matching the optimal order for smooth nonconvex stochastic optimization and improving on uniform Mix by a factor related to K and the noise level.
  • PiKE's overhead is dominated by estimating K scalar statistics every T0 steps, about 1.2–2.4% of training time at 110M–1B scale, so the method stays practical as model size and task count grow.
  • The variance term keeps noisy tasks from being oversampled; removing it drops mean accuracy, e.g., from 47.0% to 45.3% in the 750M GLaM ablation.
  • Balanced-PiKE interpolates between average-loss and worst-case-loss objectives, reducing task performance gaps at moderate τ without sacrificing overall accuracy.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • If gradient conflict grows later in training—for example, when tasks become adversarial or when fine-tuning on conflicting objectives—the descent coefficient β can turn non-positive and PiKE's theoretical rationale would no longer apply; a conflict monitor could gate the update.
  • Because PiKE only needs per-task gradient norms and variances, the same reweighting rule could be applied to fine-grained data domains, modalities, or evolving data streams, not just a handful of predefined tasks.
  • The near-tight bound suggests that in low-conflict regimes the objective decrease is essentially determined by reweighted gradient norms and noise; one testable extension is whether PiKE improves data efficiency under a fixed compute budget by reaching a target loss with fewer tokens.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 4 minor

Summary. The paper studies adaptive data mixing for multi-task LLM pretraining. It observes that gradients across languages/domains are often positively correlated rather than conflicting, and proposes PiKE, which periodically estimates per-task gradient squared norms and variances and reweights the sampling distribution accordingly. A conceptual version derives the optimal weights from an upper bound on the one-step expected loss decrease under low-conflict and bounded-alignment assumptions; the practical version, Algorithm 1, replaces the KKT minimizer with a single exponential mirror-descent step with two tunable constants, ζ1 and ζ2. A Balanced-PiKE variant based on tilted empirical risk minimization is also introduced. Experiments on mC4 and GLaM with models from 110M to 1B parameters report consistent downstream accuracy gains over Mix, FAMO, ADO, DoReMi, and static GLaM weights.

Significance. If the reported results hold, the contribution is practically valuable: adaptive reweighting from cheap per-task gradient statistics is a scalable alternative to proxy-model mixing, with O(K) memory overhead and 1.2-2.4% training-time overhead measured in Table 8. The low-conflict observation on 750M and 1B decoder-only models extends prior evidence from multilingual BERT. The conceptual per-step bound (Theorem 3.4) and the tightness example (Theorem H.4) are clean, the variance-term ablation (Table 11) supports a non-obvious design choice, and the evaluation spans multiple scales and baselines. The main weakness is that the theorem applies to the conceptual Algorithm 2, not to the Algorithm 1 whose results are reported; as presented, the theoretical guarantees do not explain the empirical method.

major comments (3)
  1. [§3.2, Algorithm 1; §3.1, Algorithm 2; Theorem 3.5] The paper's abstract and Key Feature 2 claim that PiKE provides convergence guarantees, but the guarantees apply to Algorithm 2, not to the algorithm evaluated in Section 4. Theorem 3.5 (and Appendix H, Theorem H.5) analyzes Algorithm 2, which solves the KKT system in Eq. (14) exactly at every iteration. Algorithm 1, the version used for all reported results, replaces the KKT minimizer with a single exponential mirror-descent step, updates weights only every T0=1,000 steps, substitutes tuned constants ζ1 and ζ2 for the coefficients β, L, γ, and is run with AdamW. No theorem in Appendix H provides a descent or convergence bound for Algorithm 1's update. The authors should either prove a guarantee for Algorithm 1 under explicit assumptions on ζ1, ζ2, and T0, or explicitly state that the theoretical guarantees apply only to the conceptual version and present Algorithm 1 as a heuristic approximation supported by the empirical evaluation.
  2. [Definition 3.2; Lemma H.1; Theorem H.3] Definition 3.2 is printed with squared norms on the right-hand side: ⟨∇Lj(θ),∇Lk(θ)⟩ ≤ c̄∥∇Lj(θ)∥²∥∇Lk(θ)∥². This inequality is dimensionally inconsistent and is not the inequality used anywhere in the proofs. In the second part of Lemma H.1 and in steps (c)-(d) of Eq. (18) in Theorem H.3, the argument instead applies the definition as c̄∥∇Lj(θ)∥∥∇Lk(θ)∥, which is what yields the stated γ=1+c̄(K-1) after the Cauchy-Schwarz manipulation. The definition should be corrected to ⟨∇Lj(θ),∇Lk(θ)⟩ ≤ c̄∥∇Lj(θ)∥∥∇Lk(θ)∥, and the proof and all subsequent statements should be made consistent with this version.
  3. [Definitions 3.1-3.2 and Theorem 3.5; Figures 5-6] The convergence guarantee is conditional on the gradients being c-conflicted and c̄-aligned with c < 1/(K-2+b/b_k) at every iterate. The paper provides retrospective measurements from its own training runs as evidence, but there is no certificate or monitoring procedure that verifies this condition during training, and the condition depends on the evolving weights through b/b_k. If the condition fails mid-training, β in Eq. (8) can become non-positive and the descent bound gives no information. The paper should state this limitation explicitly and propose a practical check or safeguard that can detect when the theoretical regime is being left; as written, the theorem's applicability to new or larger task mixtures is an unverified assumption.
minor comments (4)
  1. [Theorem 3.5 and Theorem H.5] The step-size condition is internally inconsistent: the main text first states η ≤ βδ/(Lσ²max/b + Lγδ), but the following 'Particularly' sentence writes η = βδ/(Lσ²max/b + Lηδ), and the appendix restatement repeats Lηδ in the denominator. The proof between Eq. (22) and Eq. (23) requires the denominator with Lγδ; please correct all occurrences and re-verify the resulting iteration count.
  2. [Section 2 headings] The headings 'Bulding Block #1', 'Bulding Block #2', and 'Bulding Block #3' contain a repeated spelling error and should read 'Building Block'.
  3. [Appendix H.3] The text near Eq. (28) contains a sentence fragment, 'By summing u Rearranging the terms, we obtain:', which should be rewritten as a complete sentence.
  4. [Appendix E, Algorithm 3] The expression for y⋆k in Appendix E writes the exponent as τLk(θ)-1 and τLk(θ)-1 in numerator and denominator, while Lemma 3.6 uses τLk(θ) without the -1. Although the -1 cancels, the notation should be aligned between the lemma, the appendix, and Algorithm 3 to avoid confusion.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the derivation is self-contained and the empirical claims are not constructed from their own outputs.

full rationale

The paper's central derivation chain is not circular. The conceptual PiKE update is obtained by minimizing the right-hand side of the one-step SGD descent bound in Theorem 3.4, where the quantities entering the bound (per-task gradient norms and variances) are estimated from the training dynamics and are not themselves defined in terms of the claimed downstream improvements. The convergence guarantee in Theorem 3.5 is a standard telescoping descent argument under the stated c-conflicted and c-bar-aligned assumptions; it does not assume the empirical conclusion that PiKE achieves better accuracy. The practical Algorithm 1 replaces the exact KKT minimizer of the conceptual version with a single mirror-descent step and two tunable constants zeta_1 and zeta_2. This is an approximation that creates a gap between the theory (which covers Algorithm 2) and the experiments (which use Algorithm 1), but it is not circularity: the tunable constants are hyperparameters selected by grid search, not fitted targets that are then reported as predictions. The low-conflict observation is an empirical measurement used to motivate Assumptions 3.1 and 3.2, and the theoretical results explicitly state those assumptions rather than silently embedding the conclusion. Balanced-PiKE reuses the known Fenchel-dual representation of tilted ERM / alpha-fairness, and this is cited to prior work rather than presented as a new unification. No load-bearing self-citation chain was found. The main correctness caveat is that the reported experiments are not covered by the formal convergence proof because Algorithm 1 is a simplified variant, but that concern is a scope issue, not a circularity issue.

Assumptions & free parameters 4 free parameters · 4 assumptions · 0 invented entities

The central derivation relies on standard smoothness and bounded-variance assumptions (Assumption 3.3), the empirically-motivated smallness of gradient conflict and alignment parameters, and the choice of Mix batching. The practical algorithm introduces tunable coefficients zeta1 and zeta2, interval T0, and the Balanced variant adds temperature tau. No new physical entities are postulated. The main fragility is that the low-conflict condition is measured, not guaranteed.

free parameters (4)
  • zeta1 = 0.01 to 0.15 (grid)
    Coefficient on per-task gradient norm squared in the weight update (Algorithm 1, line 6); tuned via grid search in Table 7.
  • zeta2 = 0.001 to 0.01 (grid)
    Coefficient on per-task gradient variance; tuned via grid search; ablation in Table 11 shows sensitivity.
  • tau = 1, 3, 5
    Fairness temperature in Balanced-PiKE; user-selected, not fitted from data.
  • T0 = 1000
    Interval for updating mixing weights; set to 1000 in all experiments; affects overhead and adaptation speed.
assumptions (4)
  • domain assumption Per-task losses are L-smooth with unbiased gradients and bounded variance sigma_k^2 (Assumption 3.3).
    Standard smoothness and variance assumption for stochastic optimization; invoked in Theorem 3.4 and all convergence results.
  • domain assumption Task gradients are c-conflicted and c-bar-aligned with small c (Definitions 3.1 and 3.2), with c < 1/(K-2+b/b_k).
    Empirically measured in Section 2 (Figures 2, 5, 6); required for beta > 0 in Theorem 3.4 and for the bound to be a descent guarantee.
  • domain assumption The Mix batching strategy (all tasks in each batch) is the appropriate base strategy.
    Justified empirically in Figure 2 and Figure 4; the analysis and algorithm are built on Mix.
  • standard math The Fenchel-dual representation of tilted ERM (Lemma 3.6) is valid for L_k > 0.
    Known duality result; used to derive Balanced-PiKE.

how reviews work

0 comments
Cite this review

Pith. "Pith review of PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts." pith.science (2026). https://pith.science/paper/VA2EYTMY

@misc{pith2026250206244,
  author       = {Pith},
  title        = {Pith review of: PiKE: Adaptive Data Mixing for Large-Scale Multi-Task Learning Under Low Gradient Conflicts},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/VA2EYTMY}},
  note         = {Machine review of arXiv:2502.06244}
}
read the original 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 multitask learning (MTL) perspective While prior work in MTL has emphasized mitigating gradient conflicts we observe that largescale pretraining scenariossuch as multilingual or multidomain trainingoften exhibit little to no gradient conflict Motivated by this observation we propose PiKE (Positive gradient interaction-based K-task weights Estimator) an adaptive data mixing algorithm that dynamically adjusts sampling weights during training PiKE exploits nonconflicting gradient interactions to minimize a neartight 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 nonadaptive mixing baselines Furthermore we extend PiKE to promote balanced learning across tasks Extensive experiments on largescale language model pretraining confirm that PiKE achieves faster convergence and improved downstream performance compared to existing approaches

Figures

Figures reproduced from arXiv: 2502.06244 by the authors.

Figure 1
Figure 1. 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.9× faster. Right (750M models, GLaM six domains): PiKE achieves a 3.4% higher accuracy over DoReMi [68]. PiKE scales efficiently with model size and number of tasks. Detailed results are in Tables 13 and 14. most are … view at source ↗
Figure 2
Figure 2. Left: Average accuracy on four downstream tasks (ArcE, CSQA, HellaSwag, PIQA) for 750M GPT-2 style models pre-trained with Mix, Round-Robin, and Random strategies. See Appendix G.1 for more. Right: Task gradient cosine similarity for a 750M GPT-2 style model pre-trained on GLaM datasets. “data1-data2” indicates gradient similarity between tasks data1 and data2. Further results in Appendix G.2. L-Lipschitz if ∥h(θ) −… view at source ↗
Figure 3
Figure 3. Adaptive mixing consistently outper￾forms static mixing in the Example 2.1. While such conflicts are well-documented in vision tasks, we observe they are rare in large-scale lan￾guage model pretraining. In particular, as shown in Figures 2, 5, and 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-co… view at source ↗
Figures from the paper (5 more)
Figure 4
Figure 4. Figure 4: Average downstream task accuracy of pretraining language models using Mix, Round [PITH_FULL_IMAGE:figures/full_fig_p020_4.png]
Figure 5
Figure 5. Figure 5: 1B models trained on multilingual mC4 datasets. [PITH_FULL_IMAGE:figures/full_fig_p020_5.png]
Figure 6
Figure 6. Figure 6: 750M models on GLaM datasets with six domains. [PITH_FULL_IMAGE:figures/full_fig_p020_6.png]
Figure 7
Figure 7. Figure 7: Eval perplexity of pretraining 270M GPT-2 style multilingual language models on mC4 [PITH_FULL_IMAGE:figures/full_fig_p021_7.png]
Figure 8
Figure 8. Figure 8: The sampling weights for each dataset during the pre-training of 1B GPT-2-style multilin [PITH_FULL_IMAGE:figures/full_fig_p025_8.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

77 extracted references · 32 canonical work pages

  1. [1]

    Abbas, K

    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:2303.09540, 2023

    arXiv 2023

  2. [2]

    J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016

    arXiv 2016

  3. [3]

    Ban and K

    H. Ban and K. Ji. Fair resource allocation in multi-task learning. arXiv preprint arXiv:2402.15638, 2024

    arXiv 2024

  4. [4]

    Bengio, J

    Y . Bengio, J. Louradour, R. Collobert, and J. Weston. Curriculum learning. InProceedings of the 26th annual international conference on machine learning, pages 41–48, 2009

    work page 2009

  5. [5]

    Y . Bisk, R. Zellers, R. Le Bras, J. Gao, and Y . Choi. Reasoning about physical commonsense in natural language, 2019

    work page 2019

  6. [6]

    Bradbury, R

    J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax

    work page 2018

  7. [7]

    Brown, B

    T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020

    1901

  8. [8]

    Z. Chen, V . Badrinarayanan, C.-Y . Lee, and A. Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International conference on machine learning, pages 794–803. PMLR, 2018

    work page 2018

Show all 77 references
  1. [9]

    Chowdhery, S

    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

    2022 arXiv

  2. [10]

    Chowdhery, S

    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.Journal of Machine Learning Research, 24(240):1–113, 2023

    2023

  3. [11]

    Clark, I

    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

    2018 arXiv

  4. [12]

    Dac Lai, C

    V . Dac Lai, C. Van Nguyen, N. T. Ngo, T. Nguyen, F. Dernoncourt, R. A. Rossi, and T. 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

    2023

  5. [13]

    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

    2016

  6. [14]

    Dehghani, J

    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.05442, 2023. 10

    2023 arXiv

  7. [15]

    Désidéri

    J.-A. Désidéri. Multiple-gradient descent algorithm (mgda) for multiobjective optimization. Comptes Rendus Mathematique, 350(5-6):313–318, 2012

    2012

  8. [16]

    J. Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018

    2018 arXiv

  9. [17]

    N. Du, Y . Huang, A. M. Dai, S. Tong, D. Lepikhin, Y . Xu, M. Krikun, Y . Zhou, A. W. Yu, O. Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning, pages 5547–5569. PMLR, 2022

    2022

  10. [18]

    Dubey, A

    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 preprint arXiv:2407.21783, 2024

    2024 arXiv

  11. [19]

    Gaffney, D

    C. Gaffney, D. Li, R. Sang, A. Jain, and H. Hu. Orbax, 2023. URL http://github.com/ google/orbax

    2023

  12. [20]

    L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020

    2020 arXiv

  13. [21]

    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

    2025 arXiv

  14. [22]

    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 arXiv:2405.14908, 2024

    2024 arXiv

  15. [23]

    Grain - feeding jax models, 2023

    Google. Grain - feeding jax models, 2023. URL http://github.com/google/grain

    2023

  16. [24]

    Y . Gu, O. Tafjord, B. Kuehl, D. Haddad, J. Dodge, and H. Hajishirzi. Olmes: A standard for language model evaluations. arXiv preprint arXiv:2406.08446, 2024

    2024 arXiv

  17. [25]

    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

    2023 arXiv

  18. [26]

    J. 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://github.com/google/ flax

    2023

  19. [27]

    Hoffmann, S

    J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. 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

    2022

  20. [28]

    Jiang, A

    Y . Jiang, A. Zhou, Z. Feng, S. Malladi, and J. Z. Kolter. Adaptive data optimization: Dynamic sample selection with scaling laws. arXiv preprint arXiv:2410.11820, 2024

    2024 arXiv

  21. [29]

    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 a tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pages 1–12, 2017

    2017

  22. [30]

    Laurençon, L

    H. Laurençon, L. Saulnier, T. Wang, C. Akiki, A. Villanova del Moral, T. Le Scao, L. V on Werra, C. Mou, E. González Ponferrada, H. Nguyen, et al. The bigscience roots corpus: A 1.6 tb composite multilingual dataset. Advances in Neural Information Processing Systems, 35: 31809...

    2022

  23. [31]

    K. Lee, D. Ippolito, A. Nystrom, C. Zhang, D. Eck, C. Callison-Burch, and N. Carlini. Dedupli- cating training data makes language models better. arXiv preprint arXiv:2107.06499, 2021

    2021 arXiv

  24. [32]

    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/1804.08838. 11

    2018 arXiv

  25. [33]

    T. Li, A. Beirami, M. Sanjabi, and V . Smith. Tilted empirical risk minimization.arXiv preprint arXiv:2007.01162, 2020

    2007 arXiv

  26. [34]

    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

    2024 arXiv

  27. [35]

    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, 34:18878–18890, 2021

    2021

  28. [36]

    B. Liu, Y . Feng, P. Stone, and Q. Liu. Famo: Fast adaptive multitask optimization.Advances in Neural Information Processing Systems, 36, 2024

    2024

  29. [37]

    H. Liu, S. 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

    2023

  30. [38]

    Q. Liu, X. Zheng, N. Muennighoff, G. Zeng, L. Dou, T. Pang, J. Jiang, and M. Lin. Regmix: Data mixture as regression for language model pre-training. arXiv preprint arXiv:2407.01492, 2024

    2024 arXiv

  31. [39]

    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

    2015

  32. [40]

    X. Liu, P. He, W. Chen, and J. Gao. Multi-task deep neural networks for natural language understanding. arXiv preprint arXiv:1901.11504, 2019

    1901 arXiv

  33. [41]

    Loshchilov and F

    I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Con- ference on Learning Representations (ICLR), 2019. https://openreview.net/forum?id= Bkg6RiCqY7

    2019

  34. [42]

    Luong, Q

    M.-T. Luong, Q. V . Le, I. Sutskever, O. Vinyals, and L. Kaiser. Multi-task sequence to sequence learning. arXiv preprint arXiv:1511.06114, 2015

    2015 arXiv

  35. [43]

    Misra, A

    I. Misra, A. Shrivastava, A. Gupta, and M. Hebert. Cross-stitch networks for multi-task learning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 3994–4003, 2016

    2016

  36. [44]

    Mo and J

    J. Mo and J. Walrand. Fair end-to-end window-based congestion control. IEEE/ACM Transac- tions on networking, 8(5):556–567, 2000

    2000

  37. [45]

    Navon, A

    A. Navon, A. Shamsian, I. Achituve, H. Maron, K. Kawaguchi, G. Chechik, and E. Fetaya. Multi-task learning as a bargaining game. arXiv preprint arXiv:2202.01017, 2022

    2022 arXiv

  38. [46]

    Penedo, Q

    G. Penedo, Q. Malartic, D. Hesslow, R. Cojocaru, H. Alobeidli, A. Cappelli, B. Pannier, E. Almazrouei, and J. Launay. The refinedweb dataset for falcon llm: Outperforming curated corpora with web data only. Advances in Neural Information Processing Systems, 36:79155– 79172, 2023

    2023

  39. [47]

    Radford, J

    A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language Models are Unsupervised Multitask Learners, 2019. https://openai.com/blog/ better-language-models/

    2019

  40. [48]

    J. W. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoffmann, F. Song, J. Aslanides, S. Henderson, R. Ring, S. Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446, 2021

    2021 arXiv

  41. [49]

    Raffel, N

    C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y . Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1–67, 2020

    2020

  42. [50]

    J. Ren, S. Rajbhandari, R. Y . Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y . He. {ZeRO-Offload}: Democratizing {Billion-Scale} model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pages 551–564, 2021. 12

    2021

  43. [51]

    R. T. Rockafellar. Convex analysis:(pms-28). Princeton university press, 2015

    2015

  44. [52]

    Ruder, J

    S. Ruder, 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

    2019

  45. [53]

    Sachdeva, B

    N. Sachdeva, B. Coleman, W.-C. Kang, J. Ni, L. Hong, E. H. Chi, J. Caverlee, J. McAuley, and D. Z. Cheng. How to train data-efficient llms. arXiv preprint arXiv:2402.09668, 2024

    2024 arXiv

  46. [54]

    Sener and V

    O. Sener and V . Koltun. Multi-task learning as multi-objective optimization.Advances in neural information processing systems, 31, 2018

    2018

  47. [55]

    Soldaini, R

    L. Soldaini, R. Kinney, A. Bhagia, D. Schwenk, D. Atkinson, R. Authur, B. Bogin, K. Chandu, J. Dumas, Y . Elazar, et al. Dolma: An open corpus of three trillion tokens for language model pretraining research. arXiv preprint arXiv:2402.00159, 2024

    2024 arXiv

  48. [56]

    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

    2021 arXiv

  49. [57]

    Talmor, J

    A. Talmor, J. Herzig, N. Lourie, and J. Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge. arXiv preprint arXiv:1811.00937, 2018

    2018 arXiv

  50. [58]

    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 arXiv:2109.10686, 2021

    2021 arXiv

  51. [59]

    G. Team, R. Anil, S. Borgeaud, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, K. Millican, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023

    2023 arXiv

  52. [60]

    G. Team, P. Georgiev, V . 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 arXiv:2403.05530, 2024

    2024 arXiv

  53. [61]

    Vandenhende, S

    S. Vandenhende, S. 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–3633, 2021

    2021

  54. [62]

    H. V . V o, 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-supervised learning: A clustering-based approach. arXiv preprint arXiv:2405.15613, 2024

    2024 arXiv

  55. [63]

    A. Wang, Y . Pruksachatkun, N. Nangia, A. Singh, J. 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

    2019

  56. [64]

    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

    2010 arXiv

  57. [65]

    Wettig, A

    A. Wettig, A. Gupta, S. Malik, and D. Chen. Qurating: Selecting high-quality data for training language models. arXiv preprint arXiv:2402.09739, 2024

    2024 arXiv

  58. [66]

    Wortsman, P

    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 preprint arXiv:2309.14322, 2023

    2023 arXiv

  59. [67]

    M. Xia, T. Gao, Z. Zeng, and D. Chen. Sheared llama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06694, 2023

    2023 arXiv

  60. [68]

    S. M. Xie, H. Pham, X. Dong, N. Du, H. Liu, Y . Lu, P. S. Liang, Q. V . Le, T. Ma, and A. W. Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2024. 13

    2024

  61. [69]

    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

    2022

  62. [70]

    L. Xue. mt5: A massively multilingual pre-trained text-to-text transformer. arXiv preprint arXiv:2010.11934, 2020

    2010 arXiv

  63. [71]

    L. Xue, A. Barua, N. Constant, R. Al-Rfou, S. Narang, M. Kale, A. Roberts, and C. Raffel. Byt5: Towards a token-free future with pre-trained byte-to-byte models. corr, abs/2105.13626. arXiv preprint arXiv:2105.13626, 2021

    2021 arXiv

  64. [72]

    E. Yang, J. Pan, X. Wang, H. Yu, L. Shen, X. Chen, L. Xiao, J. Jiang, and G. Guo. Adatask: A task-aware adaptive learning rate approach to multi-task learning. In Proceedings of the AAAI conference on artificial intelligence, volume 37, pages 10745–10753, 2023

    2023

  65. [73]

    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.16952, 2024

    2024 arXiv

  66. [74]

    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

    2020

  67. [75]

    Zellers, A

    R. Zellers, A. Holtzman, Y . Bisk, A. Farhadi, and Y . Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019

    1905 arXiv

  68. [76]

    Zhang, C

    Y . Zhang, C. Chen, T. Ding, Z. Li, R. Sun, and Z.-Q. Luo. Why transformers need adam: A hessian perspective. arXiv preprint arXiv:2402.16788, 2024. 14 A Related Work Data Curation and Selection. The effectiveness of language models heavily depends on the quality of the pre-tr...

    2024 arXiv

  69. [77]

    ratio,” which defined as ⟨Lj (θ),Lk(θ)⟩ ∥Lj (θ)∥2+∥Lk(θ)∥2 , between task gradients during language model pre-training over time. “ data1- data2

    Task gradients rarely conflict—multilingual cosine similarities seldom drop below -0.1, and GLaM domain gradients remain predominantly positive. D PiKE: Conceptual Version Here, we present the conceptual (basic) version of PiKE. As discussed in the main text, this approach lac...

Pith tools

Reviewed August 8, 2026 · model on record in the stance chip above.