Paper deep dive
Rethinking Data Mixing from the Perspective of Large Language Models
Yuanjian Xu, Tianze Sun, Changwei Xu, XinLong Zhao, Jianing Hao, Ran Chen, Yang Liu, Ruijie Xu, Stephen Chen, Guang Zhang
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 4/10/2026, 4:25:12 AM
Summary
The paper introduces DoGraph, a novel reweighting framework for large language model (LLM) training that addresses the challenge of optimal data mixing. By establishing a theoretical connection between gradient dynamics and domain distributions, the authors argue that domains should be defined by the model's internal gradient-based perception rather than human intuition. DoGraph uses random projection and K-means clustering in gradient space to dynamically partition data into domains and optimize their weights, demonstrating improved performance and generalization across various benchmarks compared to existing methods.
Entities (5)
Relation Signals (3)
DoGraph → optimizes → Data Mixing
confidence 95% · DoGraph, a reweighting framework that formulates data scheduling as a graph-constrained optimization problem.
DoGraph → improves → Generalization
confidence 92% · DoGraph consistently achieves competitive performance... improves both domain balance and generalization.
Gradient Dynamics → defines → Domain
confidence 90% · a domain should be defined from the model’s perspective... distributional differences are encoded in the geometry of gradients
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Data mixing strategy is essential for large language model (LLM) training. Empirical evidence shows that inappropriate strategies can significantly reduce generalization. Although recent methods have improved empirical performance, several fundamental questions remain open: what constitutes a domain, whether human and model perceptions of domains are aligned, and how domain weighting influences generalization. We address these questions by establishing formal connections between gradient dynamics and domain distributions, offering a theoretical framework that clarifies the role of domains in training dynamics. Building on this analysis, we introduce DoGraph, a reweighting framework that formulates data scheduling as a graph-constrained optimization problem. Extensive experiments on GPT-2 models of varying scales demonstrate that DoGraph consistently achieves competitive performance.
Tags
Links
- Source: https://arxiv.org/abs/2604.07963v1
- Canonical: https://arxiv.org/abs/2604.07963v1
Trouble viewing inline? Open PDF directly →
Full Text
47,141 characters extracted from source content.
Expand or collapse full text
Rethinking Data Mixing from the Perspective of Large Language Models Yuanjian Xu1, Tianze Sun3, Changwei Xu2, XinLong Zhao, Jianing Hao1, Ran Chen2, Yang Liu11footnotemark: 1, Ruijie Xu2, Stephen Chen2, Guang Zhang1, 1Hong Kong University of Science and Technology (Guangzhou) 2OpenCSG 3Harbin Institute of Technology University yxu085@connect, guangzhang@hkust-gz.edu.cn China Mining GroupCorresponding author. Abstract Data mixing strategy is essential for large language model (LLM) training. Empirical evidence shows that inappropriate strategies can significantly reduce generalization. Although recent methods have improved empirical performance, several fundamental questions remain open: what constitutes a domain, whether human and model perceptions of domains are aligned, and how domain weighting influences generalization. We address these questions by establishing formal connections between gradient dynamics and domain distributions, offering a theoretical framework that clarifies the role of domains in training dynamics. Building on this analysis, we introduce DoGraph, a reweighting framework that formulates data scheduling as a graph-constrained optimization problem. Extensive experiments on GPT-2 models of varying scales demonstrate that DoGraph consistently achieves competitive performance. Code and data are publicly available at https://anonymous.4open.science/r/Dograph-53B9. Rethinking Data Mixing from the Perspective of Large Language Models Yuanjian Xu1, Tianze Sun3, Changwei Xu2, XinLong Zhao†thanks: China Mining Group, Jianing Hao1, Ran Chen2, Yang Liu11footnotemark: 1, Ruijie Xu2, Stephen Chen2, Guang Zhang1,†thanks: Corresponding author. 1Hong Kong University of Science and Technology (Guangzhou) 2OpenCSG 3Harbin Institute of Technology University yxu085@connect, guangzhang@hkust-gz.edu.cn 1 Introduction Training data fundamentally determines the capability of large language models (LLMs) (Xu et al., 2023; Wettig et al., 2024; Albalak et al., ). However, domain distributions are imbalanced due to unequal data availability: web-scale corpora are abundant, whereas specialized domains remain scarce (Gao et al., 2021). This raises a key question—can we design a principled sampling strategy to mitigate such imbalance? Exhaustively searching over all possible sampling policies is infeasible, as LLM training is prohibitively expensive. To make progress, we must first answer: what does a “domain” truly mean for a LLM, and are human and model perceptions of domains aligned (Sun and others, 2025)? Prior data mixing studies have predominantly relied on domain definitions derived from human intuition. Existing approaches can be broadly categorized into two lines of work. The first derives heuristics from small- or medium-scale models and then scales them to LLMs (Liu and others, 2024; Ye et al., 2024; Fan et al., 2023; Xie et al., 2023); however, empirical evidence shows that scaling laws and domain sensitivities observed in small models do not transfer reliably to larger ones (Kang et al., 2024). The second directly performs data reweighting or optimization on LLMs, either at the sample or domain level (Sun and others, 2025; Sow et al., 2025), but often incurs prohibitive computational costs or relies on unrealistic assumptions. In this work, we argue that the optimization of LLMs continuously reshapes their domain perception, creating a mismatch between human-defined and model-internal representations (Bengio et al., 2013). Figure 1 visualizes this evolution: at initialization, samples from domains such as C4, Wikipedia, Book, and ArXiv form well-separated clusters, reflecting strong domain-specific biases. As training progresses, these clusters gradually merge into an approximately isotropic distribution, indicating that the model internalizes more domain-invariant linguistic structures (Power et al., 2022; Gao et al., 2019). Figure 1: PCA projections of gradient directions at different training epochs. Colors denote data domains (C4, Wikipedia, ArXiv, Book, etc.). Initially, gradients form distinct clusters, showing strong domain bias. Over time, they overlap, indicating that the model homogenizes its domain perception. Experiments use 20% of SlimPajama trained on GPT2-Mini. This evolving misalignment biases existing data mixing methods. To address it, we formally link domain distributions with gradient dynamics, showing how model-defined domains emerge during optimization (Koh and Liang, 2017; Fort et al., 2019). Building on this foundation, we propose DoGraph, which formulates domain scheduling as a graph-constrained reweighting problem. DoGraph models the model-perceived domains as graph nodes and learns their weights through optimization. Our main contributions are summarized as follows: 1) We theoretically establish a connection between domain distribution and gradient dynamics, and empirically validate the dynamic correction of domain representations during LLM training. 2) We propose DoGraph, a graph-constrained reweighting framework that formalizes domain scheduling as an optimization problem. DoGraph is strongly grounded in theoretical principles. 3) We conduct extensive experiments across diverse benchmarks, demonstrating consistent improvements in both performance and domain balance, which validate the competitiveness of our approach. 2 Methods In this section, we begin by redefining domains from a learning-theoretic perspective. We then establish their connection to gradients, showing that distributional differences are reflected in gradient geometry in Section 2.1. Finally, we build on this insight to propose the DoGraph. 2.1 Rethinking the Definition of Domain In NLP, the notion of domain has often been unclearly defined, particularly in the training corpora of LLMs, where such boundaries are increasingly blurred. Before developing concrete strategies for domain weighting, it is essential to first clarify what we mean by a domain. We argue that a domain should be defined from the model’s perspective, namely as the distribution of inputs it perceives, rather than from a human perspective. Definition 2.1. Let V be a finite vocabulary and =∗X=V^* the space of all token sequences. A domain is a probability space (,ℱ,PX)(X,F,P_X), where ℱF is the σ-algebra on X and PXP_X a probability measure. Two domains 1=(,ℱ,P1)D_1=(X,F,P_1) and 2=(,ℱ,P2)D_2=(X,F,P_2) are distinct iff P1≠P2P_1≠ P_2. We now formulate the definition of domain as stated in Definition 2.1. Each x∈x is a finite token sequence from V, with domains distinguished by the regions of X where their distributions PXP_X concentrate. In simple cases, such as distinguishing code from natural language, these regions are relatively easy to separate. However, in practice, many domains are much less clear-cut, with boundaries that overlap or gradually shift. Thus, domains differ through probability measures over the same space ∗V^* rather than through disjoint supports. Connection between Domain and Gradients A central question is whether domains can be inferred from observable data instead of being imposed a priori, as such assumptions inevitably risk introducing bias. Since each training sample affects learning only via its gradient, the model perceives not raw token frequencies but the geometry of gradient flows. To investigate this, we analyze a simplified self-attention structure in which the Transformer can be linearized, leading to a tractable correspondence between distributions and gradients. Theorem 2.2. Under the linearized Transformer setting, for any parameter block b∈V,Q,K,O,Wb∈\V,Q,K,O,W\ and two domains P1,P2P_1,P_2, the difference of expected gradients satisfies g¯b(P1)−g¯b(P2)=∫∇WbL(x,y;θ)(P1−P2)(dx,dy). g_b(P_1)- g_b(P_2)= _W_bL(x,y;θ)\,(P_1-P_2)(dx,dy). Moreover, this difference admits a kernel representation: ‖g¯b(P1)−g¯b(P2)‖2=MMDkb2(P1,P2),\| g_b(P_1)- g_b(P_2)\|^2=MMD_k_b^2(P_1,P_2), where kb(s,s′)=⟨gb(s),gb(s′)⟩k_b(s,s )= g_b(s),g_b(s ) is the gradient-induced kernel. Theorem 2.2 shows that distributional differences are encoded in the geometry of gradients, implying that domains can be compared through their gradient signatures rather than token-level statistics. From this perspective, a domain is defined by its expected gradient flow, and training can be understood as a continual refinement of the model’s perception of domains, with each update adjusting how distributions are represented in gradient space. 2.2 DoGraph We argue that data weighting should adapt to the model’s evolving perception of domains, rather than fixed human-defined boundaries. Building on this idea, we introduce the DoGraph, where each domain corresponds to a node in a graph. At every epoch, we collect per-sample gradients and project them into a low-dimensional subspace via random projection. Next, we apply K-means clustering in the projected gradient space to obtain model-centric partitions of the training distribution. This partition evolves over training, reflecting the changing geometry of gradients. Formally, let gi∈ℝdg_i ^d be the gradient of the i-th sample and G=[g1,…,gn]⊤∈ℝn×dG=[g_1,…,g_n] ^n× d. We apply a random projection matrix R∈ℝd×kR ^d× k with Rpq∼(0,1/k)R_pq (0,1/k), yielding g~i=R⊤gi,or G~=GR. g_i=R g_i, G=GR. By the Johnson–Lindenstrauss lemma, (1−ϵ)‖gi−gj‖22≤‖g~i−g~j‖22≤(1+ϵ)‖gi−gj‖22,(1-ε)\|g_i-g_j\|_2^2≤\| g_i- g_j\|_2^2≤(1+ε)\|g_i-g_j\|_2^2, ensuring that clustering in the projected space preserves the gradient geometry while reducing computational cost and noise. Clustering g~i\ g_i\ into m groups D1,…,Dm\D_1,…,D_m\, we compute each domain’s mean gradient as g¯j=1|Dj|∑i∈Djg~i. g_j= 1|D_j| _i∈ D_j g_i. To balance learning, we assign adaptive domain weights w=(w1,…,wm)w=(w_1,…,w_m) by solving minw∈Δm−1ℒopt(∑j=1mwjg¯j), _w∈ ^m-1\;L_opt\! ( _j=1^mw_j g_j ), where Δm−1 ^m-1 is the probability simplex. DoGraph Pipeline At each epoch, per-sample gradients are first extracted and projected into a low-dimensional subspace via random projection, then clustered into domains in the projected space. Domain mean gradients are aggregated through an optimization step that computes the optimal domain weights. The model parameters are updated with the weighted gradient, and the process repeats, allowing both the partition of domains and their relative importance to adapt continuously throughout training. The choice of the optimization objective ℒoptL_opt is discussed in the Appendix A.6. Algorithm 1 summarizes the overall procedure of DoGraph. More implementation details can be found in Appendix A.3. 3 Experiment Results In this section, we begin by outlining the experimental setup, after which we present the overall performance analysis. We further conduct a perplexity analysis and investigate how model scale influences the observed trends, with detailed results presented in Section 3.3 and Section 3.4. All main results in the paper are reported using the GPT-2 Medium. To isolate the effects of architecture and parameter scale, additional experiments with the LLaMA-1.1B model are deferred to the Appendix A.4. Sensitivity to hyperparameters and the choice of optimizer are analyzed in Appendix A.6 and Appendix A.7. Commonsense / Reasoning RC LM Avg Method HellaSwag PiQA OBQA COPA LogiQA WinoG SciQ ARC-E Lambada SlimPajama Uniform 26.1 55.5 11.7 58.0 25.7 49.9 49.0 31.4 11.6 35.4 Dynamic Loss-Based 26.6 56.8 13.8 59.0 29.8 50.1 53.3 31.7 13.4 37.2 DoReMi 26.4 55.7 12.2 59.0 27.2 49.9 53.3 32.3 12.7 36.5 DOGE 26.2 55.8 11.5 62.0 27.2 50.4 52.8 31.3 11.6 36.5 RegMix 26.1 55.6 13.2 60.0 23.7 50.0 46.6 31.7 14.0 35.7 Data Mixing Law 26.5 54.5 13.0 62.0 24.4 49.1 45.2 32.0 12.0 35.4 !20 DoGraph (Ours) 27.3 56.9 14.8 63.0 26.3 50.8 53.5 33.9 14.5 37.9 Table 1: Downstream benchmark results (accuracy %) on SlimPajama (GPT-2 Medium). Tasks grouped into Commonsense/Reasoning, Reading Comprehension, and Language Modeling. Best results highlighted in bold. 3.1 Experiments Setup Our experiments use decoder-only, Transformer-based language models Vaswani et al. (2017); Radford et al. (2019) at 210M and 300M scales. Models are trained on SlimPajama Soboleva et al. (2023), spanning seven text domains. We evaluate DoGraph on nine stable benchmarks and compare with representative baselines. Model details, training protocol, and baseline breakdown are in Appendix A.2. 3.2 Results in the Pretraining Stage As shown in Table 1, DoGraph achieves more balanced learning across domains and delivers consistent gains over all baselines. It yields the largest improvements on reasoning-oriented benchmarks, highlighting the advantage of its structured weighting mechanism in capturing logical and commonsense dependencies across domains. Moreover, the performance gains on reading comprehension tasks, which require semantic consistency and information integration, demonstrate that DoGraph’s adaptive data scheduling enhances semantic alignment and improves overall generalization. 3.3 Perplexity Analysis Table 2 shows validation perplexity on SlimPajama under various domain-mixing strategies. Uniform sampling performs moderately but fails to balance domain frequencies. Loss-based weighting and prior methods (DoReMi, DOGE) yield unstable gains, overfitting to high-resource domains and degrading on long-tail data. RegMix and Data Mixing Law worsen this trend, with higher perplexity despite larger models. DoGraph achieves the best perplexity , reflecting balanced domain integration and strong generalization. Method SlimPajama (Val PPL ↓) Uniform 4.13 DYNAMIC LOSS-BASED 3.10 DoReMi 3.30 DOGE 3.31 RegMix 4.51 Data Mixing Law 4.50 !20 DoGraph (Ours) 3.09 Table 2: Pre-training results on SlimPajama. Validation perplexity (PPL) comparison across domain-mixing strategies. Lower values indicate better generalization. 3.4 DoGraph Stability across Model Scales As shown in Figure 2, validation perplexity decreases with model scale, but the rate of improvement depends on the reweighting strategy. Uniform weighting yields consistently high perplexity, while RegMix offers partial gains that diminish as models grow. DoGraph achieves the lowest perplexity across all scales, validating its ability to dynamically balance domains. Figure 2: Perplexity across GPT-2 model sizes. 4 Conclusion We revisited data mixing for LLMs through the lens of gradient dynamics. By characterizing domain differences via gradient geometry, we proposed DoGraph, a graph-constrained reweighting framework that adaptively balances domains during training. Experiments across model scales and benchmarks show that DoGraph improves both domain balance and generalization. Our results suggest that domains should be defined by the model’s evolving representation rather than human intuition. 5 Limitations While DoGraph achieves consistent improvements across domains and already reduces computational overhead through randomized gradient projection, its efficiency can still be further optimized. Future work will explore more lightweight aggregation and weighting strategies to enhance scalability in large-scale training. References [1] A. Albalak, Y. Elazar, S. M. Xie, S. Longpre, N. Lambert, X. Wang, N. Muennighoff, B. Hou, L. Pan, H. Jeong, et al. A survey on data selection for language models. Transactions on Machine Learning Research. Cited by: §1. A. Baevski et al. (2024) DataComp: in search of the next generation of multimodal datasets. arXiv preprint arXiv:2304.14108. Cited by: §A.1. Y. Bengio, A. Courville, and P. Vincent (2013) Representation Learning: A Review and New Perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence 35 (8), p. 1798–1828. Cited by: §1. Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. (2020) PiQA: Reasoning About Physical Commonsense in Natural Language. In Proceedings of the AAAI Conference on Artificial Intelligence, Cited by: §A.2. P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord (2018) Think You Have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge. arXiv preprint arXiv:1803.05457. External Links: 1803.05457 Cited by: §A.2. S. Fan, M. Pagliardini, and M. Jaggi (2023) DOGE: Domain reweighting with generalization estimation. In Second Agent Learning in Open-Endedness Workshop, External Links: Link Cited by: §A.2, §1. S. Fort, H. Hu, and B. Lakshminarayanan (2019) Deep Ensembles: A Loss Landscape Perspective. arXiv preprint arXiv:1912.02757. External Links: 1912.02757, Link Cited by: §1. J. Gao, D. He, X. Tan, T. Qin, L. Wang, and T. Liu (2019) Representation Degeneration Problem in Training Natural Language Generation Models. arXiv preprint arXiv:1907.12009. External Links: 1907.12009 Cited by: §1. L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy (2021) The pile: an 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027. External Links: 2101.00027, Link Cited by: §A.1, §A.2, §1. L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. L. Noac’h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou (2023) A framework for few-shot language model evaluation. External Links: Document, Link Cited by: §A.2. Y. Gu, L. Dong, H. Wang, Y. Hao, Q. Dong, F. Wei, and M. Huang (2024) Data Selection via Optimal Control for Language Models. arXiv preprint arXiv:2410.07064. External Links: 2410.07064 Cited by: §A.1. 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. (2022) Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. Cited by: §A.1. F. Kang, Y. Sun, B. Wen, S. Chen, D. Song, R. Mahmood, and R. Jia (2024) AutoScale: Automatic Prediction of Compute-Optimal Data Composition for Training LLMs. arXiv preprint arXiv:2407.20177. External Links: 2407.20177, Link Cited by: §1. J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei (2020) Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. Cited by: §A.1. P. W. Koh and P. Liang (2017) Understanding Black-box Predictions via Influence Functions. In Proceedings of the 34th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 70, p. 1885–1894. Cited by: §1. B. Z. Li et al. (2024) Dolma: an open corpus of three trillion tokens for language model pretraining research. arXiv preprint arXiv:2402.00159. Cited by: §A.1. J. Liu, L. Cui, H. Liu, D. Huang, Y. Wang, and Y. Zhang (2020) LogiQA: A Challenge Dataset for Machine Reading Comprehension with Logical Reasoning. arXiv preprint arXiv:2007.08124. External Links: 2007.08124 Cited by: §A.2. N. Liu et al. (2024) RegMix: regularizing data mixtures for language model pretraining. arXiv preprint arXiv:2407.10671. Cited by: §A.1, §A.2, §A.2, §1. S. Mehta, M. H. Sekhavat, Q. Cao, M. Horton, Y. Jin, C. Sun, I. Mirzadeh, M. Najibi, D. Belenko, P. Zatloukal, and M. Rastegari (2024) OpenELM: An Efficient Language Model Family with Open Training and Inference Framework. arXiv preprint arXiv:2404.14619. External Links: 2404.14619, Document Cited by: §A.2. T. Mihaylov, P. Clark, T. Khot, and A. Sabharwal (2018) Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering. arXiv preprint arXiv:1809.02789. External Links: 1809.02789 Cited by: §A.2. D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández (2016) The LAMBADA Dataset: Word Prediction Requiring a Broad Discourse Context. arXiv preprint arXiv:1606.06031. External Links: 1606.06031 Cited by: §A.2. A. Power, Y. Burda, H. Edwards, I. Babuschkin, and V. Misra (2022) Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets. arXiv preprint arXiv:2201.02177. External Links: 2201.02177 Cited by: §1. A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever (2019) Language models are unsupervised multitask learners. Note: OpenAI BlogVersion 1, Issue 8 Cited by: §3.1. S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang (2020) Distributionally robust neural networks for group shifts: on the importance of regularization for worst-case generalization. In International Conference on Learning Representations (ICLR), Cited by: §A.1. K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi (2021) WinoGrande: An Adversarial Winograd Schema Challenge at Scale. Communications of the ACM 64 (9), p. 99–106. Cited by: §A.2. P. Sarlin, D. DeTone, T. Malisiewicz, and A. Rabinovich (2020) SuperGlue: Learning Feature Matching with Graph Neural Networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 4938–4947. Cited by: §A.2. D. Soboleva, F. Al-Khateeb, R. Myers, J. R. Steeves, J. Hestness, and N. Dey (2023) SlimPajama: A 627B Token Cleaned and Deduplicated Version of RedPajama. Note: Dataset available at: https://huggingface.co/datasets/cerebras/SlimPajama-627B Cited by: §3.1. D. Sow, H. Woisetschläger, S. Bulusu, S. Wang, H. Jacobsen, and Y. Liang (2025) Dynamic loss-based sample reweighting for improved large language model pretraining. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: §A.2, §A.2, §1. H. Sun et al. (2025) Domain2Vec: vectorizing datasets to find the optimal data mixture without training. In Proceedings of the 42nd International Conference on Machine Learning (ICML), Cited by: §A.1, §1, §1. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017) Attention Is All You Need. In Advances in Neural Information Processing Systems, p. 5998–6008. Cited by: §3.1. J. Welbl, N. F. Liu, and M. Gardner (2017) Crowdsourcing Multiple Choice Science Questions. arXiv preprint arXiv:1707.06209. External Links: 1707.06209 Cited by: §A.2. A. Wettig, A. Gupta, S. Malik, and D. Chen (2024) QuRating: selecting high-quality data for training language models. In International Conference on Machine Learning, p. 52915–52971. Cited by: §1. S. M. Xie, H. Pham, X. Dong, N. Du, H. Liu, Y. Lu, P. Liang, Q. V. Le, T. Ma, and A. W. Yu (2023) DoReMi: optimizing data mixtures speeds up language model pretraining. In NeurIPS, Cited by: §A.1, §A.2, §1. Y. Xu, Q. An, J. Zhang, P. Li, and Z. Nie (2023) Hard sample aware prompt-tuning. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 12356–12369. Cited by: §1. J. Ye, P. Liu, T. Sun, Y. Zhou, J. Zhan, and X. Qiu (2024) Data mixing laws: optimizing data mixtures by predicting language modeling performance. arXiv preprint arXiv:2403.16952. Cited by: §A.1, §A.2, §1. D. Yoon (2023) SlimPajama-6B. Note: https://huggingface.co/datasets/DKYoon/SlimPajama-6BAccessed: 2024-09-24 Cited by: §A.2. R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi (2019) HellaSwag: Can a Machine Really Finish Your Sentence?. arXiv preprint arXiv:1905.07830. External Links: 1905.07830 Cited by: §A.2. Appendix A Appendix Contents A.1 Connections to Prior Work We categorize data mixture optimization into two main paradigms: offline and online approaches. Offline approaches predefine mixture ratios before training. Early scaling-law studies (Kaplan et al., 2020; Hoffmann et al., 2022) established the relationship between model size, data volume, and compute, motivating subsequent work that explicitly models how mixture composition affects performance. Methods such as DoReMi (Xie et al., 2023), RegMix (Liu and others, 2024), and Mixing Laws (Ye et al., 2024) optimize mixture ratios using proxy models or learned predictors, improving efficiency but requiring retraining when datasets change. Other efforts focus on heuristic sample scoring to derive refined data mixtures (Gu et al., 2024), distinct from large-scale corpora that offer fixed domain ratios for benchmarking (Gao et al., 2021; Baevski and others, 2024; Li and others, 2024). Domain2Vec (Sun and others, 2025) further introduces dataset vectorization and distribution alignment, enabling mixture optimization without proxy models. Online approaches adjust mixtures adaptively during training. Representative methods such as Group-DRO (Sagawa et al., 2020) dynamically reweight domains to improve worst-case generalization under distribution shift. While effective, they rely on explicit domain labels and are costly to scale. A.2 Experimental Details Benchmarks. We evaluate our method on nine diverse downstream benchmarks to assess its real-world impact. Guided by prior work Mehta et al. (2024) and our own observations, we selected these tasks for their performance stability, excluding volatile benchmarks like RTE. The chosen tasks are HellaSwag Zellers et al. (2019), PiQA Bisk et al. (2020), OpenBookQA Mihaylov et al. (2018), Lambada Paperno et al. (2016), SciQ Welbl et al. (2017), ARC-Easy Clark et al. (2018), COPA Sarlin et al. (2020), LogiQA Liu et al. (2020), and WinoGrande Sakaguchi et al. (2021). All evaluations use the lm-eval-harness Gao et al. (2023), and we report normalized accuracy where available, otherwise standard accuracy. Baselines. To rigorously assess the effectiveness of our proposed method, DoGraph, we benchmark it against a diverse set of reweighting baselines spanning three levels of granularity. We first include the uniform mixing baseline, where all samples contribute equally, as a fundamental reference. We then compare DoGraph with state-of-the-art domain-level reweighting methods, including DoGE Fan et al. (2023), DoReMi Xie et al. (2023), Regmix Liu and others (2024), and Data Mixing Law Ye et al. (2024). Finally, to evaluate performance at a finer granularity, we incorporate a representative sample-level reweighting approach, Dynamic Loss-based Sample Reweighting Sow et al. (2025). Commonsense / Reasoning RC LM Avg Method HellaSwag PiQA OBQA COPA LogiQA WinoG SciQ ARC-E Lambada The Pile Uniform 29.5 58.8 27.3 65.8 23.9 50.5 60.3 40.0 11.7 40.9 Dynamic Loss-Based 29.0 57.7 26.4 64.3 22.8 49.3 60.0 38.9 10.2 39.9 DoReMi 29.4 58.3 27.3 67.5 26.4 52.2 61.6 40.6 12.1 41.7 DOGE 29.2 58.5 27.1 64.5 23.2 49.8 60.1 40.0 11.7 40.5 RegMix 29.2 59.3 27.3 65.2 25.8 53.1 62.8 41.7 14.2 42.1 Data Mixing Law 29.2 58.8 26.9 67.2 23.6 50.4 58.6 39.0 11.9 40.6 !20 DoGraph (Ours) 29.8 59.2 27.8 65.0 28.3 51.2 66.1 39.2 15.9 42.5 Table 3: Downstream benchmark results (accuracy %) on The Pile (LLaMA-1.1B). Tasks are grouped into Commonsense/Reasoning, Reading Comprehension, and Language Modeling. Our method, DoGraph, achieves consistently better and more balanced results across domains, demonstrating its competitiveness and generalization ability. The best results are highlighted in bold. Commonsense / Reasoning RC LM Avg Method HellaSwag PiQA OBQA COPA LogiQA WinoG SciQ ARC-E Lambada The Pile Uniform 29.6 58.8 29.4 66.0 25.9 51.1 61.0 39.1 12.6 41.5 Dynamic Loss-Based 29.3 58.1 29.6 66.1 25.0 52.5 62.7 39.9 12.2 41.7 DoReMi 29.6 58.4 29.8 66.0 24.9 51.4 61.1 40.5 12.8 41.6 DOGE 29.7 56.9 29.2 64.0 25.5 50.6 61.7 40.6 11.9 41.1 RegMix 29.4 59.5 29.4 66.5 25.1 53.6 62.5 41.2 12.3 42.2 Data Mixing Law 29.3 58.4 30.2 65.9 25.6 51.3 61.3 40.2 12.1 41.6 !20 DoGraph (Ours) 29.6 60.5 29.0 67.0 29.7 51.4 65.2 40.2 15.1 43.1 Table 4: Downstream benchmark results (accuracy %) on The Pile (LLaMA-3.2-3B). Tasks are grouped into Commonsense/Reasoning, Reading Comprehension, and Language Modeling. Our method, DoGraph, achieves consistently better and more balanced results across domains, demonstrating its competitiveness and generalization ability. The best results are highlighted in bold. Commonsense / Reasoning RC LM Avg Method HellaSwag PiQA OBQA COPA LogiQA WinoG SciQ ARC-E Lambada SlimPajama Uniform 26.0 55.4 13.8 57.2 22.8 49.3 32.6 30.6 12.0 33.3 Dynamic Loss-Based 26.2 56.1 13.2 55.3 26.0 49.2 53.8 31.8 12.6 36.2 DoReMi 26.1 55.7 12.3 53.5 26.8 48.8 52.4 30.9 12.4 35.4 DOGE 26.2 55.0 14.4 60.5 23.5 49.0 31.1 30.8 11.4 33.5 RegMix 26.0 54.3 13.3 58.0 24.1 49.8 38.7 29.8 12.5 34.1 Data Mixing Law 26.1 56.3 13.4 59.2 24.5 48.9 39.6 30.1 12.4 34.5 !20 DoGraph (Ours) 26.3 57.5 14.6 58.0 26.0 49.7 53.8 32.3 12.8 36.4 Table 5: Downstream benchmark results (accuracy %) on SlimPajama (GPT-2 Small). Tasks are grouped into Commonsense/Reasoning, Reading Comprehension, and Language Modeling. Our method, DoGraph, achieves consistently better and more balanced results across domains, demonstrating its competitiveness and generalization ability. The best results are highlighted in bold. Training Datasets. Our training data strategy is designed to align dataset scale with model capacity. For all GPT-2 models, we utilize the SlimPajama-6B dataset Yoon (2023), a 6-billion-token corpus comprising seven diverse domains: ArXiv, Books, Common Crawl, C4, GitHub, StackExchange, and Wikipedia. The byte proportion of each source is detailed in Table 6, illustrating the composition of the data mixture used for training. For all LLaMA models, we conduct our experiments using the domains of the Pile dataset Gao et al. (2021) depicted in Table 7. Due to copyright concerns, we utilize the 17 subsets available on HuggingFace that do not violate copyright issues. These datasets provide a balanced and diverse text distribution suitable for evaluating cross-domain generalization in medium-scale language models. Data Source Byte Proportion Common Crawl 54.1% C4 28.7% GitHub 4.2% Books 3.7% ArXiv 3.4% Wikipedia 3.1% StackExchange 2.8% Table 6: Byte proportion of data sources in the SlimPajama-6B dataset. Component Effective Size Pile-C 227.12 GiB PubMed Central 180.55 GiB Books3 151.44 GiB OpenWebText2 125.54 GiB ArXiv 112.42 GiB Github 95.16 GiB FreeLaw 76.73 GiB Stack Exchange 64.39 GiB USPTO Backgrounds 45.81 GiB PubMed Abstracts 38.53 GiB Gutenberg (PG-19) 27.19 GiB OpenSubtitles 19.47 GiB Wikipedia (en) 19.13 GiB DM Mathematics 15.49 GiB Ubuntu IRC 11.03 GiB BookCorpus2 9.45 GiB EuroParl 9.17 GiB HackerNews 7.80 GiB YoutubeSubtitles 7.47 GiB PhilPapers 4.76 GiB NIH ExPorter 3.79 GiB Enron Emails 1.76 GiB Table 7: Overview of the Pile dataset with datasets that are no longer available due to copyright issues marked in gray. Merged into a single column list. Model Architecture. Following prior studies Liu and others (2024); Sow et al. (2025), we consider both model architecture and model scale in our evaluation, as summarized in Table 8. Specifically, we evaluate two decoder-only Transformer models based on GPT-2 architecture and two models based on LLaMA architecture, ranging from lightweight to medium scales. GPT-2 Small GPT-2 Medium LLaMA-1.1B LLaMA-3.2-3B Parameters 210M 300M 1.1B 3B Layers 24 36 22 28 Attention Heads 16 24 32 24 Embedding Dim. 768 768 2048 8192 Hidden Dim. 3072 3072 2048 3072 Max Seq. Length 512 512 2048 131072 Table 8: Model architectures used in our experiments. Training Process. Following standardized practices in prior work, we train all models under protocols summarized in Table 9. Specifically, we adopt a linear warmup cosine schedule with identical weight decay (0.01) and gradient clipping (1.0) across all model scales, while adjusting batch size and training steps according to model capacity. This setup ensures that each model is trained sufficiently to convergence. GPT-2 Small GPT-2 Medium TinyLLaMA-1.1B LLaMA-3.2-3B Minibatch Size 48 48 64 64 Learning Rate (×10−3)(× 10^-3) 0.50 0.50 0.50 0.50 Learning Rate End (×10−4)(× 10^-4) 1.0 1.0 1.0 1.0 Warmup Steps 500 500 500 500 r 0.4 0.4 0.4 0.4 Training Steps 20,000 20,000 25,000 25,000 Total Documents Seen 960,000 960,000 1280,000 1280,000 Table 9: Training hyperparameters for GPT-2 and LLaMA models in our benchmark evaluations. A.3 DoGraph Pipeline Formalized in Algorithm 1, the process begins by projecting high-dimensional per-sample gradients gi(t)∈ℝdg_i^(t) ^d into a lower-dimensional subspace g~i(t)∈ℝk g_i^(t) ^k using a random Gaussian matrix R(t)R^(t), where we set k=5000k=5000 for both SlimPajama and The Pile to preserve the gradient manifold’s geometric properties per the Johnson-Lindenstrauss Lemma. Subsequently, we identify latent optimization structures by applying K-means clustering to these projected signals, partitioning the mini-batch into m=11m=11 model-centric domains Dj(t)\D_j^(t)\ and computing their respective centroid gradients g¯j(t) g_j^(t). Finally, importance weights w(t)∈Δm−1w^(t)∈ ^m-1 are determined by solving the auxiliary objective ℒoptL_opt, and the model parameters θ are updated via the weighted aggregate ∑j=1mwj(t)g¯j(t) _j=1^mw_j^(t) g_j^(t), effectively decoupling training dynamics from static, pre-defined domain labels. Input: Training data D, parameters θ, number of clusters m, projection dimension k, epochs T, learning rate η Output: Trained parameters θ∗θ^*, domain weights w(t)\w^(t)\ for t=1t=1 to T do Sample random projection matrix R(t)∈ℝd×kR^(t) ^d× k with Rpq(t)∼(0,1/k)R_pq^(t) (0,1/k); Compute per-sample gradients gi(t)=∇θL(xi,yi;θ(t−1))g_i^(t)= _θL(x_i,y_i;θ^(t-1)); Project gradients: g~i(t)=R(t)⊤gi(t) g_i^(t)=R^(t) g_i^(t); Cluster g~i(t)\ g_i^(t)\ into m domains D1(t),…,Dm(t)\D_1^(t),…,D_m^(t)\; Compute domain mean gradients g¯j(t)=1|Dj(t)|∑i∈Dj(t)g~i(t) g_j^(t)= 1|D_j^(t)| _i∈ D_j^(t) g_i^(t); Optimize weights w(t)=argminw∈Δm−1ℒopt(∑j=1mwjg¯j(t))w^(t)= _w∈ ^m-1L_opt\! ( _j=1^mw_j g_j^(t) ); Update model parameters: θ(t)←θ(t−1)−η∑j=1mwj(t)g¯j(t)θ^(t)←θ^(t-1)-η _j=1^mw_j^(t) g_j^(t); return θ(T),w(t)t=1Tθ^(T),\w^(t)\_t=1^T Algorithm 1 Dograph Pipeline A.4 Scaling to Larger Datasets and Model Sizes We report results on GPT-2 models from 210M to 300M parameters and a 6B-token SlimPajama subset, as shown in Table 5. DoGraph is scale-free and does not rely on any model-size–specific assumptions. All components, including gradient extraction, random projection, clustering, and domain-level optimization, operate directly on per-step gradients and thus scale linearly with model size. The method does not require proxy models, validation-model fitting, or domain-specific metadata, making it naturally compatible with billion-parameter LLMs. To further prove these, we pretrain LLaMA-1.1B and LLaMA-3B from scratch under the same DoGraph pipeline, as shown in Table 3 and Table 4. A.5 Clustering Visualization As shown in Figure 3, while human-defined domains (indicated by colors) become indistinguishable later in training, DoGraph successfully extracts m latent structures from this mixture, proving that model-centric domains are composed of heterogeneous data sources. Figure 3: Evolution of per-sample gradients and the emergence of model-centric structures. Points are colored by their original source datasets. Initially, gradients are separated by domain bias. Over time, the model homogenizes its perception of these sources, leading to significant overlap. Despite this mixing, DoGraph identifies m=11 distinct model-centric domains within the gradient space. Experiments use 20% of SlimPajama trained on GPT2-Mini. A.6 More Analysis about the Choice of Optimization Function At each training epoch, the DoGraph framework computes domain mean gradients g¯jj=1m\ g_j\_j=1^m and determines their adaptive weights w∈Δm−1w∈ ^m-1 by minimizing an auxiliary objective ℒoptL_opt. Since m (the number of domains) is typically small, this optimization occurs in a low-dimensional space and can be efficiently solved in closed or iterative form. We discuss several representative objectives and their corresponding solvers below. Gradient variance minimization. To balance the learning progress across domains, one may minimize the variance of gradient magnitudes while maintaining the global descent direction: ℒopt(w)=Varj[‖g¯j‖2]+λ‖∑j=1mwjg¯j‖22.L_opt(w)=Var_j\! [\| g_j\|_2 ]+λ \| _j=1^mw_j g_j \|_2^2. This convex quadratic problem can be solved by projected gradient descent or quadratic programming with a simplex constraint. Robust min–max objective. When robustness against hard or under-represented domains is desired, one may adopt a distributionally robust formulation: ℒopt(w)=τlog∑j=1mexp(‖g¯j‖2/τ),L_opt(w)=τ \! _j=1^m \! (\| g_j\|_2/τ ), which smoothly approximates maxj‖g¯j‖2 _j\| g_j\|_2. The optimal weights admit a closed-form softmax solution wj∝exp(‖g¯j‖2/τ)w_j (\| g_j\|_2/τ). Gradient alignment regularization. To encourage consistent update directions across domains, we define ℒopt(w)=−∑j=1mwjcos(g¯j,g¯),g¯=∑j=1mwjg¯j.L_opt(w)=- _j=1^mw_j ( g_j,\; g), g= _j=1^mw_j g_j. Although non-convex due to the dependence of g¯ g on w, it can be efficiently solved by a few fixed-point iterations: each step updates wjw_j in proportion to the cosine similarity between g¯j g_j and the current aggregate g¯ g. Domain uncertainty weighting. Alternatively, if each domain exhibits distinct gradient variability, we estimate its intra-domain variance σj2=Vari∈Dj[‖gi−g¯j‖22] _j^2=Var_i∈ D_j\! [\|g_i- g_j\|_2^2 ] and assign weights inversely proportional to it: ℒopt(w)=‖∑jwjg¯j‖22+β∑jσj2wj.L_opt(w)= \| _jw_j g_j \|_2^2+β _j _j^2w_j. This convex quadratic form admits a closed-form Newton update or can be solved by a projected Frank–Wolfe method. Table 10 summarizes the computational complexity of each optimization objective and the corresponding validation perplexity PPL. All variants share the same backbone and differ only in the choice of ℒoptL_opt. The robust softmax objective achieves the lowest computational cost, while the uncertainty-weighted variant attains the best overall performance. Table 10: Comparison of optimization objectives in DoGraph. Method Complexity PPL DoGraph (variance) O(m2)O(m^2) 3.24 DoGraph (robust softmax) O(m)O(m) 3.31 DoGraph (alignment) O(m2)O(m^2) 3.15 DoGraph (uncertainty) O(m2)O(m^2) 3.09 Among all variants, DoGraph (uncertainty) achieves the lowest perplexity, indicating that weighting domains by intra-domain gradient stability provides the most consistent optimization dynamics. A.7 Impact of Cluster Granularity m. We investigate the sensitivity of model performance to the number of clusters m. As illustrated in Figure 1, the validation perplexity exhibits a clear U-shaped trend with respect to m. Performance initially improves as m increases from 7 to 11, suggesting that moderately finer-grained, model-centric domains better capture coherent gradient structures and facilitate optimization. However, further increasing m beyond 11 leads to a significant performance degradation. This decline is likely due to over-partitioning, which fragments the gradient space and splits coherent patterns into inconsistent components, thereby weakening signal consistency. Consequently, we select m=11 as our default setting for all subsequent experiments. Figure 4: Impact of cluster granularity m on validation perplexity. The U-shaped curve demonstrates that m=11m=11 provides the optimal balance; insufficient granularity fails to resolve gradient structures, while excessive partitioning leads to signal inconsistency. A.8 Computational Efficiency Analysis As shown in Figure 5, dograph achieves state-of-the-art performance while introducing a modest and practical computational overhead. On a 2×2× H200 GPU cluster, our method completes pre-training in 20.37 hours, corresponding to a 4.51% increase in runtime compared to regmix. This incremental cost falls within the commonly accepted budget for large-scale pre-training. Figure 5: Pre-training GPT-2 Mini on SlimPajama under a 100B-Token Computational Budget. We report the total training time (GPU hours) using 2×2× NVIDIA H200 GPUs. While our dograph method introduces a sophisticated data-driven decision process, the resulting overhead is minimal (only 4.51% over regmix), while establishing a new SOTA performance baseline. The marginal increase in budget is well-justified by the superior convergence quality and data selection efficiency. A.9 Proofs Assumption A.1 (Linearized Attention Mechanism). Let X∈ℝn×dX ^n× d denote the sequence representation. We define the projected queries, keys, and values as Q=XWQQ=XW_Q, K=XWKK=XW_K, and V=XWVV=XW_V, and the scaled similarity matrix as S=1dkQK⊤S= 1 d_kQK . The row-wise softmax of S is approximated by its first-order linearization: P≈A+1τnCS,whereA=1n⊤,C=I−1n⊤.P~≈~A+ 1τ nCS, A= 1n11 ,\ C=I- 1n11 . Consequently, the attention output satisfies O=PV≈AV+TQK⊤VO=PV≈ AV+TQK V, where T=1τndkCT= 1τ n d_k\,C. Assumption A.2 (Linear Output Transformations). The attention output O is passed through two linear mappings: H=OWOH=OW_O and Z=HWZ=HW. The model prediction is obtained via Π=softmax(Z) =softmax(Z), which represents the token-level probability distribution. Assumption A.3 (Upstream Gradients and Mismatch Tensor). Given the ground-truth label matrix Y, the upstream gradients are defined as GZ=Π−YG_Z= -Y, GH=GZW⊤G_H=G_ZW , and GO=GHWO⊤G_O=G_HW_O . We further define the mismatch tensor R=(Π−Y)MR=( -Y)M, where M=W⊤WO⊤M=W W_O . Assumption A.4 (Regularity Conditions). All expectations involved in subsequent derivations are assumed to exist, and the per-sample gradients are square-integrable. Per-sample gradients and proof of Theorem 3.2 . With dL=⟨GO,dO⟩dL= G_O,dO and O≈AV+TQK⊤VO≈ AV+TQK V, the per-sample gradients are ∂L∂V=A⊤GO+KQ⊤T⊤GO,∂L∂Q=(T⊤GO)V⊤K,∂L∂K=VGO⊤TQ,∂L∂W=H⊤GZ \ aligned ∂ L∂ V&=A G_O+KQ T G_O,\\[3.0pt] ∂ L∂ Q&=(T G_O)\,V K,\\[3.0pt] ∂ L∂ K&=V\,G_O T\,Q,\\[3.0pt] ∂ L∂ W&=H G_Z aligned . ∂L∂WV=X⊤(A⊤GO+KQ⊤T⊤GO),∂L∂WQ=X⊤((T⊤GO)V⊤K),∂L∂WK=X⊤(VGO⊤TQ),∂L∂WO=O⊤GH \ aligned ∂ L∂ W_V&=X \! (A G_O+KQ T G_O ),\\[3.0pt] ∂ L∂ W_Q&=X \! ((T G_O)\,V K ),\\[3.0pt] ∂ L∂ W_K&=X \! (V\,G_O T\,Q ),\\[3.0pt] ∂ L∂ W_O&=O G_H aligned . From Assumption A.3, the upstream gradient can be written as GO=(Π−Y)W⊤WO⊤=RG_O=( -Y)W W_O =R. Substituting this into the above expressions shows that all per-sample gradients are linear functions of R: ∇WbL(x,y;θ)=b(X,Q,K,V,T)[R(x,y)], _W_bL(x,y;θ)~=~ Lin_b(X,Q,K,V,T)\,[\,R(x,y)\,], where b(⋅) Lin_b(·) denotes a matrix-valued linear operator determined only by the forward pass variables X,Q,K,V,TX,Q,K,V,T. Using the identity vec(UGV)=(V⊤⊗U)vec(G)vec(UGV)=(V \! U)vec(G), each matrix gradient can be rewritten in vectorized form as gb(s) g_b(s) :=vec(∇WbL(x,y;θ))=ℒb(x)ρ(s), =vec\! ( _W_bL(x,y;θ) )~=~L_b(x)\,ρ(s), ρ(s) ρ(s) :=vec(R(s)). =vec\! (R(s) ). Here, ℒb(x)L_b(x) absorbs all Kronecker factors (e.g., X⊤X , T⊤T , K, Q) from the explicit gradient expressions. Hence, for every parameter block b, the sample-wise gradient is a linear transformation of the mismatch vector ρ(s)ρ(s). Define the expected gradient under a data distribution P as g¯b(P):=s∼P[gb(s)] g_b(P):=E_s P[g_b(s)]. By linearity of expectation (Bochner integral in finite dimensions), g¯b(P1)−g¯b(P2)=∫gb(s)(P1−P2)(ds). g_b(P_1)- g_b(P_2)~=~ g_b(s)\,(P_1-P_2)(ds). The inner product between two per-sample gradients naturally defines kb(s,s′) k_b(s,s ) :=⟨gb(s),gb(s′)⟩ = g_b(s),\,g_b(s ) =ρ(s)⊤ℒb(x)⊤ℒb(x′)ρ(s′). =ρ(s) \,L_b(x) L_b(x )\,ρ(s ). Because kbk_b is an inner product in feature space, it is positive semidefinite. Applying Fubini–Tonelli and the bilinearity of the inner product yields ‖g¯b(P1)−g¯b(P2)‖22 \| g_b(P_1)- g_b(P_2) \|_2^2 =⟨∫gbd(P1−P2),∫gbd(P1−P2)⟩ = g_b\,d(P_1\!-\!P_2),\ g_b\,d(P_1\!-\!P_2) =∬⟨gb(s),gb(s′)⟩(P1−P2)(ds)(P1−P2)(ds′) = g_b(s),g_b(s ) \,(P_1\!-\!P_2)(ds)\,(P_1\!-\!P_2)(ds ) =P1,P1[kb]+P2,P2[kb]−2P1,P2[kb], =E_P_1,P_1[k_b]+E_P_2,P_2[k_b]-2\,E_P_1,P_2[k_b], which is exactly MMDkb2(P1,P2)MMD^2_k_b(P_1,P_2) by definition. Finally, the positive semidefiniteness of kbk_b follows from ∑i,jαiαjkb(si,sj)=‖∑iαigb(si)‖22≥0. _i,j _i _jk_b(s_i,s_j)= \| _i _ig_b(s_i) \|_2^2≥ 0. This completes the proof. ∎