Pith. sign in

REVIEW 3 major objections 4 minor 12 references

DIDS: Domain Impact-aware Data Sampling for Large Language Model Training

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

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

Pith's one-line read Domain-impact-aware data sampling outperforms full-data LLM training using only 100k examples, according to experiments on Llama-3.1-8B and nine downstream tasks.

desk verdict Empirically broad, but the central update rule as written appears to reward domain gradients that oppose the task gradient—the opposite of the paper's narrative; fix that and this could be worth publishing. read the letter →

arxiv 2504.13227 v2 pith:TUPS347F submitted 2025-04-17 cs.CL cs.AIcs.LG

Weijie Shi , Jipeng Zhang , Yaguang Wu , Jingzhi Fang , Ruiyuan Zhang , Jiajie Xu , Jia Zhu , Hao Chen
show 3 more authors
Yao Zhao Sirui Han Xiaofang Zhou
This is my paper · ORCID
classification cs.CLcs.AIcs.LG
keywords domain-leveldatasamplingmixinggradientclusteringFisherinformationmatrixlargelanguagemodeltrainingdownstreamtaskimpactlosstrajectoryweightingefficiency
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

This paper argues that the biggest untapped lever in LLM post-training is not adding data but deciding which data to sample, and proposes Domain Impact-aware Data Sampling (DIDS) to make that choice online. DIDS first repartitions the training corpus by clustering proxy-model gradients, so each "domain" is a set of examples with similar learning effects rather than similar topics. It then measures each domain's impact on each downstream task with a Fisher Information Matrix (FIM) guided quadratic form $I(D_i,S_j) = \tfrac12(\nabla\ell_{S_j}-\nabla\ell_{D_i})^\top F (\nabla\ell_{S_j}-\nabla\ell_{D_i})$, and updates sampling probabilities by combining that impact with loss-improvement trajectories. On Llama-3.1-8B and nine tasks, DIDS reports average scores of 62.3 (multi-task) and 63.7 (single-task) using only 100k samples, versus 61.2 for training on the full 929k corpus, with roughly 3.4% higher average performance claimed. The point of interest is practical: if it holds, downstream accuracy can be improved while spending about 10% of the data and a small compute overhead.

What carries the argument

The load-bearing object is the FIM-guided domain impact metric, a quadratic form $I(D_i,S_j) = \tfrac12(\nabla\ell_{S_j}-\nabla\ell_{D_i})^\top F (\nabla\ell_{S_j}-\nabla\ell_{D_i})$ over the difference between a training-domain gradient and a downstream-task gradient, with $F$ the diagonal Fisher Information Matrix $F \approx \mathbb{E}[\nabla\log p(y|\theta)\odot \nabla\log p(y|\theta)]$. It is the second-order Taylor approximation of the KL divergence between the model's predictive distribution after a domain update and after a task update. This metric feeds the sampling update in Eq. (16), where impact is multiplied by per-task loss improvements and potential, divided by previous sampling weight, then smoothed by EMA; the accompanying gradient-clustering step supplies the "domains" that make the metric meaningful.

What would settle it

Inspect the released code: for one domain-task pair where training on the domain clearly improves task loss, compute $\Delta=\nabla\ell_{S_j}-\nabla\ell_{D_i}$ and the resulting sampling-weight change; if the weight rises when $\Delta^\top F \Delta$ is large, the mechanism is inverted and the gains would need another explanation. Alternatively, run DIDS with the impact term's sign flipped and compare average scores.

Watch

Extended reading notes

Core claim

The paper's central claim is that a dynamically updated domain sampling distribution, computed from gradient-based clusters and an FIM-guided measure of how each cluster shifts the model's output distribution on target tasks, optimizes the bi-level objective in Eq. (4) better than uniform, random, or gradient-alignment baselines, and even better than full-data training at one-tenth the budget. Concretely, DIDS reports 62.3 multi-task and 63.7 single-task averages on Llama-3.1-8B across BBH, BoolQ, GSM8K, MathQA, IFEval, MMLU, PIQA, PubMedQA, and TruthfulQA, against 61.2 for full-data training and 60.2 for the best multi-task baseline. The result extends to Mixtral-7B and the OpenHermes-2.5 dataset in the appendix, where DIDS again matches or exceeds full-data training with 100k samples.

Load-bearing premise

