1 Introduction↩︎

The progress of large language models depends on large-scale supervised corpora, yet the cost of curating and training on such data compounds across fine-tuning cycles, hyperparameter sweeps, and continual learning. The real burden is not any single run but the cumulative expenditure across the full development lifecycle. Corpus size, not model size, is often the binding constraint. Reducing it is therefore a systemic priority with direct implications for cost, carbon footprint, and data governance.

Dataset distillation (DD), first proposed by [1], replaces a large dataset \(\mathcal{D}\) with a smaller surrogate \(\tilde{\mathcal{D}} \ll \mathcal{D}\) such that a model trained on \(\tilde{\mathcal{D}}\) performs comparably to one trained on \(\mathcal{D}\). Put another way: what is the smallest textbook that teaches the same exam? This places DD squarely within data-centric AI, complementing model compression by operating on the data rather than the model itself. Early DD methods work in image space, optimizing synthetic pixels via meta-learning, gradient or trajectory matching [1]–[3]. Text breaks these assumptions: non-differentiable token decoding blocks gradient-based synthesis, and embedding spaces are tightly coupled to specific architectures. Adapting DD to text has therefore required fundamentally different strategies.

Existing text DD methods [4]–[8] fall short in at least one of the following: at extreme compression, uniform weighting exhausts the budget on uninformative samples; their objectives lack task alignment or global optimization; and their outputs are embeddings, neither auditable nor transferable. We survey this landscape in §2.1 and show that no prior method simultaneously addresses all of the above. These gaps motivate three desiderata for TAKE: (1) knowledge-based weighting — weighting samples by their task-aligned downstream contribution, directing the budget toward the most informative ones; (2) a dataset-level distillation objective — optimizing over the full training distribution rather than batch-by-batch, ensuring global coverage; and (3) human-readable output — a distilled corpus that is directly inspectable, auditable, and transferable across architectures.

We propose Trajectory-Aware Knowledge Estimation (TAKE), the first text DD method to satisfy all three desiderata. We address the first via influence functions [9], which quantify each sample’s counterfactual effect on the downstream loss and provide a theoretically grounded, task-aligned weighting scheme (§3). However, scoring at a single checkpoint introduces the hard-sample bias (HSB): at convergence, noisy samples dominate influence scores while clean and moderate samples, the most valuable for a robust distilled dataset, are suppressed. TAKE corrects this by integrating influence scores along the full training trajectory into a per-sample knowledge score that captures clean and moderate samples while down-weighting noisy ones. Then, TAKE fine-tunes an LLM to generate a pool of human-readable candidate instances. The knowledge scores and synthetic candidates jointly feed a discrete Optimal Transport (OT) objective that globally selects the distilled corpus. Concretely, we contribute:

• Theory: We formalize the hard-sample bias (HSB) in the text DD setting, prove that single-checkpoint influence scores are biased toward noisy samples, and derive a formal gap bound for knowledge-reweighted distribution matching.

• Method: TAKE is the first text DD method to jointly address sample weighting and distributional coverage—via trajectory-integrated knowledge scores and a discrete OT objective aligned with the DD goal.

• Empirical: TAKE matches or exceeds prior text DD methods across six language benchmarks at extreme compression (0.1% or 20 instances/class), producing human-readable distilled corpora that generalize across diverse backbone families.

The remainder of this paper is organised as follows: § 2 reviews related work; § 3 presents theoretical motivation; § 4 details TAKE and our DD pipeline; § 5 reports experiments and discussion; § 6 reflects limitations and ethical concerns; § 7 concludes.

2 Related Works↩︎

2.1 Text Dataset Distillation↩︎

Dataset distillation was introduced by [1] as a bi-level meta-learning problem: synthesize a small set of images such that a model trained on them matches one trained on the full dataset. Subsequent methods improve scalability and fidelity via gradient matching [2], distribution matching [10], and trajectory matching [3], [11]. These methods share a common assumption: data is continuous and differentiable, so the synthetic set can be directly optimised via backpropagation. Text violates the differentiability assumption: token decoding is non-differentiable, outputs must be semantically coherent, and embedding spaces are tightly coupled to specific architectures.

Adapting DD to text has therefore required fundamentally different strategies, from meta-learning on synthetic embeddings to LLM-based generation of readable corpora, each generation trading one limitation for another. SLDD and DDTC [4], [5] directly adapted the meta-learning formulation [1], optimizing synthetic text embeddings so that a model trained on them minimizes task loss. Of the two, only SLDD partially recovers interpretability by mapping distilled embeddings to per-token nearest bag-of-words. DDAL [6] optimizes the attention labels of a BERT model to match class-wise dataset statistics, achieving partial task alignment but producing outputs that are both architecture-locked and uninterpretable. A second generation of methods addressed interpretability by rethinking the pipeline entirely. DiLM [7] generates readable candidate texts via an LLM and selects among them by gradient matching over the training corpus, followed by k-centre selection; but matching gradients at discrete checkpoints is a proxy for task loss, leaving task alignment only partial. DaLLME [8] learns an inverse mapping from distilled embeddings back to synthetic text, recovering interpretability and attempting global coverage via k-centroids clustering; but with no task-loss objective, task alignment is absent. Table 1 maps this progression: each generation addressed some gaps while leaving others open.

8pt

A subtler gap cuts across all methods: every existing approach applies uniform weighting, treating every training sample as equally informative. At moderate distillation budgets this is defensible, but extreme compression is precisely the regime where distillation matters most, and there the assumption fails. Samples vary substantially in their influence on learning (clean vs.noisy, easy vs.hard), yet no prior method accounts for this structure. Without importance weighting, the distillation budget is exhausted on uninformative samples, biasing selection away from the boundary cases that matter most for robustness.

2.2 Influence Functions and the Hard-Sample Bias↩︎

Influence functions (IF) [9], [12] quantify the leave-one-out effect of a training point via \(\mathcal{I}(z) = \nabla^\top H_\theta^{-1} \nabla\), but depend on a single checkpoint and require expensive second-order computation. Trajectory-based methods [13]–[15] extend attribution across checkpoints but are test-conditioned: designed to attribute a fixed test prediction. The self-influence approximation adapts them to corpus-level scoring by using each training point as its own query. However, self-influence is evaluated only once at the final parameters, introducing a systematic bias that corrupts the corpus-level scores.