The reported gains depend on the sampling-update formula meaning what the prose says; as written, the formula appears to upweight domains whose updates oppose the downstream task, so the implementation, not the metric, is the load-bearing premise.

Editorial extensions

If this is right

  • If the central claim holds, a 100k-sample impact-weighted subset produces higher average downstream accuracy than the full 929k-corpus training run, so post-training budgets can drop by roughly an order of magnitude.
  • The appendix replicates the pattern on Mixtral-7B and on the OpenHermes-2.5 dataset, which suggests the sampling rule transfers across model families and data distributions rather than memorizing one benchmark suite.
  • The gradient-based domain repartition places algebra and code together by learning effect, and replacing that repartition with BERT semantic clustering costs over a point of average accuracy; this makes behavioral clustering a load-bearing part of the method.
  • Dynamic updates that combine impact with loss trajectories beat every static mixing ratio in the grid-search comparison, so the optimal mixture is not constant over training.
  • DIDS retains competitive scores on four unseen out-of-distribution tasks (WMT16, TLDR, ARC-Challenge, MBPP), indicating that the optimized sampling distribution does not merely overfit the nine observed evaluation tasks.

Reading between the lines

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

  • Beyond the paper: the sign of the impact term deserves a direct ablation, because $I(D_i,S_j)$ vanishes when the domain and task gradients agree and grows when they oppose; as written, Equation (16) would reward opposing updates, so the reported gains may be carried by the clustering, loss-trajectory, or smoothing components rather than by the FIM term.
  • Beyond the paper: the diagonal-FIM approximation and the exponential-decay loss model are both local assumptions that the appendices acknowledge; the method's transfer to out-of-distribution tasks is consistent with the sampling acting as an implicit regularizer, which could be tested by comparing validation-loss trajectories.
  • Beyond the paper: because the impact metric is a local quadratic form, it should become less reliable exactly when updates are large; a practical extension would be to recompute the FIM on the updated model more often or to blend in direct task-accuracy estimates.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

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 proposes DIDS, a domain-level data sampling method for LLM post-training. It has three components: (1) gradient-based domain repartitioning using a small proxy model, top-k gradient retention, and Johnson-Lindenstrauss random projection followed by k-means clustering; (2) a Fisher Information Matrix (FIM) guided impact metric, obtained as a second-order Taylor approximation of the KL divergence between model predictions after a domain update and after a downstream-task update; and (3) a periodic sampling-probability update that combines the FIM impact scores with current loss improvement and a predicted future loss improvement, smoothed by an exponential moving average. Experiments on Llama-3.1-8B and Mixtral-7B with Tulu-3 and OpenHermes-2.5 report that 100k samples selected by DIDS outperform full-data training and several baselines on a nine-task suite, with additional ablations, OOD evaluations, and computational cost analyses in the appendices.

Significance. If the reported results hold, DIDS would be a practical data-efficiency contribution: it reduces post-training data by roughly an order of magnitude while matching or exceeding full-data performance, and the public code would support reproducibility. The paper also contains a genuine theoretical fragment, the KL/FIM Taylor expansion in Section 4.2, and a commendably extensive experimental appendix covering multiple models, datasets, ablations, OOD generalization, and cost breakdowns. However, the central mechanism as written has an internal inconsistency that must be resolved before the empirical claims can be attributed to the proposed algorithm.