We refer to this as the hard-sample bias (HSB): at convergence, clean samples have near-zero gradient by definition, so noisy or hard samples dominate influence scores. This failure mode is well-documented in adjacent work: [16] show that networks fit clean patterns before noisy ones; dataset cartography [17] confirms that single-checkpoint statistics conflate difficulty with noise; and Co-teaching [18] documents systematic overselection of noisy samples in gradient-based curricula. In standard training, HSB can often be mitigated via regularization; in the extreme low-data regime of DD, however, it amplifies: biased selection suppresses the informative samples essential for model robustness. To our knowledge, no prior text DD method has identified or corrected for it.

2.3 Optimal Transport for Distribution Matching↩︎

Distribution matching has been proposed as a tractable DD surrogate [19]: if the distilled and training sets induce equal expected gradients along the optimization trajectory, the trained models converge to comparable parameters. This is a sufficient condition, but not an equivalence (the formal gap bound appears in Theorem 1). Common divergences have practical shortcomings at the extreme budgets relevant here. MMD [20] lacks sample-level assignment structure, requiring auxiliary scoring heuristics that reintroduce the global optimization gap that distribution matching is meant to close. Kullback–Leibler and Jensen–Shannon divergences require density estimation, which is unreliable in the low-data regime [21].

OT in the Monge-Kantorovich sense avoids both pitfalls: it operates natively on discrete measures, respects the metric geometry of the embedding space, and produces an explicit transport plan that directly assigns training samples to prototypes—no auxiliary heuristic needed. Entropic regularization via Sinkhorn-Knopp [22] makes this tractable at corpus scale. OT has been widely applied in NLP and machine learning [23]–[26], but never to text DD, and never combined with a knowledge-reweighted source distribution in any DD setting. TAKE fills both gaps: \(\kappa_n\) defines a non-uniform source measure, and the discrete entropic OT plan selects prototypes covering the knowledge-weighted training distribution without any auxiliary selection step.

3 Theoretical Motivation↩︎

3.1 Reweighted Distribution Matching for Dataset Distillation↩︎

The dataset distillation (DD) objective seeks \(\tilde{\mathcal{D}}\) with \(|\tilde{\mathcal{D}}| = K \ll N\) minimizing the loss on the original data: \[\min_{\tilde{\mathcal{D}}}\; \mathbb{E}_{\theta_0 \sim p(\theta_0)} \left[ \mathbb{E}_{z \sim \mathcal{D}} \left[ \ell\!\left(z;\; F(\theta_0;\, \tilde{\mathcal{D}}, \eta, T)\right) \right] \right] \label{eq:dataset-distillation}\tag{1}\] where \(F(\theta_0; \tilde{\mathcal{D}}, \eta, T)\) denotes the model obtained by running \(T\) steps of gradient descent with learning rate \(\eta\) on \(\tilde{\mathcal{D}}\) from initialization \(\theta_0\). This bi-level problem is intractable for discrete text. We relax it to distribution matching via the gradient-field equivalence argument [19]: if \(\tilde{\mathcal{D}}\) induces the same expected gradient field as \(\mathcal{D}\) throughout training, the two runs converge to comparable parameters. Since \(\nabla_\theta \mathbb{E}_{z \sim P}[\ell(z;\theta)] = \mathbb{E}_{z \sim P}[\nabla_\theta \ell(z;\theta)]\), matching distributions in a task-relevant feature space controls the gradient field, making distribution matching a tractable surrogate for Eq. 1 . This yields the distribution matching (DM) objective: \[\min_{|\tilde{\mathcal{D}}| = K}\; d\!\left(P_{\tilde{\mathcal{D}}},\; P_{\mathcal{D}}\right) \label{eq:distribution-matching}\tag{2}\] for a divergence \(d\) in a task-relevant feature space. Eq. 2 is a surrogate, not an equivalence; the formal gap bound appears as Theorem 1 once \(d\) is specified.

At extreme compression (\(\rho = K/N \ll 1\)), uniform matching allocates prototypes proportional to density — but not all regions contribute equally to learning. Dense easy regions consume budget that would be better spent on moderate and boundary samples near decision boundaries, which carry the most information for model robustness. We instead reweight the source by per-sample importance: define \(w_n \geq 0\), \(\sum_n w_n = 1\), inducing \(P^w = \sum_n w_n \delta_{z_n}\), yielding the reweighted distribution matching (RDM) objective: \[\min_{|\tilde{\mathcal{D}}| = K}\; d\!\left(P_{\tilde{\mathcal{D}}},\; P^w\right) \label{eq:reweighted-distribution-matching}\tag{3}\] The weights \(w_n\) must reflect downstream learning contribution, not geometry alone, motivating gradient-based estimation. Influence self-scores (§2.2) are the natural candidate: they are task-aligned by construction and theoretically grounded. We now show that the simplest instantiation—\(w_n \propto \mathcal{I}_\text{self}(z_n)\) at a single checkpoint, select top-\(K\)—fails on both counts: the weights are biased toward noisy samples (Proposition 1), and independent top-\(K\) selection ignores distributional coverage (Remark 3).

3.2 The Naive Instantiation Fails: Two Open Problems↩︎

Define the influence self-score at checkpoint \(\hat{\theta}\) as \[\mathcal{I}_\text{self}(z_n; \hat{\theta}) = \nabla_\theta \ell(z_n; \hat{\theta})^\top H_{\hat{\theta}}^{-1} \nabla_\theta \ell(z_n; \hat{\theta}), \label{eq:influence-self}\tag{5}\] where \(H_{\hat{\theta}} = \frac{1}{N}\sum_n \nabla^2_\theta \ell(z_n; \hat{\theta})\) is the empirical Hessian, specializing [9] by setting the test point equal to the training point. High \(\mathcal{I}_\text{self}\) identifies atypical examples: samples that are difficult to fit and exert outsized influence on the model’s parameters, which in practice are the noisy ones. This suggests a naive pipeline: set \(w_n \propto \mathcal{I}_\text{self}(z_n; \hat{\theta})\) and return \(\arg\max_{|\mathcal{A}|=K} \sum_{i \in \mathcal{A}} w_i\). We identify two orthogonal failures of this strategy.

The first failure is a bias in the weights themselves.

Because clean gradients vanish at convergence by definition of a local minimum, the bias is structural: it cannot be corrected by post-hoc rescaling of single-checkpoint scores, but only by integrating gradient signals across the training trajectory. The binary partition simplifies exposition; in practice the bias is monotone in \(\|\nabla\|\), as the spectral bounds above show directly.