major comments (3)
  1. [§4.2–4.3, Eqs. (12) and (16)] The core mechanism is internally inconsistent as written. Eq. (12) defines I(D_i,S_j) = (1/2)(∇ℓ_Sj − ∇ℓ_Di)^T F (∇ℓ_Sj − ∇ℓ_Di). Since F is positive semidefinite, I is zero when the domain gradient equals the task gradient and grows as the two gradients diverge or oppose. Eq. (16) then multiplies this quantity by (l_c + l_p), both of which are nonnegative when the task loss is improving, and adds the product to the sampling weight. The update therefore upweights domains whose gradient updates are most different from the task update, which is the opposite of the alignment story in the prose. No derivation is given showing that this update optimizes the Eq. (4) bi-level objective, and l_c = ΔL(S_j) is a task-level scalar that is identical for every domain, so it cannot attribute the observed task improvement to a specific domain. If the implementation follows the text, the reported gains cannot be explained by the stated mechanism; the authors need to either correct the sign/direction of the metric, replace it with an alignment-based measure, or supply a derivation connecting Eq. (16) to Eq. (4).
  2. [§5.1.2 and Algorithm 1] There is a potential evaluation circularity that needs to be resolved. The downstream tasks S_j are used in Algorithm 1 to compute per-domain gradients, the FIM, and the loss trajectories that drive the sampling update, and the same nine tasks are then reported as the evaluation scores in Table 1. Section 5.1.2 mentions partitioning a validation subset for the baselines, but it does not clearly state whether DIDS itself also restricts its observable task data to a separate held-out subset. If the reweighting uses any of the evaluation data or the same task distributions as the final evaluation, the reported gains partly reflect fitting to the evaluation suite. The OOD results in Appendix J are helpful, but the main claim in Section 5.2 still needs a clean separation between what is observed for sampling and what is reported as performance.
  3. [Abstract and §4.3] The paper claims 'theoretical guarantees' and 'theoretical foundations' for the FIM-guided impact metric, but the only derivation is the Taylor approximation of KL divergence in Section 4.2 and Appendix E. There is no theorem showing that the sampling update in Eq. (16) optimizes the stated objective in Eq. (4), nor any convergence or optimality guarantee for the resulting mixture. The word 'guarantees' should be reserved for what is actually proven, or the claims should be toned down to describe the metric as a heuristic with a Taylor-based justification.
minor comments (4)
  1. [Table 4] The sample counts and percentages in Table 4 do not sum to the claimed totals: the listed components sum to 889,344 samples rather than 939,344, and the percentages sum to roughly 94.5% rather than 100%. The table should be reconciled with the dataset manifest.
  2. [§4.2, Eqs. (10)–(12) and Appendix E] The notation for the Fisher matrix is inconsistent: Eq. (11) uses the elementwise product E[∇log p ⊙ ∇log p], the figure uses a similar form, and Eq. (17) in Appendix E uses the outer product ∇log p (∇log p)^T. The diagonal approximation should be stated explicitly and used consistently throughout.
  3. [Appendix E.2] Appendix E.2 says the KL approximation is 'positive definite and symmetric (when properly scaled).' With a diagonal F the quadratic form 1/2 Δ^T F Δ is positive semidefinite, not necessarily positive definite, since F can have zero diagonal entries.
  4. [§5.5.1 and Figure 2a] The sentence describing the 5-update setting says DIDS achieves 58.2, 'comparable to the random sampling baseline at 58.9,' which is a roundabout way of saying it underperforms random sampling by 0.7 points; the text should state this directly so readers can assess the sensitivity of the method.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the sampling weights are task-aware by design, and the OOD evaluations provide independent support.

full rationale

DIDS's sampling weights are updated using FIM-guided gradient differences and loss improvements measured on validation subsets of the downstream tasks (Sec. 5.1.2), and the main table reports accuracy on those same tasks. This is task-aware validation-set tuning rather than a derivation that makes the reported accuracy equal to the fitted inputs by construction: test accuracy is not identical to validation loss l_c or l_p, nor to the gradient-distance I. The loop is further broken by Appendix J, which evaluates on unseen OOD tasks (WMT16, TLDR, ARC-Challenge, MBPP) and finds DIDS highest average (53.5 vs 52.6 for DGA and 51.0 for Random), an externally falsifiable result not used in fitting the sampling weights. The only load-bearing internal issue is a sign/direction concern in Eq. (12)/(16): since I is a positive quadratic form in (∇ℓ_Sj - ∇ℓ_Di), it is zero when domain and task gradients are aligned and largest when they oppose, which is a mechanistic correctness risk rather than a circularity. No self-citation chain, imported uniqueness theorem, or ansatz-by-citation is used; stated limitations (diagonal FIM approximation, exponential-decay loss-trajectory assumption) do not create circularity.

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

The central claim rests on several tuning choices (domain count, EMA coefficient, update interval, projection dimension, top-k fraction, proxy model size) and on modeling assumptions about small updates, diagonal FIM, proxy representativeness, exponential loss decay, and the disjointness of the observable validation subset. The exponential decay parameters are fitted to the evaluation-task loss curves, which is the clearest free-parameter dependence. No new physical or architectural entities are introduced.