Even with unbiased weights, a second failure remains: top-\(K\) selection is blind to distributional coverage.

The two failures compound: biased \(w_n\) misdirects a population-level solver toward noisy regions; correct \(w_n\) fed into an independent selector still produces redundancy. They respectively yield two open problems:

1. Biased weighting. Estimate \(w_n\) by integrating gradient signals across trajectory \(\{\theta_t\}_{t=0}^{T-1}\), reducing \(\Delta_\mathrm{HSB}\)—which we will address in §4.2 (formalized as Proposition 2; Appendix A.2).

2. Independent selection. Solve Eq. 3 directly for a metric \(d\) that jointly optimises coverage of \(P^w\)—which we will address in §4.4.

4 Methodology↩︎

4.1 Problem Setup and Framework Overview↩︎

TAKE instantiates the RDM objective (Eq. 3 ) by solving the two open problems from §3.2. Let \(\mathcal{D} = \{z_n\}_{n=1}^{N}\) denote the training corpus, \(\mathcal{D}_\text{synth} = \{\tilde{z}_c\}_{c=1}^{M}\) the synthetic candidate pool, and \(\tilde{\mathcal{D}} \subset \mathcal{D}_\text{synth}\) the distilled set of budget \(K\). Indices \(n \in [N]\) and \(t \in [T]\) range over training samples and checkpoints respectively. TAKE proceeds in three stages:

1. Score—train a logistic probe on \(\mathcal{D}\) across \(T\) checkpoints, compute influence matrix \(I \in \mathbb{R}^{N \times T}\), apply reciprocal reweighting and temporal weighting to obtain trajectory-aware knowledge scores \(\kappa_n\).

2. Generate—fine-tune a small language model on \(\mathcal{D}\) and sample to produce the synthetic candidate pool \(\mathcal{D}_\text{synth}\).

3. Distill—solve discrete OT with source weights \(\kappa_n\) over \(\mathcal{D}_\text{synth}\) to obtain the distilled corpus \(\tilde{\mathcal{D}}\) of budget \(K\).

The three stages are decoupled and independently replaceable, with \(\phi\) shared across Stages 1 and 3 as a natural consistency choice. The following sections detail each in turn.

4.2 Gradient-Based Influence Along the Trajectory↩︎

The hard-sample bias (HSB; §3.2) arises because single-checkpoint influence scores overemphasize noisy or difficult samples, whose gradients remain large at convergence. The fix is to integrate gradient information across the full training trajectory, producing a trajectory-aware knowledge score \(\kappa_n\) per sample. We build \(\kappa_n\) in two steps: constructing the influence matrix \(I\) (§4.2.1), then collapsing it into a per-sample scalar score via reciprocal reweighting and temporal weighting (§4.2.2). The resulting \(\kappa_n\) provides a task-aligned, bias-corrected weight for \(P^w\).

4.2.1 Influence Matrix↩︎

Our goal is to construct \(I \in \mathbb{R}^{N \times T}\) whose entry \(I_{(n,t)}\) reflects how much sample \(z_n\) contributes to learning at checkpoint \(\theta_t\). The influence self-score (§3.2, Eq. 5 ) is the natural candidate, but requires Hessian inversion at a converged checkpoint and is undefined mid-trajectory. We address both limitations by adapting the per-sample gradient norm from TracIn [13] to the self-influence setting, querying training points instead of test points: \[I_{(n,t)} = \|\nabla_\theta \ell(z_n;\,\theta_t)\|. \label{eq:influence-matrix}\tag{7}\] As the \(H=I\) special case of the influence self-score (Remark 2), \(I_{(n,t)}\) exhibits the same HSB structure—the bias that \(\kappa_n\) corrects via reciprocal inversion and trajectory integration (Proposition 2). Scores are column-wise normalized \(I_{(n,t)} \leftarrow I_{(n,t)} / \sum_{n'} I_{(n',t)}\) for comparability across checkpoints.

In practice, we score samples with a linear head \(W\) atop a pre-trained, frozen sentence encoder \(\phi\), giving \(\theta = (W, \phi)\). The signal quality comes not from the linear head but from \(\phi(x_n)\): \(\phi\) is pre-trained on large-scale corpora and \(W\) is fine-tuned on the full task corpus, so the representations, and hence the gradient norms, are both semantically rich and task-aligned by construction. Following the standard practice in scalable influence estimation [27], we approximate full-model influence via the last layer—here, the linear head \(W\): \[\nabla_\theta \ell(z_n;\,\theta_t) \approx \nabla_W \ell(z_n;\,\theta_t), \label{eq:last-layer}\tag{8}\] trading a small bias for tractability. By the Goodfellow outer-product identity [28], \(\nabla_W \ell = g_n \phi(x_n)^\top\) where \(g_n\) is the output gradient. The Goodfellow trick vectorizes the batch backward pass, avoiding per-sample loops. Adapting to a new task is trivial: retrain \(W\) on the new corpus; \(\phi\) and the influence pipeline are unchanged.

4.2.2 Knowledge Score and Temporal Weighting↩︎

To correct the HSB, we require a score that is high for well-fitted samples and low for noisy ones—the opposite of what \(I_{(n,t)}\) directly provides. Since a well-fitted sample has a small gradient norm, taking the reciprocal naturally inverts this ranking. We further weight each checkpoint by a decreasing temporal kernel \(k : \{0,\ldots,T-1\} \to \mathbb{R}_{>0}\), giving: \[\kappa_n = \sum_{t=0}^{T-1} \frac{k(t)}{I_{(n,t)}}, \label{eq:kappa}\tag{9}\]

A decreasing kernel is the natural choice: early checkpoints capture easy samples fitting rapidly, mid checkpoints offer the strongest signal, and late checkpoints see the gap narrow as memorization sets in [16]. Thus, concentrating mass on the early-to-mid phase and discounting late memorization directly corrects the HSB structure identified in Proposition 1. By default we use the exponential kernel \(k(t) = e^{-\lambda t / T}\) (\(\lambda > 0\)), with \(\lambda\) treated as a hyperparameter. Alternative kernels (linear, cosine) are drop-in replacements, and Proposition 2 guarantees bias reduction holds for any strictly decreasing \(k\). The resulting \(\kappa_n\) summarizes each sample’s cumulative, bias-corrected learning contribution across the full trajectory, resolving open problem [open-problem-weight]. Normalized source weights \(w_n = \kappa_n / \sum_{n'} \kappa_{n'}\) define the reweighted distribution \(P^w\) and serve as input to the OT selection step (§4.4).

4.3 Synthetic Candidate Pool Generation↩︎

Rather than selecting directly from \(\mathcal{D}\), TAKE draws \(\tilde{\mathcal{D}}\) from a synthetic candidate pool \(\mathcal{D}_\text{synth}\). Direct subsampling has two problems: at extreme compression, verbatim records raise privacy concerns, and the budget \(K\) is too small to cover the full training support. We generate \(\mathcal{D}_\text{synth}\) by fine-tuning a small language model on \(\mathcal{D}\) and sampling from it. The fine-tuned LM generalizes beyond memorized training instances, and stochastic sampling produces a diverse pool that extends beyond the training support. \(\mathcal{D}_\text{synth}\) serves solely as the selection pool and is never appended to downstream training data.

TAKE does not assume \(P_{\mathcal{D}_\text{synth}} = P_{\mathcal{D}}\); this distributional gap is beneficial. First, coverage expansion fills undersampled regions of \(\mathcal{D}\), improving generalization of the distilled corpus. Second, model fluency bias suppresses noisy samples that HSB would otherwise over-select, yielding a cleaner candidate pool. Third, the distribution mismatch physically decouples \(\tilde{\mathcal{D}}\) from original records, reducing verbatim exposure risk. OT handles the distributional gap gracefully: it routes mass toward \(P^w\) regardless of pool distribution, so \(P_{\tilde{\mathcal{D}}}\) approximates \(P^w\) provided \(\mathcal{D}_\text{synth}\) covers the support of \(P_{\mathcal{D}}\) (Assumption 1):

Under Assumption 1, \(W_1(P^w, P_{\mathcal{D}_\text{synth}}) \leq r\), so the optimal transport cost is at most \(r\) before any prototype selection. A practitioner can monitor coverage directly by reporting the mean nearest-neighbor distance \(\bar{r} = \frac{1}{N}\sum_n \min_c \|\phi(z_n) - \phi(\tilde{z}_c)\|\).

4.4 Distributional Prototype Selection via Discrete Optimal Transport↩︎

Both \(\mathcal{D}\) and \(\mathcal{D}_\text{synth}\) are finite sample sets, so the matching problem is naturally over empirical measures. Discrete OT operates directly on \(\mu = \sum_n w_n \delta_{z_n}\) and \(\nu = \frac{1}{|\mathcal{D}_\text{synth}|}\sum_c \delta_{\tilde{z}_c}\), requiring no density estimation and yielding an explicit transport plan with direct prototype-level assignments.

Let \(C_{nc} = \|\phi(z_n) - \phi(\tilde{z}_c)\|^2\) be the cost matrix in a task-relevant embedding space. Bare Kantorovich OT, \(\min_{\gamma \in \Pi(\mu,\nu)} \langle C, \gamma \rangle\), produces a sparse, non-differentiable plan and scales as \(O(N^3)\). We add entropic regularization to obtain a strictly convex, smooth objective solvable via Sinkhorn–Knopp iterations at \(O(N \cdot |\mathcal{D}_\text{synth}|)\) per step: \[\gamma^* = \arg\min_{\gamma \in \Pi(\mu,\nu)} \langle C, \gamma \rangle + \varepsilon \cdot \mathrm{KL}(\gamma \| \mu \otimes \nu). \label{eq:ot}\tag{11}\] The soft coupling distributes mass across nearby candidates rather than hard one-to-one assignment, reducing sensitivity to embedding errors and encouraging diversity in \(\tilde{\mathcal{D}}\). The \(K\) candidates with highest received mass \(\sum_n \gamma^*_{nc}\) are selected as \(\tilde{\mathcal{D}}\), resolving open problem [open-problem-select]. The following theorem closes the theoretical chain from RDM (Eq. 3 ) back to DD (Eq. 1 ), bounding the downstream loss gap in terms of the OT cost.

The bound is stated with respect to \(P^w\) rather than the uniform \(P_\mathcal{D}\) — this is intentional, as \(P^w\) is the designed optimization target of the RDM objective; since \(P^w\) reweights \(P_\mathcal{D}\) over the same support, minimizing \(W_1(P^w, P_{\tilde{\mathcal{D}}})\) controls coverage of \(P_\mathcal{D}\) up to the reweighting gap bounded in Remark 1.

Together, Propositions 1–2 and Theorem 1 establish the full theoretical basis for TAKE: HSB is structural and unavoidable at a single checkpoint (Proposition 1); trajectory integration strictly reduces it (Proposition 2); and minimizing the OT cost directly controls the downstream loss gap under \(P^w\) (Theorem 1; the gap to \(P_\mathcal{D}\) is bounded by Remark 1). Algorithm 1 consolidates the full TAKE pipeline.



5 Experiments↩︎

5.1 Setup↩︎

We evaluate on two linguistic tasks: AG News (\(N=120\) K), IMDb (\(N=25\) K), SST-2 (\(N=67\) K) for classification; MNLI-m (\(N=393\) K), QNLI (\(N=105\) K), and QQP (\(N=364\) K) for natural language inference (NLI). We compare against four baselines:

• Random: uniform random selection at the target budget.

• EDA [30]: word-level augmentation (10\(\times\)) applied to the Random set.

• DiLM [7]: LLM-generated candidates selected by gradient matching across k-centre selection. It evaluated on both classification and NLI.

• DaLLME [8]: synthetic texts inverted from distilled embeddings. It evaluated on classification only.

We evaluate TAKE at two budgets (20/cls and 0.1%) to match DiLM’s and DaLLME’s reported settings, respectively. To match downstream evaluation, downstream models for classification are Logistic Regression (LR) with TF-IDF [31] features, TextCNN [32] and TextRNN [33] with GloVe [34] word embeddings, and BERT [35]. For NLI, we evaluate on BERT and Siamese LR. Siamese LR is a variant of LR: two input branches to independently encode sentence pairs, and a logistic head over the concatenated representations, tailored as a weak learner for NLI. The primary metric is downstream accuracy on classification/NLI tasks, reported as mean\(_{\pm{\text{s.e.}}}\) across five runs.