free parameters (7)
  • Domain count k = 72
    Chosen via parameter analysis in Section 5.5.4; the text states performance plateaus near 72 domains and adds little at 1152, but there is no principled selection criterion.
  • Top-k gradient fraction and retained layer fraction = top 10% of elements; last 10% of layers
    Computational shortcuts introduced in Section 4.1 with no evidence that the retained gradients preserve clustering quality.
  • Random projection target dimension s = 1024
    Set in Appendix B.4 as a clustering-manageable JL dimension; the paper gives no sensitivity analysis for this choice.
  • EMA momentum coefficient beta = 0.1
    Set in Appendix B.4 with practical guidance 0.1-0.3 in Appendix F, without derivation or sensitivity curves.
  • Domain update interval tau = 4000 steps
    Set in Appendix B.4; Fig 2a shows 5-95 updates during training, which is inconsistent with the stated 100k-sample, batch-size-128 budget.
  • Exponential decay fit parameters (a, b, c) = estimated via curve fitting on task loss history
    Eq. (14) fits an exponential decay to each downstream task's loss curve; these parameters are calibrated on the same tasks used for evaluation, so the derived potential l_p is not an independent predictor.
  • Proxy model size = 500M
    Selected from Fig 3a after showing minimal impact on final score; a design choice rather than a fitted constant.
assumptions (7)
  • domain assumption Gradient updates are small enough for the second-order Taylor expansion of KL divergence to be accurate.
    Invoked in Section 4.2 before Eq. (8) and revisited in the Limitations section; the paper acknowledges the approximation may miss parameter interactions.
  • domain assumption The diagonal FIM approximation at theta is close to the FIM at theta_Di.
    Stated in Section 4.2 after Eq. (10) and listed as a limitation; the paper argues the difference is negligible for small updates, which is an assumption rather than a proven bound.
  • domain assumption Gradients from a 500M proxy model are representative of gradients of the 8B target model for clustering purposes.
    Used throughout Section 4.1 and Appendix D.3; Fig 3a gives empirical support, but the assumption is not theoretically justified.
  • domain assumption The effect of multi-step training on a domain can be summarized by a single unit-size gradient step (theta_Di = theta + grad_l_Di).
    The metric definition in Eq. (7) ignores learning rate, optimizer state, and multi-step dynamics; no analysis of this reduction is provided.
  • domain assumption A one-step exponential decay curve (Eq. 14) models downstream task loss trajectories.
    Stated in Section 4.3 and acknowledged as a limitation in the Conclusions; the authors note the pattern may not hold for all learning scenarios.
  • domain assumption The observable validation subset used for reweighting is representative and disjoint from the evaluated examples.
    Section 5.1.2 partitions a subset from the downstream validation set for reweighting, but the text does not state whether this subset is excluded from the reported evaluation.
  • standard math Johnson-Lindenstrauss random projection approximately preserves the gradient geometry needed for k-means clustering.
    Appendix A cites the JL lemma; the paper assumes the projected distances are sufficient for meaningful clustering of training effects.

how reviews work

0 comments
Cite this review

Pith. "Pith review of DIDS: Domain Impact-aware Data Sampling for Large Language Model Training." pith.science (2026). https://pith.science/paper/TUPS347F

@misc{pith2026250413227,
  author       = {Pith},
  title        = {Pith review of: DIDS: Domain Impact-aware Data Sampling for Large Language Model Training},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/TUPS347F}},
  note         = {Machine review of arXiv:2504.13227}
}
read the original abstract

Large language models (LLMs) are commonly trained on multi-domain datasets, where domain sampling strategies significantly impact model performance due to varying domain importance across downstream tasks. Existing approaches for optimizing domain-level sampling strategies struggle with maintaining intra-domain consistency and accurately measuring domain impact. In this paper, we present Domain Impact-aware Data Sampling (DIDS). To ensure intra-domain consistency, a gradient clustering algorithm is proposed to group training data based on their learning effects, where a proxy language model and dimensionality reduction are employed to reduce computational overhead. To accurately measure domain impact, we develop a Fisher Information Matrix (FIM) guided metric that quantifies how domain-specific parameter updates affect the model's output distributions on downstream tasks, with theoretical guarantees. Furthermore, to determine optimal sampling ratios, DIDS combines both the FIM-guided domain impact assessment and loss learning trajectories that indicate domain-specific potential, while accounting for diminishing marginal returns. Extensive experiments demonstrate that DIDS achieves 3.4% higher average performance while maintaining comparable training efficiency. The code is available at https://github.com/shiweijiezero/DIDS.

Figures

Figures reproduced from arXiv: 2504.13227 by the authors.

Figure 1
Figure 1. Overview of DIDS’s three-step process: (1) Domain repartition using gradient-based clustering with [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Effects of update frequency and irrelevant [PITH_FULL_IMAGE:figures/full_fig_p008_2.png] view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

12 extracted references · 12 canonical work pages

  1. [1]

    Spherical symmetry: For any orthogonal ma- tricesA,B∈O(d),RART andRBRT have iden- tical distributions

    work page

  2. [2]

    Orthogonality: The rows of R are approxi- mately orthogonal

    work page

  3. [3]

    The projection matrixR∈ Rh×s is constructed as follows:

    Unit length: Each row ofR is normalized to unit length. The projection matrixR∈ Rh×s is constructed as follows:

    work page

  4. [6]

    In ACL, Bangkok, Thailand

    Llamafactory: Unified efficient fine-tuning of 100+ language models. In ACL, Bangkok, Thailand. Association for Computational Linguistics. Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, et al. 2024. Lima: Less is more for align- ment. Advances in Neural Information Processing Systems, ...

    work page 2024

  5. [10]

    Generate entriesRij independently according to: Rij = ( +1/ √ t with probability 1/2 −1/ √ t with probability 1/2 where t = Ω(k/ε2) for dimension reduction pa- rameterk and error toleranceε

    work page

  6. [11]

    Normalize each column to ensure unit length: ˜Rj =Rj/||Rj||2 Dataset Samples Percentage (%) Tulu 3 Persona MATH 149,960 16.0 Evol CodeAlpaca 107,276 11.4 FLAN v2 89,982 9.6 NuminaMath-TIR 64,312 6.8 Tulu 3 Persona GSM 49,980 5.3 Tulu 3 WildGuardMix 50,000 5.3 Tulu 3 WildJailbreak 50,000 5.3 Tulu 3 Persona Python 34,999 3.7 Tulu 3 Persona IF 29,980 3.2 Tul...

    work page

  7. [12]

    DIDS automates this process and im- proves upon static optimal ratios through dynamic adaptation

    and Llama3 (Touvron et al., 2023), where mixed ratios are carefully tuned through extensive grid search. DIDS automates this process and im- proves upon static optimal ratios through dynamic adaptation. H Domain Partitioning Robustness Analysis To evaluate DIDS’s sensitivity to domain partition- ing quality, we conducted destructive experiments by artific...

    work page 2023

  8. [323]

    Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guil- laume Leclerc, and Aleksander Madry

    Springer. Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guil- laume Leclerc, and Aleksander Madry. 2023. Trak: Attributing model behavior at scale. arXiv. Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. 2020. Estimating training data influ- ence by tracing gradient descent. Advances in Neural Information Processing Systems, 33:19920–...

    work page 2023

Show all 12 references
  1. [2020]

    Measuring massive multitask language under- standing. arXiv. Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William W Cohen, and Xinghua Lu. 2019. Pubmedqa: A dataset for biomedical research question answering. arXiv. Feiyang Kang, Yifan Sun, Bingbing Wen, Si Chen, Dawn Song, Rafid ...

    2019

  2. [2024]

    Yuan Ge, Yilun Liu, Chi Hu, Weibin Meng, Shimin Tao, Xiaofeng Zhao, Hongxia Ma, Li Zhang, Box- ing Chen, Hao Yang, et al

    A framework for few-shot language model evaluation. Yuan Ge, Yilun Liu, Chi Hu, Weibin Meng, Shimin Tao, Xiaofeng Zhao, Hongxia Ma, Li Zhang, Box- ing Chen, Hao Yang, et al. 2024. Clustering and ranking: Diversity-preserved instruction selection through expert-aligned quality ...

    2024

  3. [2025]

    In ECCV, pages 304–

    Mm1: methods, analysis and insights from multimodal llm pre-training. In ECCV, pages 304–

  4. [2430]

    Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al

    PMLR. Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. 2020. Piqa: Reasoning about physical com- monsense in natural language. In AAAI, volume 34, pages 7432–7439. Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutan...

    2020 arXiv

Pith tools

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