In terms of implementation, TAKE extracts knowledge scores \(\kappa_n\) by running \(T=20\) training checkpoints with exponential kernel \(\lambda^* = T/T_m\), using all-MiniLM-L6-v2 [36] as the frozen encoder \(\phi\). The synthetic pool uses Gemma-3 [37] (270 M) fine-tuned on \(\mathcal{D}\) to generate \(|\mathcal{D}_\text{synth}|=5{,}000\) samples. OT selection embeds via the same \(\phi\) and uses Sinkhorn with \(\varepsilon=0.05\) for 200 iterations. All experiments fit on a single NVIDIA V100 (40 GB).

5.2 Classification↩︎

Accuracy gains and their source. TAKE (0.1%) outperforms DaLLME on LR by \(+2.5\) on AG News and \(+1.8\) on IMDb. The gap is widest at the LR level and narrows for neural backbones, consistent with TAKE’s backbone-agnostic selection—trajectory scores are computed independently of the downstream model. TAKE minimizes the OT cost in embedding space, directly controlling the downstream loss gap under \(P^w\) via Theorem 1 (the gap to the original \(P_\mathcal{D}\) is bounded by Remark 1). Since \(\kappa_n\) scores are computed independently of the evaluation backbone, this coverage of \(P^w\) transfers across model families.

Stability across runs. The wide standard errors for Random (IMDb LR \(\pm5.5\), BERT \(\pm4.1\)) reveal that random selection at 0.1% can land anywhere from a coherent corpus to a degenerate one. TAKE’s errors are at least \(5\times\) narrower (\(\pm0.1\)–\(2.2\)), a direct consequence of OT-based selection: the transport plan distributes selected prototypes across \(P^w\) rather than drawing them independently, suppressing the variance that Random incurs by chance.

SST-2 vs.DiLM. At 20/cls, TAKE matches DiLM on BERT (\(81.15\) vs.\(80.30\), within s.e.) while also providing results for LR, TextCNN, and TextRNN where DiLM has none—indicating that the OT selection step generalizes across backbone families without retraining the scorer.

5.3 Natural Language Inference↩︎

Gains and their interpretation. DaLLME has no published NLI experiments and is excluded from this comparison. In comparison with DiLM, TAKE outperforms DiLM by \(+2.5\) on MNLI-m and \(+3.4\) on QQP on a BERT learner, and provides competitive results for QNLI. The LR gap (\(+8.5\) on MNLI-m, \(+5.8\) on QNLI, \(+7.6\) on QQP relative to Random) is particularly informative: Siamese LR has no capacity to compensate for a poor corpus—every percentage point of accuracy here must come from the quality of \(\tilde{\mathcal{D}}\) itself. That TAKE’s Siamese LR exceeds the Random BERT baseline on both QQP (\(60.04\) vs.\(59.10\)) and QNLI (\(57.96\) vs.\(57.19\)) suggests the selected sentence pairs are genuinely more informative, not merely better distributed.

The NLI ceiling gap and what it implies. All distilled methods fall short of the full-dataset ceiling (\(86.8\) / \(82.5\) / \(89.5\) for MNLI-m, QNLI, QQP), a gap larger than on single-sentence tasks. NLI imposes a joint constraint: the synthetic pool must contain pairs whose relationship correctly reflects each class label at extreme budget (\(\approx\!60\) pairs for MNLI-m, \(\approx\!40\) for QQP and QNLI, at 20/cls). This is a harder coverage requirement than single-sentence classification, and the synthetic pool must cover all relationship types within it. This suggests that pool generation and budget are the binding constraints, rather than scoring or selection.

Stability. On MNLI-m and QQP, TAKE’s errors (\(\pm0.1\)–\(0.5\)) are \(5\)–\(7\times\) narrower than DiLM’s (\(\pm2.2\)–\(2.6\)), mirroring the pattern on classification. This is consistent with the OT transport plan enforcing coverage of \(P^w\) rather than selecting samples independently: deterministic coverage implies that any two runs of TAKE over the same pool yield nearly identical corpora, whereas DiLM’s gradient-matching step introduces variance through its dependence on a randomly initialized model.

5.4 Component Contributions↩︎

To isolate the contribution of each design choice, we ablate TAKE’s components on AG News and QQP as representatives for the classification and NLI tasks, in Table ¿tbl:tab:results-abla-components?. TAKE improves over Random by \(+10.7\) on AG News LR and \(+10.6\) on BERT; \(k\)-means, by comparison, accounts for only part of this gain (\(+6.5\) LR, \(+1.9\) BERT), while our discrete OT selection step alone (TAKE without \(\kappa\)) already achieves \(+9.4\) LR and \(+5.0\) BERT from the same baseline. The same pattern holds on QQP: \(k\)-means gains \(+2.9\) / \(+3.4\) (Siamese LR / BERT) whereas TAKE without \(\kappa\) gains \(+6.4\) / \(+5.7\). Adding \(\kappa\) knowledge scores delivers the remaining improvement, modestly for Siamese LR (\(+1.2\)) but substantially for BERT (\(+5.6\) on AG News, \(+2.6\) on QQP), suggesting that \(\kappa\) primarily benefits models with enough capacity to exploit the harder, informative examples that high-\(\kappa\) samples represent.

5.5 Kernels for Temporal Weighting↩︎

Table ¿tbl:tab:results-abla-kernels? ablates the choice of temporal kernel defined in §4.2.2. The last-checkpoint baseline is the critical exception: it falls below the no-\(\kappa\) baseline across all backbones, most severely on BERT (\(-2.7\) on AG News, \(-1.5\) on QQP). This is precisely the failure mode identified in Proposition 2: when \(\kappa\) is extracted at convergence only, the learner has fully memorized difficult samples, so the hard-sample bias is intensified rather than corrected, and the selected corpus is skewed towards harder, less representative examples. All other trajectory-based strategies lessen this effect and outperform the no-\(\kappa\) baseline (OT only), with decreasing kernels (linear, cosine, exponential) doing so most strongly by concentrating mass on the early-to-mid training phase where the bias-correcting signal is strongest. The exponential kernel (TAKE default) achieves the best result in all cases, consistent with the theoretical preference for early-phase concentration of influence.

6 Limitations & Ethical Considerations↩︎

We identify three limitations of the current method. (1) Gradient norm proxy. \(I_{(n,t)}\) uses plain gradient norms rather than Fisher-preconditioned scores; the Goodfellow trick [28] keeps per-sample cost at \(O(C + d_\phi)\) per checkpoint. K-FAC [38] or EKFAC [39] preconditioning are natural extensions when calibrated magnitude scores are needed. (2) Encoder dependence. The OT cost matrix inherits the geometry of \(\phi\); a mismatched encoder can distort transport costs. The \(L\)-Lipschitz assumption on \(\ell\) (Theorem 1) holds approximately when \(\phi\) is task-aligned and embeddings are normalized, but remains an approximation for cross-entropy over discrete text. (3) Synthetic pool coverage. Pool diversity is bounded by the generator’s domain coverage. Fine-tuning on \(\mathcal{D}_\text{train}\) mitigates this in practice.

Two ethical considerations arise in the deployment of TAKE. (1) Bias propagation. Distillation concentrates the statistical properties of \(\mathcal{D}_\text{train}\), including demographic biases. We recommend corpus auditing before distillation. (2) Privacy risk. Language models fine-tuned on sensitive data may memorize training examples, a risk that synthetic generation alone does not fully eliminate. We recommend distance-to-closest-record (DCR) screening of \(\mathcal{D}_\text{synth}\) to verify that synthetic samples are not verbatim copies; for regulated domains, differential privacy at the fine-tuning stage provides a formal guarantee against memorization.

7 Conclusion↩︎

We identified the hard-sample bias (HSB) as the core failure mode of gradient-based text dataset distillation: at convergence, noisy samples dominate single-checkpoint influence scores, suppressing the informative ones that matter most. TAKE corrects this by integrating gradient norms across the training trajectory into per-sample knowledge scores, used as non-uniform source weights in a discrete Optimal Transport objective to select a compact, human-readable corpus. Theoretically, we provide end-to-end guarantees: the hard-sample bias is structural, trajectory integration strictly reduces it, and the downstream loss gap is bounded by the OT cost. TAKE outperforms baselines across six language benchmarks at extreme compression (0.1% or 20 samples/class), generalizing across diverse backbone families. Together, these results position text dataset distillation as a practical, data-centric tool for efficient and privacy-aware NLP, with direct implications for coreset construction and continual learning in resource-constrained settings.

8 Proofs↩︎

8.1 Proof of Proposition 1 (Hard-Sample Bias)↩︎

Setup. Let \(\mathcal{I}_\text{self}(z) = \nabla^\top H_{\hat{\theta}}^{-1} \nabla\), where \(\nabla \equiv \nabla_\theta \ell(z;\hat{\theta})\) and \(H_{\hat{\theta}}\) is the empirical Hessian, assumed symmetric positive definite (SPD) with \(\lambda_{\min}(H_{\hat{\theta}}) > 0\) (this holds whenever the model is not at a degenerate critical point). Partition \(\mathcal{D}\) into \(\mathcal{D}_\text{clean}\) and \(\mathcal{D}_\text{noisy}\) such that \(\|\nabla\| \leq \epsilon\) for all \(z \in \mathcal{D}_\text{clean}\) and \(\|\nabla\| \geq \delta\) for all \(z \in \mathcal{D}_\text{noisy}\), with \(\delta \gg \epsilon\).

Step 1 — Bound \(\mathcal{I}_\text{self}\) for each partition. Since \(H_{\hat{\theta}}\) is SPD with eigenvalues in \([\lambda_{\min}(H_{\hat{\theta}}),\, \lambda_{\max}(H_{\hat{\theta}})]\), the Rayleigh quotient gives: \[\frac{\|\nabla\|^2}{\lambda_{\max}(H_{\hat{\theta}})} \;\leq\; \mathcal{I}_\text{self}(z) \;\leq\; \frac{\|\nabla\|^2}{\lambda_{\min}(H_{\hat{\theta}})}.\] Substituting the gradient-norm assumptions (\(\|\nabla\|^2 \leq \epsilon^2\) for clean samples; \(\|\nabla\|^2 \geq \delta^2\) for noisy samples) and taking expectations (the bounds are uniform over each partition, so expectation preserves them): \[\begin{align} \mathbb{E}_{z \sim \mathcal{D}_\text{clean}}[\mathcal{I}_\text{self}(z)] &\leq \frac{\epsilon^2}{\lambda_{\min}(H_{\hat{\theta}})}, \\ \mathbb{E}_{z \sim \mathcal{D}_\text{noisy}}[\mathcal{I}_\text{self}(z)] &\geq \frac{\delta^2}{\lambda_{\max}(H_{\hat{\theta}})}. \end{align}\]

Step 2 — Lower bound on \(\Delta_\text{HSB}\). \[\Delta_\text{HSB} := \mathbb{E}_{z \sim \mathcal{D}_\text{noisy}}[\mathcal{I}_\text{self}(z;\hat{\theta})] - \mathbb{E}_{z \sim \mathcal{D}_\text{clean}}[\mathcal{I}_\text{self}(z;\hat{\theta})] \;\geq\; \frac{\delta^2}{\lambda_{\max}(H_{\hat{\theta}})} - \frac{\epsilon^2}{\lambda_{\min}(H_{\hat{\theta}})}.\] Substituting \(\lambda_{\min}(H_{\hat{\theta}}) = \lambda_{\max}(H_{\hat{\theta}})/\mathrm{cond}(H_{\hat{\theta}})\): \[\Delta_\text{HSB} \;\geq\; \frac{\delta^2 - \mathrm{cond}(H_{\hat{\theta}})\,\epsilon^2}{\lambda_{\max}(H_{\hat{\theta}})} \;>\; 0 \qquad \text{whenever} \quad \frac{\delta}{\epsilon} > \sqrt{\mathrm{cond}(H_{\hat{\theta}})}.\] Under the interpolation regime (an overparameterised model trained to near-zero training loss), \(\epsilon \to 0\): the bound simplifies to \(\delta^2/\lambda_{\max}(H_{\hat{\theta}}) > 0\) and the positivity condition is trivially satisfied (\(\delta/\epsilon \to \infty\)). \(\square\)

8.2 Proof of Proposition 2 (Trajectory Integration Reduces HSB)↩︎

Proposition 2. Under smooth loss and bounded gradient noise, the trajectory-integrated knowledge score \(\kappa_n\) produces a strictly smaller noisy-over-clean gap than the single last-checkpoint baseline \(\kappa^\dagger_n = 1/I_{(n,T-1)}\): \[\mathbb{E}_{z \sim \mathcal{D}_\text{noisy}}[\kappa_n] - \mathbb{E}_{z \sim \mathcal{D}_\text{clean}}[\kappa_n] \;<\; \mathbb{E}_{z \sim \mathcal{D}_\text{noisy}}[\kappa^\dagger_n] - \mathbb{E}_{z \sim \mathcal{D}_\text{clean}}[\kappa^\dagger_n].\]

Memorization onset \(T_m\). In the main paper (§4.2.2) we define \(T_m\) as the first checkpoint at which noisy-sample gradient norms begin to decrease. The formal continuous analogue is: \[T_m := \inf\!\left\{t : \mathbb{E}_{z \sim \mathcal{D}_\text{noisy}}\!\left[\|\nabla_\theta \ell(z;\theta_t)\|\right] < \mathbb{E}_{z \sim \mathcal{D}_\text{noisy}}\!\left[\|\nabla_\theta \ell(z;\theta_0)\|\right]\right\}. \label{eq:Tm}\tag{13}\] Assumptions.

• Monotone gradient decay for clean samples. For \(z \in \mathcal{D}_\text{clean}\), \(I_{(n,t)}\) is non-increasing in \(t\), converging toward zero (holds under \(L\)-smooth loss and a sufficiently small learning rate). The reciprocal \(1/I_{(n,t)}\) is therefore non-decreasing.

• Bounded gradient for noisy samples within the proof horizon. For \(z \in \mathcal{D}_\text{noisy}\), \(I_{(n,t)} \geq \delta > 0\) for all \(t \leq T_m\), where \(T_m\) is the memorization onset defined in Section 4.2.2 of the main paper. Beyond \(T_m\), noisy gradients decay; the decreasing kernel discounts these late checkpoints. With \(\lambda^* = T/T_m\), the kernel weight at \(T_m+1\) relative to \(t^*\) is \(e^{-\lambda^*(T_m+1-t^*)/T}\), which decreases as \(T_m \to T\), making the post-memorization contribution \(S_2\) negligible relative to \(S_1\).

• Post-memorization baseline inflation. At \(t = T-1 > T_m\), clean gradients are near zero (by A1) and noisy gradients have dropped under memorization. By symmetry near convergence, the noisy reciprocals exceed the clean reciprocals: \(\Delta_{T-1} := \mathbb{E}_\text{noisy}[1/I_{(n,T-1)}] - \mathbb{E}_\text{clean}[1/I_{(n,T-1)}] > 0\).

• Positive temporal kernel. \(k(t) > 0\) for all \(t\); decreasing so \(k(t^*) > k(T-1)\) for any \(t^* < T-1\).

Proof.

Define \(\Delta_t := \mathbb{E}_\text{noisy}[1/I_{(n,t)}] - \mathbb{E}_\text{clean}[1/I_{(n,t)}]\), so \[\mathbb{E}_\text{noisy}[\kappa_n] - \mathbb{E}_\text{clean}[\kappa_n] = \sum_{t=0}^{T-1} k(t)\,\Delta_t.\]

Step 1 — Sign structure. For \(1 \leq t \leq T_m\): (A1) and (A2) together give \(\Delta_t \leq 0\), strictly negative at some \(t^* \leq T_m\). For \(t > T_m\): \(\Delta_t\) may be positive (memorization onset decreases noisy norms).

Step 2 — Weighted sum. Split at \(T_m\): \[S_1 = \sum_{t=1}^{T_m} k(t)\,\Delta_t \;<\; 0 \qquad \text{and} \qquad S_2 = \sum_{t=T_m+1}^{T-1} k(t)\,\Delta_t \;\geq\; 0.\] Since the kernel is decreasing, \(S_2 \leq k(T_m+1) \cdot \sum_{t > T_m} \Delta_t\), which is dominated by \(|S_1|\) for sufficiently large \(\lambda\). Hence \(\sum_t k(t)\,\Delta_t < 0\).

Step 3 — Baseline gap. By (A3), \(\mathbb{E}_\text{noisy}[\kappa^\dagger_n] - \mathbb{E}_\text{clean}[\kappa^\dagger_n] = \Delta_{T-1} > 0\).

Combining. \[\mathbb{E}_\text{noisy}[\kappa_n] - \mathbb{E}_\text{clean}[\kappa_n] \;<\; 0 \;<\; \Delta_{T-1} = \mathbb{E}_\text{noisy}[\kappa^\dagger_n] - \mathbb{E}_\text{clean}[\kappa^\dagger_n]. \quad \square\]

8.3 Distribution Matching Gap Bound (Theorem 1)↩︎

Cited in Section 4.4 of the main paper. Closes the chain from (RDM) back to (DD).

Goal. Let \(f^*_\mathcal{S}\) be the model trained on \(\mathcal{S}\) and \(\mathcal{L}(\mathcal{S}) = \mathbb{E}_{z \sim P^w}[\ell(z; f^*_\mathcal{S})]\) be the downstream loss under \(P^w\). Show: \[|\mathcal{L}(\tilde{\mathcal{D}}) - \mathcal{L}(\mathcal{D})| \;\leq\; 2L \cdot W_1(P^w, P_{\tilde{\mathcal{D}}}) + \beta_K.\]

Assumptions.

• \(\ell(\cdot;\,f)\) is \(L\)-Lipschitz in the embedding metric \(\|\phi(\cdot) - \phi(\cdot)\|\) for all \(f\).

• The model class has uniform stability constant \(\beta_K = O(1/K)\) [29]: \(\sup_z |\ell(z; f^*_{\tilde{\mathcal{D}}}) - \ell(z; f^*_\mathcal{D})| \leq \beta_K\). Standard ERM algorithms (logistic regression, SVM) satisfy this for bounded loss; as \(K \to \infty\), \(\beta_K \to 0\).

• The ground metric for \(W_1\) is \(c(z,z') = \|\phi(z) - \phi(z')\|\).

Step 1 — Kantorovich–Rubinstein bound for a fixed model. By Kantorovich–Rubinstein duality [40] under (B3), for any fixed \(f\): \[\left| \mathbb{E}_{z \sim P^w}[\ell(z; f)] - \mathbb{E}_{z \sim P_{\tilde{\mathcal{D}}}}[\ell(z; f)] \right| \;\leq\; L \cdot W_1(P^w, P_{\tilde{\mathcal{D}}}).\]

Step 2 — Triangle inequality decomposition. \(\mathcal{L}(\tilde{\mathcal{D}})\) uses \(f^*_{\tilde{\mathcal{D}}}\) while \(\mathcal{L}(\mathcal{D})\) uses \(f^*_\mathcal{D}\); both are evaluated under \(P^w\) but with different models, so Step 1 does not apply directly. Insert two intermediates and apply the triangle inequality twice: \[\begin{align} |\mathcal{L}(\tilde{\mathcal{D}}) - \mathcal{L}(\mathcal{D})| &\leq \underbrace{ \left| \mathbb{E}_{P^w}[\ell(z;f^*_{\tilde{\mathcal{D}}})] - \mathbb{E}_{P_{\tilde{\mathcal{D}}}}[\ell(z;f^*_{\tilde{\mathcal{D}}})] \right| }_{\text{(I) distribution shift, fixed }f^*_{\tilde{\mathcal{D}}}} \notag \\ &\quad + \underbrace{ \left| \mathbb{E}_{P_{\tilde{\mathcal{D}}}}[\ell(z;f^*_{\tilde{\mathcal{D}}})] - \mathbb{E}_{P_{\tilde{\mathcal{D}}}}[\ell(z;f^*_\mathcal{D})] \right| }_{\text{(II) model difference, fixed }P_{\tilde{\mathcal{D}}}} \notag \\ &\quad + \underbrace{ \left| \mathbb{E}_{P_{\tilde{\mathcal{D}}}}[\ell(z;f^*_\mathcal{D})] - \mathbb{E}_{P^w}[\ell(z;f^*_\mathcal{D})] \right| }_{\text{(III) distribution shift, fixed }f^*_\mathcal{D}}. \end{align}\] Terms (I) and (III) each satisfy Step 1 with \(f = f^*_{\tilde{\mathcal{D}}}\) and \(f = f^*_\mathcal{D}\) respectively, so each is bounded by \(L \cdot W_1(P^w, P_{\tilde{\mathcal{D}}})\). Term (II) is bounded pointwise by (B2): \(\text{(II)} \leq \beta_K = O(1/K)\).

Combining. \[|\mathcal{L}(\tilde{\mathcal{D}}) - \mathcal{L}(\mathcal{D})| \;\leq\; 2L \cdot W_1(P^w, P_{\tilde{\mathcal{D}}}) + \beta_K.\] Minimizing \(W_1(P^w, P_{\tilde{\mathcal{D}}})\) via the OT objective directly minimizes the dominant term; \(\beta_K \to 0\) as \(K \to \infty\), a guarantee that k-means and greedy coreset methods lack. \(\square\)

References: Kantorovich & Rubinstein (1958); Villani (2009, Ch. 6).

9 Entropic Regularization for Tractability↩︎

This section extends the derivation in Section 4.4 of the main paper. The discrete OT problem is defined over empirical measures \(\mu = P^w = \sum_n w_n \delta_{z_n}\) (TAKE-weighted source) and \(\nu = \frac{1}{|\mathcal{D}_\text{synth}|}\sum_c \delta_{\tilde{z}_c}\) (uniform target over the synthetic pool).

Transportation cost (Kantorovich form). \[\mathcal{C}(\gamma) = \int c(x,\hat{x})\,d\gamma(x,\hat{x}), \qquad c(x,\hat{x}) = \|\phi(x) - \phi(\hat{x})\|^2, \quad \gamma \in \Pi(\mu,\nu).\]

Entropic regularization. The bare Kantorovich problem is \(O(N^3)\) complexity and non-differentiable. Adding an entropic penalty yields a strictly convex, smooth objective: \[\min_{\gamma \in \Pi(\mu,\nu)} \langle C, \gamma \rangle + \varepsilon \cdot \mathrm{KL}(\gamma, \mu \otimes \nu), \qquad \mathrm{KL}(\gamma, \mu \otimes \nu) = \int \log\frac{d\gamma}{d(\mu \otimes \nu)}\,d\gamma,\] where \(\mu \otimes \nu\) is the product (independence) reference coupling. Smaller \(\varepsilon\) recovers the exact OT solution; larger \(\varepsilon\) smooths the plan toward the independent coupling, encouraging diversity in the selected prototypes.

Optional prior coupling. If a task-specific prior \(\alpha(x,\hat{x})\) is available (e.g., a topic-conditional measure), replace \(\mu \otimes \nu\) with \(\alpha\) to regularize toward domain knowledge.

Sinkhorn–Knopp solver. The solution is computed by alternating row/column normalization on the kernel matrix \(\mathbf{K}_{nc} = e^{-C_{nc}/\varepsilon}\), yielding \(O(N \cdot |\mathcal{D}_\text{synth}|)\) per iteration.

Prototype extraction. The distilled set \(\tilde{\mathcal{D}}\) is formed by selecting the top-\(K\) candidates ranked by received mass \(\sum_n \gamma^*_{nc}\).

10 Implementation Details↩︎

Datasets. Table 2 reports the task type and label set for each benchmark.

Backbone models. Classification experiments use well-known classifiers such as Logistic Regression, TextCNN, TextRNN, and BERT. NLI experiments use Siamese Logistic Regression and BERT. These architectures are chosen as standard, reproducible benchmarks spanning the major NLP paradigms while remaining tractable at ablation study scale.

Hyperparameters. All experiments run on a single NVIDIA V100 40 GB GPU using PyTorch, PyTorch Lightning, HuggingFace Transformers, and the Python Optimal Transport (POT) library. Influence scoring uses all-MiniLM-L6-v2 [36], [41] as the frozen sentence encoder \(\phi\). We use \(T = 20\) checkpoints, balancing trajectory coverage and compute. The kernel decay is set to \(\lambda^* = T / T_m\), concentrating mass on the pre-memorization phase as derived in Section 4.2.2. For the Sinkhorn solver we use \(\varepsilon = 0.05\) with 200 iterations, which produces soft plans with sufficient diversity and converges for all tested pool sizes. The synthetic pool size is fixed at \(|\mathcal{D}_\text{synth}| = 5{,}000\), chosen based on empirical saturation across all tested domains. The language model generator is Gemma-3-270M, which is small, fast, and sufficient for in-domain fluency at this pool size.

References↩︎